From a0a92f5e7ea43917763577368fc4243c57a37d2c Mon Sep 17 00:00:00 2001 From: Flowseal Date: Sat, 10 Oct 2026 18:54:09 +0300 Subject: [PATCH] fix socketpair EOF --- proxy/cf_h2.py | 10 +- proxy/event_loop.py | 84 +++++++++++ proxy/tg_ws_proxy.py | 2 + tests/test_cf_h2.py | 23 +++ tests/test_event_loop.py | 278 +++++++++++++++++++++++++++++++++++++ tests/test_h2_transport.py | 3 +- 6 files changed, 396 insertions(+), 4 deletions(-) create mode 100644 proxy/event_loop.py create mode 100644 tests/test_event_loop.py diff --git a/proxy/cf_h2.py b/proxy/cf_h2.py index 96da4f6..42c379d 100644 --- a/proxy/cf_h2.py +++ b/proxy/cf_h2.py @@ -525,10 +525,14 @@ class _HttpChannel: async def _recover(self) -> None: while not self.closed: + wakeup = asyncio.create_task(self.recovery_wakeup.wait()) try: - await asyncio.wait_for(self.recovery_wakeup.wait(), REPLAY_CHECK_SECONDS) - except asyncio.TimeoutError: - pass + # wait_for can swallow cancellation when the wakeup completes + # at the same time on Python 3.8-3.11 (CPython #86296). + await asyncio.wait((wakeup,), timeout=REPLAY_CHECK_SECONDS) + finally: + wakeup.cancel() + await asyncio.gather(wakeup, return_exceptions=True) self.recovery_wakeup.clear() now = time.monotonic() self.lane.log_wait(now) diff --git a/proxy/event_loop.py b/proxy/event_loop.py new file mode 100644 index 0000000..2da3b10 --- /dev/null +++ b/proxy/event_loop.py @@ -0,0 +1,84 @@ +import asyncio +import logging +import signal +import socket +import sys +import threading + + +log = logging.getLogger('tg-mtproto-proxy') + + +def _replace_self_pipe(loop): + reader, writer = socket.socketpair() + try: + reader.setblocking(False) + writer.setblocking(False) + if threading.current_thread() is threading.main_thread(): + previous_fd = signal.set_wakeup_fd(writer.fileno()) + if previous_fd != loop._csock.fileno(): + signal.set_wakeup_fd(previous_fd) + except BaseException: + reader.close() + writer.close() + raise + + old_reader, old_writer = loop._ssock, loop._csock + loop._ssock, loop._csock = reader, writer + old_reader.close() + old_writer.close() + + +class _SelfPipeGuard: + """Work around https://github.com/python/cpython/issues/156333""" + + def __init__(self, loop): + self.loop = loop + self.read_self_pipe = loop._loop_self_reading + + def __call__(self, future=None): + loop = self.loop + if future is not None: + if loop._self_reading_future is not future: + return + try: + data = future.result() + except OSError as exc: + reason = str(exc) + except BaseException: + return self.read_self_pipe(future) + else: + if data: + return self.read_self_pipe(future) + reason = 'EOF' + + loop._self_reading_future = None + self._recover(reason) + return + return self.read_self_pipe(future) + + def _recover(self, reason): + loop = self.loop + if loop.is_closed() or loop._self_reading_future is not None: + return + try: + _replace_self_pipe(loop) + except OSError as exc: + log.warning('Windows event loop wakeup socket recovery failed: %s; ' + 'retry in 1s', exc) + loop.call_later(1, self._recover, reason) + return + log.warning('Windows event loop wakeup socket lost (%s); recreated', reason) + self.read_self_pipe() + + +def install_self_pipe_guard(loop): + if sys.platform != 'win32' or not isinstance(loop, asyncio.ProactorEventLoop): + return + if not isinstance(loop._loop_self_reading, _SelfPipeGuard): + guard = _SelfPipeGuard(loop) + loop._loop_self_reading = guard + future = loop._self_reading_future + if future is not None: + future.remove_done_callback(guard.read_self_pipe) + future.add_done_callback(guard) diff --git a/proxy/tg_ws_proxy.py b/proxy/tg_ws_proxy.py index a181983..34e3f5f 100644 --- a/proxy/tg_ws_proxy.py +++ b/proxy/tg_ws_proxy.py @@ -29,6 +29,7 @@ from .balancer import balancer from .pool import ws_pool, cf_worker_pool from .cf_h2 import CfH2Pool from .network_debug import log_ws_flow +from .event_loop import install_self_pipe_guard from ._aes import Cipher, algorithms, modes @@ -360,6 +361,7 @@ _client_tasks: Set[asyncio.Task] = set() async def _run(stop_event: Optional[asyncio.Event] = None): global _server_instance, _server_stop_event, cf_h2_pool + install_self_pipe_guard(asyncio.get_running_loop()) _server_stop_event = stop_event reset_tcp_backoff() diff --git a/tests/test_cf_h2.py b/tests/test_cf_h2.py index 5ac5513..4b7c0d1 100644 --- a/tests/test_cf_h2.py +++ b/tests/test_cf_h2.py @@ -580,6 +580,29 @@ class HttpMultiplexTest(unittest.IsolatedAsyncioTestCase): await healthy.send(b'h' * 40, False) self.assertEqual(await asyncio.wait_for(healthy.receive(), 1), b'r' * 65536) + async def test_recovery_cancellation_is_not_lost_when_wakeup_completes(self): + channel = self._channel() + wait = channel.recovery_wakeup.wait + calls = 0 + + async def wake_and_cancel(): + nonlocal calls + calls += 1 + if calls == 1: + recovery.cancel() + return True + return await wait() + + with patch.object(channel.recovery_wakeup, 'wait', side_effect=wake_and_cancel): + recovery = asyncio.create_task(channel._recover()) + try: + done, _ = await asyncio.wait((recovery,), timeout=.2) + self.assertIn(recovery, done) + self.assertTrue(recovery.cancelled()) + finally: + recovery.cancel() + await asyncio.gather(recovery, return_exceptions=True) + async def test_slow_receiver_queue_overflow_does_not_break_other_channels(self): async def handler(request): return httpx.Response(200, content=request.content, extensions={'http_version': b'HTTP/2'}) diff --git a/tests/test_event_loop.py b/tests/test_event_loop.py new file mode 100644 index 0000000..f0c419f --- /dev/null +++ b/tests/test_event_loop.py @@ -0,0 +1,278 @@ +import asyncio +import signal +import socket +import sys +import threading +import time +import unittest +from unittest.mock import AsyncMock, Mock, patch + +from proxy import event_loop + + +async def _wait_for_replacement(loop, old_writer, timeout=1): + async def wait(): + while loop._csock is old_writer: + await asyncio.sleep(.005) + await asyncio.wait_for(wait(), timeout) + + +def _wake_later(loop, callback): + def send(): + time.sleep(.03) + loop.call_soon_threadsafe(callback) + thread = threading.Thread(target=send, daemon=True) + thread.start() + return thread + + +@unittest.skipUnless(sys.platform == 'win32', 'Windows proactor self-pipe') +class WindowsSelfPipeTest(unittest.TestCase): + def setUp(self): + self.loop = asyncio.ProactorEventLoop() + self.addCleanup(self.loop.close) + + def test_eof_replaces_sockets_and_stops_idle_callbacks(self): + loop = self.loop + read = Mock(wraps=loop._loop_self_reading) + loop._loop_self_reading = read + + async def run(): + # Install after the original callback has armed a read, as can + # happen when the proxy starts in an already running event loop. + await asyncio.sleep(.01) + event_loop.install_self_pipe_guard(loop) + callback = loop._loop_self_reading + event_loop.install_self_pipe_guard(loop) + self.assertIs(loop._loop_self_reading, callback) + old_reader, old_writer = loop._ssock, loop._csock + internal_fds = loop._internal_fds + with patch.object(event_loop.socket, 'socketpair', wraps=socket.socketpair) as pair: + old_writer.shutdown(socket.SHUT_WR) + await _wait_for_replacement(loop, old_writer) + self.assertEqual(pair.call_count, 1) + self.assertEqual(old_reader.fileno(), -1) + self.assertEqual(old_writer.fileno(), -1) + self.assertEqual(loop._internal_fds, internal_fds) + self.assertFalse(loop._self_reading_future.done()) + read.reset_mock() + await asyncio.sleep(.05) + read.assert_not_called() + + completed = loop.create_future() + thread = _wake_later(loop, lambda: completed.set_result(True)) + started = time.monotonic() + try: + self.assertTrue(await asyncio.wait_for(completed, 1)) + self.assertLess(time.monotonic() - started, .5) + finally: + thread.join(timeout=1) + + with self.assertLogs(event_loop.log, 'WARNING') as messages: + loop.run_until_complete(run()) + self.assertEqual(len(messages.output), 1) + self.assertIn('EOF', messages.output[0]) + + def test_repeated_eof_preserves_existing_tcp_connection(self): + async def run(): + loop = self.loop + event_loop.install_self_pipe_guard(loop) + finished = loop.create_future() + + async def echo(reader, writer): + try: + while True: + data = await reader.read(128) + if not data: + break + writer.write(data) + await writer.drain() + finally: + writer.close() + await writer.wait_closed() + finished.set_result(None) + + server = await asyncio.start_server(echo, '127.0.0.1', 0) + writer = None + try: + port = server.sockets[0].getsockname()[1] + reader, writer = await asyncio.open_connection('127.0.0.1', port) + for index in range(5): + old_writer = loop._csock + old_writer.shutdown(socket.SHUT_WR) + await _wait_for_replacement(loop, old_writer) + payload = f'packet {index}'.encode() + writer.write(payload) + await writer.drain() + data = await asyncio.wait_for(reader.readexactly(len(payload)), 1) + self.assertEqual(data, payload) + finally: + if writer is not None: + writer.close() + await writer.wait_closed() + await asyncio.wait_for(finished, 1) + server.close() + await server.wait_closed() + + with self.assertLogs(event_loop.log, 'WARNING'): + self.loop.run_until_complete(run()) + + def test_reset_rebuilds_but_cancelled_and_stale_callbacks_do_not(self): + async def run(): + loop = self.loop + event_loop.install_self_pipe_guard(loop) + await asyncio.sleep(0) + old_writer = loop._csock + pending = loop._self_reading_future + pending.cancel() + failed = loop.create_future() + failed.set_exception(ConnectionResetError('self-pipe reset')) + loop._self_reading_future = failed + loop._loop_self_reading(failed) + self.assertIsNot(loop._csock, old_writer) + replacement = loop._csock + with patch.object(loop, 'call_exception_handler') as report: + loop._loop_self_reading(failed) # queued callback from the old pipe + loop._loop_self_reading(pending) # cancellation during shutdown + report.assert_not_called() + self.assertIs(loop._csock, replacement) + await asyncio.sleep(.01) + + with self.assertLogs(event_loop.log, 'WARNING') as messages: + self.loop.run_until_complete(run()) + self.assertEqual(len(messages.output), 1) + + def test_reset_on_read_armed_before_installation_is_recovered(self): + async def run(): + loop = self.loop + await asyncio.sleep(.01) + event_loop.install_self_pipe_guard(loop) + old_writer = loop._csock + loop._self_reading_future.set_exception(ConnectionResetError('self-pipe reset')) + await _wait_for_replacement(loop, old_writer) + self.assertFalse(loop._self_reading_future.done()) + + with self.assertLogs(event_loop.log, 'WARNING') as messages: + self.loop.run_until_complete(run()) + self.assertEqual(len(messages.output), 1) + + def test_recovery_in_background_proxy_thread(self): + errors = [] + + async def run(): + loop = asyncio.get_running_loop() + event_loop.install_self_pipe_guard(loop) + await asyncio.sleep(0) + old_writer = loop._csock + old_writer.shutdown(socket.SHUT_WR) + await _wait_for_replacement(loop, old_writer) + self.assertFalse(loop._self_reading_future.done()) + + def worker(): + try: + asyncio.run(run()) + except BaseException as exc: + errors.append(exc) + + with self.assertLogs(event_loop.log, 'WARNING'): + thread = threading.Thread(target=worker, daemon=True) + thread.start() + thread.join(timeout=2) + self.assertFalse(thread.is_alive()) + self.assertFalse(errors, errors) + + def test_allocation_failure_retries_without_spinning(self): + async def run(): + loop = self.loop + event_loop.install_self_pipe_guard(loop) + await asyncio.sleep(0) + original_pair = socket.socketpair + attempts = 0 + + def allocate(): + nonlocal attempts + attempts += 1 + if attempts == 1: + raise OSError('temporarily out of sockets') + return original_pair() + + with patch.object(event_loop.socket, 'socketpair', side_effect=allocate): + old_reader, old_writer = loop._ssock, loop._csock + old_writer.shutdown(socket.SHUT_WR) + await asyncio.sleep(.05) + self.assertEqual(attempts, 1) + self.assertIs(loop._csock, old_writer) + self.assertNotEqual(old_reader.fileno(), -1) + self.assertIsNone(loop._self_reading_future) + await _wait_for_replacement(loop, old_writer, timeout=2) + self.assertEqual(attempts, 2) + + with self.assertLogs(event_loop.log, 'WARNING'): + self.loop.run_until_complete(run()) + + def test_signal_wakeup_fd_is_moved_only_when_owned_by_the_loop(self): + async def run(): + loop = self.loop + event_loop.install_self_pipe_guard(loop) + await asyncio.sleep(0) + external_reader, external_writer = socket.socketpair() + external_writer.setblocking(False) + try: + for external in (False, True): + with self.subTest(external=external): + if external: + signal.set_wakeup_fd(external_writer.fileno()) + old_writer = loop._csock + old_writer.shutdown(socket.SHUT_WR) + await _wait_for_replacement(loop, old_writer) + fd = signal.set_wakeup_fd(-1) + expected = external_writer if external else loop._csock + self.assertEqual(fd, expected.fileno()) + signal.set_wakeup_fd(loop._csock.fileno()) + finally: + external_reader.close() + external_writer.close() + + with self.assertLogs(event_loop.log, 'WARNING'): + self.loop.run_until_complete(run()) + + +@unittest.skipUnless(sys.platform == 'win32', 'Windows proactor self-pipe') +class ProxySelfPipeTest(unittest.IsolatedAsyncioTestCase): + async def test_proxy_without_h2_can_stop_from_thread_after_eof(self): + from proxy import tg_ws_proxy + from proxy.config import ProxyConfig + + config = ProxyConfig(port=0, dc_redirects={}, pool_size=0, cfproxy_h2_media=False) + stop = asyncio.Event() + loop = asyncio.get_running_loop() + with patch.object(tg_ws_proxy, 'proxy_config', config), \ + patch.object(tg_ws_proxy, 'start_cfproxy_domain_refresh'), \ + patch.object(tg_ws_proxy.ws_pool, 'warmup', AsyncMock()), \ + patch.object(tg_ws_proxy.cf_worker_pool, 'warmup', AsyncMock()), \ + self.assertLogs(event_loop.log, 'WARNING'): + task = asyncio.create_task(tg_ws_proxy._run(stop)) + thread = None + try: + async def listening(): + while tg_ws_proxy._server_instance is None: + await asyncio.sleep(.005) + await asyncio.wait_for(listening(), 1) + self.assertIsNone(tg_ws_proxy.cf_h2_pool) + old_writer = loop._csock + old_writer.shutdown(socket.SHUT_WR) + await _wait_for_replacement(loop, old_writer) + thread = _wake_later(loop, stop.set) + started = time.monotonic() + await asyncio.wait_for(task, 1) + self.assertLess(time.monotonic() - started, .5) + self.assertIsNone(tg_ws_proxy._server_instance) + finally: + if thread is not None: + thread.join(timeout=1) + task.cancel() + await asyncio.gather(task, return_exceptions=True) + + +if __name__ == '__main__': + unittest.main() diff --git a/tests/test_h2_transport.py b/tests/test_h2_transport.py index 41742e2..c7d368b 100644 --- a/tests/test_h2_transport.py +++ b/tests/test_h2_transport.py @@ -1082,7 +1082,8 @@ class H2WireTest(unittest.IsolatedAsyncioTestCase): self.assertEqual(len(self.connections), 1) failures = [record for record in captured.records if record.levelname in ('WARNING', 'ERROR')] self.assertEqual(len(failures), 1) - self.assertIn('reset: 2', failures[0].getMessage()) + # h2 4.1 (Python 3.8) formats IntEnum values by name. + self.assertRegex(failures[0].getMessage(), r'reset: (?:2|ErrorCodes\.INTERNAL_ERROR)\b') self.assertIn('down=20', failures[0].getMessage()) self.assertIn('sid=1', failures[0].getMessage())