Add a test for proper waiter handling in start_tls

Author: 1st1

Lib/test/test_asyncio/test_sslproto.py

@@ -359,6 +359,67 @@ async def client(addr):
                 asyncio.wait_for(client(srv.addr),
                                  loop=self.loop, timeout=self.TIMEOUT))
 
+    def test_start_tls_slow_client_cancel(self):
+        HELLO_MSG = b'1' * self.PAYLOAD_SIZE
+
+        client_context = test_utils.simple_client_sslcontext()
+        server_waits_on_handshake = self.loop.create_future()
+
+        def serve(sock):
+            sock.settimeout(self.TIMEOUT)
+
+            data = sock.recv_all(len(HELLO_MSG))
+            self.assertEqual(len(data), len(HELLO_MSG))
+
+            try:
+                self.loop.call_soon_threadsafe(
+                    server_waits_on_handshake.set_result, None)
+                data = sock.recv_all(1024 * 1024)
+            except ConnectionAbortedError:
+                pass
+            finally:
+                sock.close()
+
+        class ClientProto(asyncio.Protocol):
+            def __init__(self, on_data, on_eof):
+                self.on_data = on_data
+                self.on_eof = on_eof
+                self.con_made_cnt = 0
+
+            def connection_made(proto, tr):
+                proto.con_made_cnt += 1
+                # Ensure connection_made gets called only once.
+                self.assertEqual(proto.con_made_cnt, 1)
+
+            def data_received(self, data):
+                self.on_data.set_result(data)
+
+            def eof_received(self):
+                self.on_eof.set_result(True)
+
+        async def client(addr):
+            await asyncio.sleep(0.5, loop=self.loop)
+
+            on_data = self.loop.create_future()
+            on_eof = self.loop.create_future()
+
+            tr, proto = await self.loop.create_connection(
+                lambda: ClientProto(on_data, on_eof), *addr)
+
+            tr.write(HELLO_MSG)
+
+            await server_waits_on_handshake
+
+            with self.assertRaises(asyncio.TimeoutError):
+                await asyncio.wait_for(
+                    self.loop.start_tls(tr, proto, client_context),
+                    0.5,
+                    loop=self.loop)
+
+        with self.tcp_server(serve, timeout=self.TIMEOUT) as srv:
+            self.loop.run_until_complete(
+                asyncio.wait_for(client(srv.addr), loop=self.loop, timeout=10))
+
     def test_start_tls_server_1(self):
         HELLO_MSG = b'1' * self.PAYLOAD_SIZE