@@ -422,3 +422,312 @@ def mock_connect_impl(*args: Any, **kwargs: Any) -> MagicMock:
422422 assert conn is not None
423423 assert mock_remove .called
424424 assert mock_rmdir .called
425+
426+
427+ def test_proxy_happy_path_sequential () -> None :
428+ """Test that _proxy forwards sequential request/response traffic cleanly without errors."""
429+ local_client , local_proxy = socket .socketpair (
430+ socket .AF_UNIX , socket .SOCK_STREAM
431+ )
432+ remote_proxy , remote_server = socket .socketpair (
433+ socket .AF_UNIX , socket .SOCK_STREAM
434+ )
435+
436+ t = threading .Thread (
437+ target = _proxy , args = (local_proxy , remote_proxy ), daemon = True
438+ )
439+ t .start ()
440+
441+ # Step 1: Client sends query
442+ local_client .sendall (b"SELECT 1;" )
443+ received_query = remote_server .recv (1024 )
444+ assert received_query == b"SELECT 1;"
445+
446+ # Step 2: Server sends response
447+ remote_server .sendall (b"RESULT 1" )
448+ received_resp = local_client .recv (1024 )
449+ assert received_resp == b"RESULT 1"
450+
451+ # Step 3: Client closes connection cleanly
452+ local_client .shutdown (socket .SHUT_RDWR )
453+ local_client .close ()
454+ remote_server .close ()
455+ t .join (timeout = 2.0 )
456+ assert not t .is_alive ()
457+
458+
459+ def test_proxy_backpressure_and_clean_teardown () -> None :
460+ """Demonstrate that Cloud SQL's single-threaded selector proxy exits cleanly
461+
462+ without deadlocking on teardown.
463+ """
464+ local_client , local_proxy = socket .socketpair (
465+ socket .AF_UNIX , socket .SOCK_STREAM
466+ )
467+ remote_proxy , remote_server = socket .socketpair (
468+ socket .AF_UNIX , socket .SOCK_STREAM
469+ )
470+
471+ t_proxy = threading .Thread (
472+ target = _proxy , args = (local_proxy , remote_proxy ), daemon = True
473+ )
474+ t_proxy .start ()
475+
476+ # Step 1: Client sends query
477+ local_client .sendall (b"SELECT 1;" )
478+ assert remote_server .recv (1024 ) == b"SELECT 1;"
479+
480+ # Step 2: Server sends response
481+ remote_server .sendall (b"RESULT 1" )
482+ assert local_client .recv (1024 ) == b"RESULT 1"
483+
484+ # Step 3: Client closes connection while server is open
485+ local_client .close ()
486+
487+ # Single-threaded proxy terminates immediately without deadlocking
488+ t_proxy .join (timeout = 2.0 )
489+ assert (
490+ not t_proxy .is_alive ()
491+ ), "Cloud SQL proxy must terminate cleanly without deadlocks"
492+
493+ remote_server .close ()
494+
495+
496+ def test_proxy_pending_oserror_in_loop () -> None :
497+ """Test that _proxy pending loop handles OSError on pending() and continues."""
498+ local_client , local_proxy = mockable_socketpair ()
499+ remote_proxy , remote_server = mockable_socketpair ()
500+
501+ # Add pending mock to real socket. First call returns 10, second raises OSError.
502+ remote_proxy .pending = MagicMock (side_effect = [10 , OSError ("pending failed" )])
503+
504+ # Send data to be read by recv in pending loop
505+ remote_server .sendall (b"pending data" )
506+
507+ # Start proxy in thread because it will block on select() after pending fails
508+ t = threading .Thread (target = _proxy , args = (local_proxy , remote_proxy ), daemon = True )
509+ t .start ()
510+
511+ # Verify local_client received the pending data
512+ assert local_client .recv (1024 ) == b"pending data"
513+
514+ # Now trigger proxy exit by closing local_client
515+ local_client .close ()
516+
517+ t .join (timeout = 2.0 )
518+ assert not t .is_alive ()
519+
520+ remote_server .close ()
521+
522+
523+ def test_proxy_pending_non_int () -> None :
524+ """Test that _proxy pending check handles non-int return values from pending()."""
525+ local_client , local_proxy = mockable_socketpair ()
526+ remote_proxy , remote_server = mockable_socketpair ()
527+
528+ # First call returns non-int. Should return False (goes to select)
529+ remote_proxy .pending = MagicMock (return_value = "not an int" )
530+
531+ # We need to run it in a thread because it will block on select
532+ t = threading .Thread (target = _proxy , args = (local_proxy , remote_proxy ), daemon = True )
533+ t .start ()
534+
535+ # Trigger exit.
536+ local_client .close ()
537+ t .join (timeout = 2.0 )
538+ assert not t .is_alive ()
539+ remote_server .close ()
540+
541+
542+ def test_proxy_pending_non_int_in_loop () -> None :
543+ """Test that _proxy pending loop handles non-int return values from pending() in loop."""
544+ local_client , local_proxy = mockable_socketpair ()
545+ remote_proxy , remote_server = mockable_socketpair ()
546+
547+ # Second call returns non-int.
548+ remote_proxy .pending = MagicMock (side_effect = [10 , "not an int" ])
549+ remote_server .sendall (b"data" )
550+
551+ t = threading .Thread (target = _proxy , args = (local_proxy , remote_proxy ), daemon = True )
552+ t .start ()
553+
554+ assert local_client .recv (1024 ) == b"data"
555+
556+ local_client .close ()
557+ t .join (timeout = 2.0 )
558+ assert not t .is_alive ()
559+ remote_server .close ()
560+
561+
562+ def test_proxy_pending_recv_eof () -> None :
563+ """Test that _proxy pending loop handles remote recv EOF."""
564+ local_client , local_server = mockable_socketpair ()
565+ remote_client , remote_server = mockable_socketpair ()
566+
567+ remote_client .pending = MagicMock (return_value = 10 )
568+ # Mock recv to return EOF
569+ remote_client .recv = MagicMock (return_value = b"" )
570+
571+ # Run proxy. It should call forward_pending, which calls recv, gets EOF,
572+ # and returns True. This breaks the loop and proxy exits.
573+ _proxy (local_server , remote_client )
574+
575+ # Sockets should be closed
576+ assert local_server .fileno () == - 1
577+ assert remote_client .fileno () == - 1
578+
579+ local_client .close ()
580+ remote_server .close ()
581+
582+
583+ def test_proxy_select_oserror () -> None :
584+ """Test that _proxy loop handles OSError in select()."""
585+ local_client , local_server = mockable_socketpair ()
586+ remote_client , remote_server = mockable_socketpair ()
587+
588+ with patch (
589+ "google.cloud.sql.connector.psycopg.selectors.DefaultSelector"
590+ ) as mock_sel_cls :
591+ mock_sel = MagicMock ()
592+ mock_sel .select .side_effect = OSError ("select failed" )
593+ mock_sel_cls .return_value = mock_sel
594+
595+ # Run proxy. It should catch OSError on select() and exit loop to finally block.
596+ _proxy (local_server , remote_client )
597+
598+ assert local_server .fileno () == - 1
599+ assert remote_client .fileno () == - 1
600+
601+ local_client .close ()
602+ remote_server .close ()
603+
604+
605+ def test_proxy_finally_cleanup_errors () -> None :
606+ """Test that _proxy finally block ignores OSErrors during socket shutdown/close."""
607+ local_client , local_server = mockable_socketpair ()
608+ remote_client , remote_server = mockable_socketpair ()
609+
610+ real_local_shutdown = local_server .shutdown
611+ def mock_local_shutdown (how ):
612+ try :
613+ real_local_shutdown (how )
614+ except OSError :
615+ pass
616+ raise OSError ("shutdown failed" )
617+ local_server .shutdown = MagicMock (side_effect = mock_local_shutdown )
618+
619+ real_local_close = local_server .close
620+ def mock_local_close ():
621+ real_local_close ()
622+ raise OSError ("close failed" )
623+ local_server .close = MagicMock (side_effect = mock_local_close )
624+
625+ real_remote_shutdown = remote_client .shutdown
626+ def mock_remote_shutdown (how ):
627+ try :
628+ real_remote_shutdown (how )
629+ except OSError :
630+ pass
631+ raise OSError ("shutdown failed" )
632+ remote_client .shutdown = MagicMock (side_effect = mock_remote_shutdown )
633+
634+ real_remote_close = remote_client .close
635+ def mock_remote_close ():
636+ real_remote_close ()
637+ raise OSError ("close failed" )
638+ remote_client .close = MagicMock (side_effect = mock_remote_close )
639+
640+ # Trigger exit by closing client
641+ local_client .close ()
642+
643+ # Run proxy. It should handle the exceptions in finally block.
644+ _proxy (local_server , remote_client )
645+
646+ remote_server .close ()
647+
648+
649+ @patch ("google.cloud.sql.connector.psycopg._proxy" )
650+ def test_accept_and_proxy_cleanup_errors (mock_proxy_fn : MagicMock ) -> None :
651+ """Test that accept_and_proxy thread handles errors during cleanup on exception."""
652+ # Make _proxy raise an exception to trigger the except block in _accept_and_proxy
653+ mock_proxy_fn .side_effect = Exception ("proxy crash" )
654+
655+ mock_unix_conn = MagicMock (spec = socket .socket )
656+ mock_unix_conn .shutdown .side_effect = OSError ("unix conn shutdown failed" )
657+ mock_unix_conn .close .side_effect = OSError ("unix conn close failed" )
658+
659+ mock_remote_sock = MagicMock (spec = ssl .SSLSocket )
660+ mock_remote_sock .shutdown .side_effect = OSError ("remote shutdown failed" )
661+
662+ # Mock socket.socket to return a mock local_sock that returns our mock_unix_conn
663+ real_socket = socket .socket
664+ mock_local_sock = MagicMock ()
665+ mock_local_sock .accept .return_value = (mock_unix_conn , ("path" ,))
666+
667+ def socket_side_effect (family , type , proto = 0 , fileno = None ):
668+ if family == socket .AF_UNIX :
669+ return mock_local_sock
670+ return real_socket (family , type , proto , fileno )
671+
672+ event = threading .Event ()
673+ def remote_close_fn ():
674+ event .set ()
675+ raise OSError ("remote close failed" )
676+ mock_remote_sock .close .side_effect = remote_close_fn
677+
678+ with patch ("socket.socket" , side_effect = socket_side_effect ), patch (
679+ "psycopg.connect"
680+ ) as mock_psycopg_connect :
681+ mock_psycopg_connect .return_value = MagicMock ()
682+
683+ connect (
684+ "127.0.0.1" ,
685+ mock_remote_sock ,
686+ user = "test_user" ,
687+ db = "test_db" ,
688+ password = "test_password" ,
689+ )
690+
691+ # Wait for the accept_and_proxy thread to finish cleanup
692+ assert event .wait (timeout = 2.0 )
693+
694+ # Verify mock_unix_conn shutdown and close were called
695+ mock_unix_conn .shutdown .assert_called_once_with (socket .SHUT_RDWR )
696+ mock_unix_conn .close .assert_called_once ()
697+ mock_remote_sock .shutdown .assert_called_once_with (socket .SHUT_RDWR )
698+ mock_remote_sock .close .assert_called_once ()
699+
700+
701+ @patch ("psycopg.connect" )
702+ def test_connect_wrapper_failure_cleanup_errors (mock_psycopg_connect : MagicMock ) -> None :
703+ """Test that connect wrapper ignores OSErrors during cleanup on connection failure."""
704+ mock_remote_sock = MagicMock (spec = ssl .SSLSocket )
705+ mock_remote_sock .close .side_effect = OSError ("remote close failed" )
706+
707+ mock_psycopg_connect .side_effect = Exception ("connection failed simulated" )
708+
709+ mock_local_sock = MagicMock ()
710+ mock_local_sock .close .side_effect = OSError ("local close failed" )
711+
712+ real_socket = socket .socket
713+ def socket_side_effect (family , type , proto = 0 , fileno = None ):
714+ if family == socket .AF_UNIX :
715+ return mock_local_sock
716+ return real_socket (family , type , proto , fileno )
717+
718+ with (
719+ patch ("socket.socket" , side_effect = socket_side_effect ),
720+ pytest .raises (Exception , match = "connection failed simulated" ),
721+ ):
722+ connect (
723+ "127.0.0.1" ,
724+ mock_remote_sock ,
725+ user = "test_user" ,
726+ db = "test_db" ,
727+ password = "test_password" ,
728+ )
729+
730+ mock_local_sock .close .assert_called_once ()
731+ assert mock_remote_sock .close .call_count >= 1
732+
733+
0 commit comments