diff --git a/google/cloud/sql/connector/psycopg.py b/google/cloud/sql/connector/psycopg.py index a1d9f3da..c9a12c97 100644 --- a/google/cloud/sql/connector/psycopg.py +++ b/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,6 +53,8 @@ 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. @@ -59,8 +62,17 @@ def _forward_bytes( 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: diff --git a/tests/unit/test_psycopg.py b/tests/unit/test_psycopg.py index e419d3cf..c745ccb0 100644 --- a/tests/unit/test_psycopg.py +++ b/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()