fix(postgresql-proxy): prevent SSL COPY stalls by draining nonblocking reads

Author: nik-localstack

postgresql_proxy/proxy.py

@@ -156,6 +156,7 @@ def accept_wrapper(self, sock: socket.socket):
             events=events,
             context=context,
         )
+        conn.ssl_handshake_done = not isinstance(clientsocket, ssl.SSLSocket)
 
         pg_conn = self._create_pg_connection(address, context)
 
@@ -195,9 +196,20 @@ def _handle_ssl_negotiation(
         Returns the SSL-wrapped socket if negotiation succeeds, or the original socket.
         """
 
-        # Peek at the first 8 bytes to check for SSLRequest
-        # Using MSG_PEEK so we don't consume the data if it's not SSLRequest
-        data = client_socket.recv(8, socket.MSG_PEEK)
+        # Peek at the first 8 bytes to check for SSLRequest.
+        # Wait briefly to avoid race conditions where the request arrives just after accept.
+        previous_timeout = client_socket.gettimeout()
+        try:
+            client_socket.settimeout(0.2)
+            data = client_socket.recv(8, socket.MSG_PEEK)
+        except socket.timeout:
+            return client_socket
+        except BlockingIOError:
+            return client_socket
+        except OSError:
+            return client_socket
+        finally:
+            client_socket.settimeout(previous_timeout)
 
         if len(data) == 8:
             length = int.from_bytes(data[:4], "big")
@@ -208,14 +220,68 @@ def _handle_ssl_negotiation(
                 client_socket.recv(8)
                 # Send 'S' to indicate SSL is supported
                 client_socket.send(b"S")
-                # Wrap socket with SSL
-                ssl_socket = ssl_context.wrap_socket(client_socket, server_side=True)
-                LOG.debug("SSL handshake completed for PostgreSQL connection")
+                # Wrap without immediate handshake so a stalled client can't block accept loop.
+                ssl_socket = ssl_context.wrap_socket(
+                    client_socket,
+                    server_side=True,
+                    do_handshake_on_connect=False,
+                )
+                LOG.debug("SSL requested, deferring TLS handshake to selector loop")
                 return ssl_socket
 
         # Not an SSLRequest, return original socket
         return client_socket
 
+    def _set_write_interest(self, conn: connection.Connection, enabled: bool):
+        """Enable or disable EVENT_WRITE for a connection while preserving read interest."""
+        try:
+            selector_key = self.selector.get_key(conn.sock)
+        except KeyError:
+            return
+
+        current_events = selector_key.events
+        if enabled:
+            new_events = current_events | selectors.EVENT_WRITE
+        else:
+            new_events = current_events & ~selectors.EVENT_WRITE
+
+        if new_events != current_events:
+            self.selector.modify(conn.sock, new_events, data=conn)
+            conn.events = new_events
+
+    def _flush_outgoing(
+        self,
+        source_conn: connection.Connection,
+        source_sock: socket.socket,
+        target_conn: connection.Connection | None,
+    ):
+        if not target_conn or not target_conn.out_bytes:
+            return
+
+        try:
+            while target_conn.out_bytes:
+                LOG.debug('sending to %s:\n%s', target_conn.name, target_conn.out_bytes)
+                sent = target_conn.sock.send(target_conn.out_bytes)
+                if sent == 0:
+                    break
+                target_conn.sent(sent)
+        except ssl.SSLWantWriteError:
+            self._set_write_interest(target_conn, True)
+            return
+        except ssl.SSLWantReadError:
+            self._set_write_interest(target_conn, True)
+            return
+        except OSError:
+            # If one side is closed, close the other one
+            # this can happen in the case where the client disconnects, and postgres still return a response
+            # we then read the response then close the PG side of the socket.
+            LOG.debug('error sending to %s: connection closed', target_conn.name)
+            self._unregister_conn(source_conn)
+            source_sock.close()
+            return
+
+        self._set_write_interest(target_conn, bool(target_conn.out_bytes))
+
     def service_connection(self, key: SelectorKeyProxy, mask):
         """
         This method proxies the messages between socket. It will use properties of the Connection object to
@@ -227,37 +293,58 @@ def service_connection(self, key: SelectorKeyProxy, mask):
         """
         sock = key.fileobj
         conn = key.data
+
+        # Drive TLS handshake in non-blocking mode so one slow client cannot block others.
+        if isinstance(sock, ssl.SSLSocket) and not getattr(conn, "ssl_handshake_done", False):
+            try:
+                sock.do_handshake()
+                conn.ssl_handshake_done = True
+                self._set_write_interest(conn, bool(conn.out_bytes))
+            except ssl.SSLWantReadError:
+                return
+            except ssl.SSLWantWriteError:
+                self._set_write_interest(conn, True)
+                return
+            except OSError as e:
+                LOG.debug('%s SSL handshake failed %s: %s', conn.name, conn.address, e)
+                self._unregister_conn(conn)
+                sock.close()
+                return
+
         if mask & selectors.EVENT_READ:
             LOG.debug('%s can receive', conn.name)
-            try:
-                recv_data = sock.recv(4096)  # Should be ready to read
-                if recv_data:
-                    LOG.debug('%s received data:\n%s', conn.name, recv_data)
-                    conn.received(recv_data)
-                else:
+            while True:
+                try:
+                    if recv_data := sock.recv(4096):
+                        LOG.debug('%s received data:\n%s', conn.name, recv_data)
+                        conn.received(recv_data)
+                        # Keep draining bytes in the same readiness cycle until recv indicates no immediate data.
+                        continue
+
                     self._unregister_conn(conn)
                     LOG.debug('%s connection closing %s', conn.name, conn.address)
                     # A file object shall be unregistered prior to being closed.
                     sock.close()
-            except OSError as e:
-                # it means the socket was closed by peer
-                LOG.debug('%s connection closed by peer %s: %s', conn.name, conn.address, e)
-                self._unregister_conn(conn)
+                    return
+                except ssl.SSLWantReadError:
+                    break
+                except ssl.SSLWantWriteError:
+                    self._set_write_interest(conn, True)
+                    break
+                except BlockingIOError:
+                    break
+                except OSError as e:
+                    # it means the socket was closed by peer
+                    LOG.debug('%s connection closed by peer %s: %s', conn.name, conn.address, e)
+                    self._unregister_conn(conn)
+                    return
 
-        next_conn = conn.redirect_conn
-        if next_conn and next_conn.out_bytes:
-            try:
-                while next_conn.out_bytes:
-                    LOG.debug('sending to %s:\n%s', next_conn.name, next_conn.out_bytes)
-                    sent = next_conn.sock.send(next_conn.out_bytes)
-                    next_conn.sent(sent)
-            except OSError:
-                # If one side is closed, close the other one
-                # this can happen in the case where the client disconnects, and postgres still return a response
-                # we then read the response then close the PG side of the socket.
-                LOG.debug('error sending to %s: connection closed', next_conn.name)
-                self._unregister_conn(conn)
-                sock.close()
+            next_conn = conn.redirect_conn
+            if next_conn and next_conn.out_bytes:
+                self._flush_outgoing(conn, sock, next_conn)
+
+        if mask & selectors.EVENT_WRITE:
+            self._flush_outgoing(conn, sock, conn)
 
     def listen(self, max_connections: int = 8):
         """