Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 9 additions & 7 deletions Lib/asyncio/sslproto.py
Original file line number Diff line number Diff line change
Expand Up @@ -406,13 +406,15 @@ def connection_lost(self, exc):
if self._app_transport is not None:
self._app_transport._closed = True

if self._state != SSLProtocolState.DO_HANDSHAKE:
if (
self._app_state == AppProtocolState.STATE_CON_MADE or
self._app_state == AppProtocolState.STATE_EOF
):
self._app_state = AppProtocolState.STATE_CON_LOST
self._loop.call_soon(self._app_protocol.connection_lost, exc)
# The app protocol gets connection_lost() only if it got
# connection_made(). After start_tls() that is already the case
# while the handshake is in progress.
if (
self._app_state == AppProtocolState.STATE_CON_MADE or
self._app_state == AppProtocolState.STATE_EOF
):
self._app_state = AppProtocolState.STATE_CON_LOST
self._loop.call_soon(self._app_protocol.connection_lost, exc)
self._set_state(SSLProtocolState.UNWRAPPED)
self._transport = None
self._app_transport = None
Expand Down
81 changes: 79 additions & 2 deletions Lib/test/test_asyncio/test_sslproto.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@

import logging
import socket
import struct
import unittest
import weakref
from test import support
Expand Down Expand Up @@ -153,6 +154,29 @@ def test_connection_lost(self):
test_utils.run_briefly(self.loop)
self.assertIsInstance(waiter.exception(), ConnectionAbortedError)

def test_connection_lost_during_handshake(self):
# The app protocol gets connection_lost() if and only if it got
# connection_made(). With start_tls() (call_connection_made=False)
# it is already connected when the handshake starts.
for call_connection_made in (True, False):
with self.subTest(call_connection_made=call_connection_made):
app_proto = mock.Mock(spec=asyncio.Protocol)
ssl_proto = sslproto.SSLProtocol(
self.loop, app_proto, test_utils.dummy_ssl_context(),
waiter=None, call_connection_made=call_connection_made)
self.connection_made(
ssl_proto,
do_handshake=mock.Mock(side_effect=ssl.SSLWantReadError)
)
exc = ConnectionResetError()
ssl_proto.connection_lost(exc)
test_utils.run_briefly(self.loop)
app_proto.connection_made.assert_not_called()
if call_connection_made:
app_proto.connection_lost.assert_not_called()
else:
app_proto.connection_lost.assert_called_once_with(exc)

def test_connection_lost_when_busy(self):
# gh-118950: SSLProtocol.connection_lost not being called when OSError
# is thrown on asyncio.write.
Expand Down Expand Up @@ -570,9 +594,10 @@ def serve(sock):
sock.close()

class ClientProto(asyncio.Protocol):
def __init__(self, on_data, on_eof):
def __init__(self, on_data, on_eof, on_con_lost):
self.on_data = on_data
self.on_eof = on_eof
self.on_con_lost = on_con_lost
self.con_made_cnt = 0

def connection_made(proto, tr):
Expand All @@ -586,14 +611,18 @@ def data_received(self, data):
def eof_received(self):
self.on_eof.set_result(True)

def connection_lost(self, exc):
self.on_con_lost.set_result(exc)

async def client(addr):
await asyncio.sleep(0.5)

on_data = self.loop.create_future()
on_eof = self.loop.create_future()
on_con_lost = self.loop.create_future()

tr, proto = await self.loop.create_connection(
lambda: ClientProto(on_data, on_eof), *addr)
lambda: ClientProto(on_data, on_eof, on_con_lost), *addr)

tr.write(HELLO_MSG)

Expand All @@ -604,11 +633,59 @@ async def client(addr):
self.loop.start_tls(tr, proto, client_context),
0.5)

self.assertIsNone(
await asyncio.wait_for(on_con_lost, support.LOOPBACK_TIMEOUT))

with self.tcp_server(serve, timeout=self.TIMEOUT) as srv:
self.loop.run_until_complete(
asyncio.wait_for(client(srv.addr),
timeout=support.SHORT_TIMEOUT))

def _test_start_tls_streams_failed_handshake(self, serve, expected_exc,
**kwargs):
# wait_closed() must not hang if the connection is lost while
# the handshake started by start_tls() is still in progress.
client_context = test_utils.simple_client_sslcontext()

async def client(addr):
reader, writer = await asyncio.open_connection(*addr)
with self.assertRaises(expected_exc):
await writer.start_tls(client_context, **kwargs)
writer.close()
with self.assertRaises(expected_exc):
await writer.wait_closed()

with self.tcp_server(serve, timeout=self.TIMEOUT) as srv:
self.loop.run_until_complete(
asyncio.wait_for(client(srv.addr),
timeout=support.SHORT_TIMEOUT))

def test_start_tls_streams_handshake_reset(self):
def serve(sock):
sock.settimeout(self.TIMEOUT)
sock.recv(1024 * 1024) # ClientHello
# Send RST instead of FIN.
sock.setsockopt(socket.SOL_SOCKET, socket.SO_LINGER,
struct.pack('ii', 1, 0))
sock.close()

self._test_start_tls_streams_failed_handshake(
serve, ConnectionError)

def test_start_tls_streams_handshake_timeout(self):
def serve(sock):
sock.settimeout(self.TIMEOUT)
try:
while sock.recv(1024 * 1024):
pass
except ConnectionError:
pass
finally:
sock.close()

self._test_start_tls_streams_failed_handshake(
serve, ConnectionAbortedError, ssl_handshake_timeout=0.1)

@socket_helper.skip_if_tcp_blackhole
def test_start_tls_server_1(self):
HELLO_MSG = b'1' * self.PAYLOAD_SIZE
Expand Down
51 changes: 51 additions & 0 deletions Lib/test/test_asyncio/test_streams.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@

import asyncio
from test.test_asyncio import utils as test_utils
from test import support
from test.support import socket_helper


Expand Down Expand Up @@ -861,6 +862,56 @@ async def run_test():
self.loop.run_until_complete(run_test())
self.assertEqual(messages, [])

@unittest.skipIf(ssl is None, 'No ssl module')
def test_start_tls_connection_lost_during_handshake(self):
# The protocol must get connection_lost() when the connection is
# lost while the start_tls() handshake is in progress, so that
# wait_closed() completes instead of hanging forever.

async def server_handler(client_reader, client_writer):
line = await client_reader.readline()
self.assertEqual(line, b"STARTTLS\n")
client_writer.write(b"go ahead\n")
await client_writer.drain()
with self.assertRaises(ConnectionAbortedError):
await client_writer.start_tls(
test_utils.simple_server_sslcontext(),
ssl_handshake_timeout=0.1)
client_writer.close()
with self.assertRaises(ConnectionAbortedError):
await asyncio.wait_for(client_writer.wait_closed(),
support.SHORT_TIMEOUT)
handler_done.set()

async def client(addr):
reader, writer = await asyncio.open_connection(*addr)
writer.write(b"STARTTLS\n")
await writer.drain()
line = await reader.readline()
self.assertEqual(line, b"go ahead\n")
# Never send a ClientHello: the server's handshake times out.
data = await reader.read()
self.assertEqual(data, b"")
writer.close()
await writer.wait_closed()

async def run_test():
server = await asyncio.start_server(
server_handler, socket_helper.HOSTv4, 0)
server_addr = server.sockets[0].getsockname()

await client(server_addr)
await asyncio.wait_for(handler_done.wait(),
support.SHORT_TIMEOUT)
server.close()
await server.wait_closed()

handler_done = asyncio.Event()
messages = []
self.loop.set_exception_handler(lambda loop, ctx: messages.append(ctx))
self.loop.run_until_complete(run_test())
self.assertEqual(messages, [])

def test_streamwriter_start_tls_updates_reader_transport(self):
reader = asyncio.StreamReader(loop=self.loop)
protocol = asyncio.StreamReaderProtocol(reader, loop=self.loop)
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,5 @@
Fix :meth:`asyncio.loop.start_tls` and :meth:`asyncio.StreamWriter.start_tls`
not calling the protocol's :meth:`~asyncio.BaseProtocol.connection_lost`
when the connection is lost, the handshake times out or the call is
cancelled during the TLS handshake. :meth:`asyncio.StreamWriter.wait_closed`
no longer waits forever after such a failure.
Loading