fix socketpair EOF

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