22
33import logging
44import socket
5+ import struct
56import unittest
67import weakref
78from test import support
@@ -153,6 +154,29 @@ def test_connection_lost(self):
153154 test_utils .run_briefly (self .loop )
154155 self .assertIsInstance (waiter .exception (), ConnectionAbortedError )
155156
157+ def test_connection_lost_during_handshake (self ):
158+ # The app protocol gets connection_lost() if and only if it got
159+ # connection_made(). With start_tls() (call_connection_made=False)
160+ # it is already connected when the handshake starts.
161+ for call_connection_made in (True , False ):
162+ with self .subTest (call_connection_made = call_connection_made ):
163+ app_proto = mock .Mock (spec = asyncio .Protocol )
164+ ssl_proto = sslproto .SSLProtocol (
165+ self .loop , app_proto , test_utils .dummy_ssl_context (),
166+ waiter = None , call_connection_made = call_connection_made )
167+ self .connection_made (
168+ ssl_proto ,
169+ do_handshake = mock .Mock (side_effect = ssl .SSLWantReadError )
170+ )
171+ exc = ConnectionResetError ()
172+ ssl_proto .connection_lost (exc )
173+ test_utils .run_briefly (self .loop )
174+ app_proto .connection_made .assert_not_called ()
175+ if call_connection_made :
176+ app_proto .connection_lost .assert_not_called ()
177+ else :
178+ app_proto .connection_lost .assert_called_once_with (exc )
179+
156180 def test_connection_lost_when_busy (self ):
157181 # gh-118950: SSLProtocol.connection_lost not being called when OSError
158182 # is thrown on asyncio.write.
@@ -570,9 +594,10 @@ def serve(sock):
570594 sock .close ()
571595
572596 class ClientProto (asyncio .Protocol ):
573- def __init__ (self , on_data , on_eof ):
597+ def __init__ (self , on_data , on_eof , on_con_lost ):
574598 self .on_data = on_data
575599 self .on_eof = on_eof
600+ self .on_con_lost = on_con_lost
576601 self .con_made_cnt = 0
577602
578603 def connection_made (proto , tr ):
@@ -586,14 +611,18 @@ def data_received(self, data):
586611 def eof_received (self ):
587612 self .on_eof .set_result (True )
588613
614+ def connection_lost (self , exc ):
615+ self .on_con_lost .set_result (exc )
616+
589617 async def client (addr ):
590618 await asyncio .sleep (0.5 )
591619
592620 on_data = self .loop .create_future ()
593621 on_eof = self .loop .create_future ()
622+ on_con_lost = self .loop .create_future ()
594623
595624 tr , proto = await self .loop .create_connection (
596- lambda : ClientProto (on_data , on_eof ), * addr )
625+ lambda : ClientProto (on_data , on_eof , on_con_lost ), * addr )
597626
598627 tr .write (HELLO_MSG )
599628
@@ -604,11 +633,59 @@ async def client(addr):
604633 self .loop .start_tls (tr , proto , client_context ),
605634 0.5 )
606635
636+ self .assertIsNone (
637+ await asyncio .wait_for (on_con_lost , support .LOOPBACK_TIMEOUT ))
638+
639+ with self .tcp_server (serve , timeout = self .TIMEOUT ) as srv :
640+ self .loop .run_until_complete (
641+ asyncio .wait_for (client (srv .addr ),
642+ timeout = support .SHORT_TIMEOUT ))
643+
644+ def _test_start_tls_streams_failed_handshake (self , serve , expected_exc ,
645+ ** kwargs ):
646+ # wait_closed() must not hang if the connection is lost while
647+ # the handshake started by start_tls() is still in progress.
648+ client_context = test_utils .simple_client_sslcontext ()
649+
650+ async def client (addr ):
651+ reader , writer = await asyncio .open_connection (* addr )
652+ with self .assertRaises (expected_exc ):
653+ await writer .start_tls (client_context , ** kwargs )
654+ writer .close ()
655+ with self .assertRaises (expected_exc ):
656+ await writer .wait_closed ()
657+
607658 with self .tcp_server (serve , timeout = self .TIMEOUT ) as srv :
608659 self .loop .run_until_complete (
609660 asyncio .wait_for (client (srv .addr ),
610661 timeout = support .SHORT_TIMEOUT ))
611662
663+ def test_start_tls_streams_handshake_reset (self ):
664+ def serve (sock ):
665+ sock .settimeout (self .TIMEOUT )
666+ sock .recv (1024 * 1024 ) # ClientHello
667+ # Send RST instead of FIN.
668+ sock .setsockopt (socket .SOL_SOCKET , socket .SO_LINGER ,
669+ struct .pack ('ii' , 1 , 0 ))
670+ sock .close ()
671+
672+ self ._test_start_tls_streams_failed_handshake (
673+ serve , ConnectionError )
674+
675+ def test_start_tls_streams_handshake_timeout (self ):
676+ def serve (sock ):
677+ sock .settimeout (self .TIMEOUT )
678+ try :
679+ while sock .recv (1024 * 1024 ):
680+ pass
681+ except ConnectionError :
682+ pass
683+ finally :
684+ sock .close ()
685+
686+ self ._test_start_tls_streams_failed_handshake (
687+ serve , ConnectionAbortedError , ssl_handshake_timeout = 0.1 )
688+
612689 @socket_helper .skip_if_tcp_blackhole
613690 def test_start_tls_server_1 (self ):
614691 HELLO_MSG = b'1' * self .PAYLOAD_SIZE
0 commit comments