fix(psycopg): prevent proxy deadlock on TLS 1.3 post-handshake messages (GoogleCloudPlatform#1494)

Author: kgala2

google/cloud/sql/connector/psycopg.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import contextlib
 import logging
 import os
 import platform
@@ -52,15 +53,26 @@ def _forward_bytes(
     dst: socket.socket,
     src_name: str,
     dst_name: str,
+    *,
+    nonblocking: bool = False,
 ) -> bool:
     """Read a chunk of bytes from src and send it to dst.
 
     Returns:
         bool: True if bytes were forwarded, False on EOF or OSError.
     """
     try:
-        data = src.recv(_BUFFER_SIZE)
+        if nonblocking:
+            src.setblocking(False)
+        try:
+            data = src.recv(_BUFFER_SIZE)
+        finally:
+            if nonblocking:
+                with contextlib.suppress(OSError):
+                    src.setblocking(True)
     except OSError as e:
+        if nonblocking and isinstance(e, (ssl.SSLWantReadError, BlockingIOError)):
+            return True
         logger.debug(f"Failed to receive from {src_name} socket: {e}")
         return False
     if not data:
@@ -118,7 +130,7 @@ def _proxy(local: socket.socket, remote: "ssl.SSLSocket") -> None:
                 ):
                     return
                 if key.data == "remote" and not _forward_bytes(
-                    remote, local, "remote", "local"
+                    remote, local, "remote", "local", nonblocking=True
                 ):
                     return
     except OSError as e:

tests/unit/test_psycopg.py

@@ -813,3 +813,49 @@ def socket_side_effect(family, type, proto=0, fileno=None):
 
     mock_local_sock.close.assert_called_once()
     assert mock_remote_sock.close.call_count >= 1
+
+
+def test_proxy_tls13_post_handshake_want_read() -> None:
+    """Test that _proxy does not deadlock or exit on SSLWantReadError (e.g. TLS 1.3 NewSessionTicket)."""
+    local_client, local_server = socket.socketpair()
+    remote_client, remote_server = socket.socketpair()
+
+    class MockTLS13SSLSocket:
+        def __init__(self, sock: socket.socket) -> None:
+            self._sock = sock
+            self._want_read_once = True
+
+        def pending(self) -> int:
+            return 0
+
+        def recv(self, bufsize: int, flags: int = 0) -> bytes:
+            if self._want_read_once:
+                self._want_read_once = False
+                # Drain the 1-byte dummy ticket wakeup from the underlying socket
+                self._sock.recv(1)
+                raise ssl.SSLWantReadError()
+            return self._sock.recv(bufsize, flags)
+
+        def __getattr__(self, name: str) -> Any:
+            return getattr(self._sock, name)
+
+    wrapped_remote = MockTLS13SSLSocket(remote_client)
+
+    # Simulate TLS 1.3 NewSessionTicket arriving on remote BEFORE client sends StartupMessage
+    remote_server.sendall(b"T")
+
+    proxy_thread = threading.Thread(
+        target=_proxy, args=(local_server, wrapped_remote), daemon=True
+    )
+    proxy_thread.start()
+
+    # Client sends StartupMessage and receives response without deadlocking
+    local_client.sendall(b"StartupMessage")
+    assert remote_server.recv(1024) == b"StartupMessage"
+    remote_server.sendall(b"AuthenticationOk")
+    assert local_client.recv(1024) == b"AuthenticationOk"
+
+    local_client.close()
+    proxy_thread.join(timeout=2.0)
+    assert not proxy_thread.is_alive()
+    remote_server.close()