From 78800870bc9fb97ab85fc78590833992b3e8dec4 Mon Sep 17 00:00:00 2001
From: Stenzek <stenzek@gmail.com>
Date: Sun, 21 Jul 2024 14:22:55 +1000
Subject: [PATCH] Sockets: Properly handle POLLHUP events

Stops client socket leaking.
---
 src/util/sockets.cpp | 36 ++++++++++++++++++++++++++++++++----
 src/util/sockets.h   |  3 +++
 2 files changed, 35 insertions(+), 4 deletions(-)

diff --git a/src/util/sockets.cpp b/src/util/sockets.cpp
index 9492a76a1..0c5e82b33 100644
--- a/src/util/sockets.cpp
+++ b/src/util/sockets.cpp
@@ -486,10 +486,17 @@ bool SocketMultiplexer::PollEventsWithTimeout(u32 milliseconds)
     PendingSocketPair& psp = triggered_sockets[i];
 
     // fire events
-    if (psp.second & (POLLIN | POLLHUP | POLLERR))
-      psp.first->OnReadEvent();
-    if (psp.second & POLLOUT)
-      psp.first->OnWriteEvent();
+    if (psp.second & (POLLHUP | POLLERR))
+    {
+      psp.first->OnHangupEvent();
+    }
+    else
+    {
+      if (psp.second & POLLIN)
+        psp.first->OnReadEvent();
+      if (psp.second & POLLOUT)
+        psp.first->OnWriteEvent();
+    }
 
     psp.first.~shared_ptr();
   }
@@ -558,6 +565,12 @@ void ListenSocket::OnReadEvent()
 
 void ListenSocket::OnWriteEvent()
 {
+  ERROR_LOG("Unexpected OnWriteEvent() in ListenSocket {}", m_local_address.ToString());
+}
+
+void ListenSocket::OnHangupEvent()
+{
+  ERROR_LOG("Unexpected OnHangupEvent() in ListenSocket {}", m_local_address.ToString());
 }
 
 StreamSocket::StreamSocket(SocketMultiplexer& multiplexer, SocketDescriptor descriptor)
@@ -770,6 +783,21 @@ void StreamSocket::OnWriteEvent()
   // shouldn't be called
 }
 
+void StreamSocket::OnHangupEvent()
+{
+  std::unique_lock lock(m_lock);
+  if (!m_connected)
+    return;
+
+  m_multiplexer.SetNotificationMask(this, m_descriptor, 0);
+  m_multiplexer.RemoveClientSocket(this);
+  closesocket(m_descriptor);
+  m_descriptor = INVALID_SOCKET;
+  m_connected = false;
+
+  OnDisconnected(Error::CreateString("Connection closed by peer."));
+}
+
 BufferedStreamSocket::BufferedStreamSocket(SocketMultiplexer& multiplexer, SocketDescriptor descriptor,
                                            size_t receive_buffer_size /* = 16384 */,
                                            size_t send_buffer_size /* = 16384 */)
diff --git a/src/util/sockets.h b/src/util/sockets.h
index b01918664..2a4375059 100644
--- a/src/util/sockets.h
+++ b/src/util/sockets.h
@@ -79,6 +79,7 @@ public:
 protected:
   virtual void OnReadEvent() = 0;
   virtual void OnWriteEvent() = 0;
+  virtual void OnHangupEvent() = 0;
 
   SocketMultiplexer& m_multiplexer;
   SocketDescriptor m_descriptor;
@@ -194,6 +195,7 @@ public:
 protected:
   void OnReadEvent() override final;
   void OnWriteEvent() override final;
+  void OnHangupEvent() override final;
 
 private:
   SocketMultiplexer::CreateStreamSocketCallback m_accept_callback;
@@ -231,6 +233,7 @@ protected:
 
   virtual void OnReadEvent() override;
   virtual void OnWriteEvent() override;
+  virtual void OnHangupEvent() override;
 
   void CloseWithError();