diff --git a/proxy/fake_tls.py b/proxy/fake_tls.py index 41abcca..153783c 100644 --- a/proxy/fake_tls.py +++ b/proxy/fake_tls.py @@ -284,17 +284,21 @@ async def probe_appdata_len(domain: str, port: int = 443, pass client_hello = outgoing.read() + # Share one deadline across connecting, sending, and reading the flight. + loop = asyncio.get_running_loop() + deadline = loop.time() + timeout reader, writer = await asyncio.wait_for( asyncio.open_connection(domain, port), timeout=timeout) writer.write(client_hello) - await writer.drain() + await asyncio.wait_for(writer.drain(), timeout=deadline - loop.time()) flight = bytearray() - loop = asyncio.get_running_loop() - deadline = loop.time() + timeout while True: + remaining = deadline - loop.time() + if remaining <= 0: + raise asyncio.TimeoutError chunk = await asyncio.wait_for( - reader.read(16384), timeout=max(0.1, deadline - loop.time())) + reader.read(16384), timeout=remaining) if not chunk: return None flight.extend(chunk) diff --git a/tests/test_fake_tls.py b/tests/test_fake_tls.py index f6e3955..e26b21d 100644 --- a/tests/test_fake_tls.py +++ b/tests/test_fake_tls.py @@ -8,6 +8,7 @@ import struct import tempfile import time import unittest +from unittest.mock import AsyncMock, Mock, patch from proxy import fake_tls from proxy.fake_tls import ( @@ -337,6 +338,38 @@ def _self_signed_context() -> ssl.SSLContext: @unittest.skipUnless(getattr(ssl, 'HAS_TLSv1_3', False), 'needs TLS 1.3 in ssl') class ProbeAppdataLenTest(unittest.TestCase): + def test_continuous_records_do_not_extend_probe_timeout(self): + async def scenario(): + async def read(_size): + await asyncio.sleep(0.01) + return _record(TLS_RECORD_CCS, 1) + + reader = Mock(read=AsyncMock(side_effect=read)) + writer = Mock(drain=AsyncMock()) + with patch.object(fake_tls.asyncio, 'open_connection', + AsyncMock(return_value=(reader, writer))): + result = await asyncio.wait_for( + probe_appdata_len('localhost', timeout=0.1), timeout=2) + self.assertIsNone(result) + self.assertGreater(reader.read.await_count, 1) + writer.close.assert_called_once() + + asyncio.run(scenario()) + + def test_blocked_drain_obeys_probe_timeout(self): + async def scenario(): + reader = Mock(read=AsyncMock()) + writer = Mock(drain=AsyncMock(side_effect=asyncio.Event().wait)) + with patch.object(fake_tls.asyncio, 'open_connection', + AsyncMock(return_value=(reader, writer))): + result = await asyncio.wait_for( + probe_appdata_len('localhost', timeout=0.1), timeout=2) + self.assertIsNone(result) + reader.read.assert_not_called() + writer.close.assert_called_once() + + asyncio.run(scenario()) + def test_measures_first_encrypted_record_of_real_tls13_server(self): async def scenario(): server = await asyncio.start_server(