fix socketpair EOF

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)
+                # wait_for can swallow cancellation when the wakeup completes
-            except asyncio.TimeoutError:
+                # at the same time on Python 3.8-3.11 (CPython #86296).
-                pass
+                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)

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)

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()

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'})

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()

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())