Sockets: Properly handle POLLHUP events

Stops client socket leaking.
stenzek · Jul 21, 2024 · 7880087 · 7880087
1 parent 6b78364
commit 7880087
Show file tree

Hide file tree

Showing 2 changed files with 35 additions and 4 deletions.
diff --git 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
@@ -79,6 +79,7 @@ class BaseSocket : public std::enable_shared_from_this<BaseSocket>
 protected:
   virtual void OnReadEvent() = 0;
   virtual void OnWriteEvent() = 0;
+  virtual void OnHangupEvent() = 0;
 
   SocketMultiplexer& m_multiplexer;
   SocketDescriptor m_descriptor;
@@ -194,6 +195,7 @@ class ListenSocket final : public BaseSocket
 protected:
   void OnReadEvent() override final;
   void OnWriteEvent() override final;
+  void OnHangupEvent() override final;
 
 private:
   SocketMultiplexer::CreateStreamSocketCallback m_accept_callback;
@@ -231,6 +233,7 @@ class StreamSocket : public BaseSocket
 
   virtual void OnReadEvent() override;
   virtual void OnWriteEvent() override;
+  virtual void OnHangupEvent() override;
 
   void CloseWithError();