From 5510b31972390ed4392717e0b65909750e2bb8ba Mon Sep 17 00:00:00 2001
From: Liam <byteslice@airmail.cc>
Date: Tue, 30 Jan 2024 10:29:05 -0500
Subject: [PATCH] internal_network: only poll for accept on blocking sockets

---
 src/core/internal_network/network.cpp | 30 +++++++++++++++------------
 src/core/internal_network/sockets.h   |  3 +++
 2 files changed, 20 insertions(+), 13 deletions(-)

diff --git a/src/core/internal_network/network.cpp b/src/core/internal_network/network.cpp
index a983f23ea..7785c1d16 100644
--- a/src/core/internal_network/network.cpp
+++ b/src/core/internal_network/network.cpp
@@ -693,20 +693,23 @@ std::pair<SocketBase::AcceptResult, Errno> Socket::Accept() {
     sockaddr_in addr;
     socklen_t addrlen = sizeof(addr);
 
-    std::vector<WSAPOLLFD> host_pollfds{
-        WSAPOLLFD{fd, POLLIN, 0},
-        WSAPOLLFD{GetInterruptSocket(), POLLIN, 0},
-    };
+    const bool wait_for_accept = !is_non_blocking;
+    if (wait_for_accept) {
+        std::vector<WSAPOLLFD> host_pollfds{
+            WSAPOLLFD{fd, POLLIN, 0},
+            WSAPOLLFD{GetInterruptSocket(), POLLIN, 0},
+        };
 
-    while (true) {
-        const int pollres =
-            WSAPoll(host_pollfds.data(), static_cast<ULONG>(host_pollfds.size()), -1);
-        if (host_pollfds[1].revents != 0) {
-            // Interrupt signaled before a client could be accepted, break
-            return {AcceptResult{}, Errno::AGAIN};
-        }
-        if (pollres > 0) {
-            break;
+        while (true) {
+            const int pollres =
+                WSAPoll(host_pollfds.data(), static_cast<ULONG>(host_pollfds.size()), -1);
+            if (host_pollfds[1].revents != 0) {
+                // Interrupt signaled before a client could be accepted, break
+                return {AcceptResult{}, Errno::AGAIN};
+            }
+            if (pollres > 0) {
+                break;
+            }
         }
     }
 
@@ -913,6 +916,7 @@ Errno Socket::SetRcvTimeo(u32 value) {
 
 Errno Socket::SetNonBlock(bool enable) {
     if (EnableNonBlock(fd, enable)) {
+        is_non_blocking = enable;
         return Errno::SUCCESS;
     }
     return GetAndLogLastError();
diff --git a/src/core/internal_network/sockets.h b/src/core/internal_network/sockets.h
index 4ba51f62c..3a32dff75 100644
--- a/src/core/internal_network/sockets.h
+++ b/src/core/internal_network/sockets.h
@@ -166,6 +166,9 @@ public:
     bool IsOpened() const override;
 
     void HandleProxyPacket(const ProxyPacket& packet) override;
+
+private:
+    bool is_non_blocking = false;
 };
 
 std::pair<s32, Errno> Poll(std::vector<PollFD>& poll_fds, s32 timeout);