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

Co-authored-by: GitHub Copilot <copilot@github.com>

Author: nik-localstack

postgresql_proxy/connection.py

@@ -16,6 +16,8 @@ def __init__(self, sock, address, name, events, context):
         self.out_bytes = b''
         self.in_bytes = b''
         self.terminated = False
+        self.ssl_handshake_done = False
+        self.ssl_negotiation_pending = False
 
     def parse_length(self, length_bytes):
         return int.from_bytes(length_bytes, 'big')

postgresql_proxy/proxy.py

@@ -128,14 +128,8 @@ def accept_wrapper(self, sock: socket.socket):
         :return:
         """
 
-        # Accept the raw connection
+        # Accept the raw connection and switch to non-blocking immediately.
         clientsocket, address = sock.accept()
-
-        # Check if SSL is enabled for this proxy
-        if self.ssl_context:
-            # Handle SSL negotiation - must happen before setblocking(False)
-            clientsocket = self._handle_ssl_negotiation(clientsocket, self.ssl_context)
-
         clientsocket.setblocking(False)
         self.num_clients += 1
         sock_name = f"{self.instance_config.listen.name}_{self.num_clients}"
@@ -156,6 +150,9 @@ def accept_wrapper(self, sock: socket.socket):
             events=events,
             context=context,
         )
+        # SSL startup is handled in the selector loop to avoid blocking accept.
+        conn.ssl_negotiation_pending = self.ssl_context is not None
+        conn.ssl_handshake_done = self.ssl_context is None
 
         pg_conn = self._create_pg_connection(address, context)
 
@@ -180,41 +177,121 @@ def accept_wrapper(self, sock: socket.socket):
         self._register_conn(conn)
         self._register_conn(pg_conn)
 
-    def _handle_ssl_negotiation(
-        self, client_socket: socket.socket, ssl_context: ssl.SSLContext
-    ) -> socket.socket:
+    def _try_start_client_tls(self, conn: connection.Connection) -> bool:
+        """Resolve PostgreSQL SSLRequest in non-blocking mode.
+
+        Returns True when startup mode is decided (SSL or plain), False when more
+        bytes are needed.
         """
-        Handle PostgreSQL SSL negotiation on an accepted socket.
 
-        PostgreSQL SSL flow:
-        1. Client sends SSLRequest (8 bytes): length (4) + code 80877103 (4)
-        2. Server responds 'S' (SSL supported) or 'N' (not supported)
-        3. If 'S', TLS handshake follows
-        4. After TLS, normal PostgreSQL protocol begins
+        if not conn.ssl_negotiation_pending:
+            return True
 
-        Returns the SSL-wrapped socket if negotiation succeeds, or the original socket.
-        """
+        sock = conn.sock
+        try:
+            data = sock.recv(8, socket.MSG_PEEK)
+        except BlockingIOError:
+            return False
+        except OSError as exc:
+            LOG.debug("%s SSL startup read failed %s: %s", conn.name, conn.address, exc)
+            self._unregister_conn(conn)
+            sock.close()
+            return False
+
+        if len(data) == 0:
+            self._unregister_conn(conn)
+            sock.close()
+            return False
+
+        if len(data) < 8:
+            return False
+
+        length = int.from_bytes(data[:4], "big")
+        code = int.from_bytes(data[4:8], "big")
+        conn.ssl_negotiation_pending = False
+
+        if length != 8 or code != 80877103:
+            # Plain startup packet.
+            conn.ssl_handshake_done = True
+            return True
+
+        try:
+            # Consume SSLRequest and acknowledge TLS support.
+            sock.recv(8)
+            sock.sendall(b"S")
+            ssl_sock = self.ssl_context.wrap_socket(
+                sock,
+                server_side=True,
+                do_handshake_on_connect=False,
+            )
+            ssl_sock.setblocking(False)
 
-        # 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)
+            if self._debug:
+                self._registered_conn.discard(f"{conn.name}-{conn.sock.fileno()}")
+            self.selector.unregister(conn.sock)
+            conn.sock = ssl_sock
+            self._register_conn(conn)
+            conn.ssl_handshake_done = False
+            LOG.debug("SSL requested, deferring TLS handshake to selector loop")
+            return True
+        except OSError as exc:
+            LOG.debug("%s SSL startup upgrade failed %s: %s", conn.name, conn.address, exc)
+            self._unregister_conn(conn)
+            sock.close()
+            return False
+
+    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
 
-        if len(data) == 8:
-            length = int.from_bytes(data[:4], "big")
-            code = int.from_bytes(data[4:8], "big")
+        current_events = selector_key.events
+        if enabled:
+            new_events = current_events | selectors.EVENT_WRITE
+        else:
+            new_events = current_events & ~selectors.EVENT_WRITE
 
-            if length == 8 and code == 80877103:  # SSLRequest code
-                # Consume the SSLRequest
-                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")
-                return ssl_socket
+        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
 
-        # Not an SSLRequest, return original socket
-        return client_socket
+        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:
+                    # send() returned 0: socket closed or buffer full. Enable write interest
+                    # so the next writable event will retry the send.
+                    self._set_write_interest(target_conn, True)
+                    return
+                target_conn.sent(sent)
+        except (BlockingIOError, ssl.SSLWantWriteError):
+            self._set_write_interest(target_conn, True)
+            return
+        except ssl.SSLWantReadError:
+            self._set_write_interest(target_conn, False)
+            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):
         """
@@ -227,37 +304,63 @@ def service_connection(self, key: SelectorKeyProxy, mask):
         """
         sock = key.fileobj
         conn = key.data
+
+        if conn.ssl_negotiation_pending:
+            if not self._try_start_client_tls(conn):
+                return
+            sock = conn.sock
+
+        # Drive TLS handshake in non-blocking mode so one slow client cannot block others.
+        if isinstance(sock, ssl.SSLSocket) and not conn.ssl_handshake_done:
+            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):
         """