multiplexing media through cfproxy

tests/test_h2_routing.py

@@ -0,0 +1,75 @@
+"""H2 belongs to CF fallback, never to the direct route."""
+import unittest
+from types import SimpleNamespace
+from unittest.mock import AsyncMock, Mock, patch
+
+from proxy import bridge, tg_ws_proxy
+from proxy.config import proxy_config
+from proxy.utils import PROTO_TAG_INTERMEDIATE
+
+
+class H2RoutingTest(unittest.IsolatedAsyncioTestCase):
+    async def test_direct_pool_hit_never_opens_h2(self):
+        reader = Mock()
+        writer = Mock()
+        writer.get_extra_info.return_value = None
+        writer.wait_closed = AsyncMock()
+        ws = Mock(send=AsyncMock())
+        pool = SimpleNamespace(open=AsyncMock())
+        with patch.object(proxy_config, 'dc_redirects', {2: '149.154.167.51'}), \
+                patch.object(proxy_config, 'force_test_dc', False), \
+                patch.object(tg_ws_proxy, 'cf_h2_pool', pool), \
+                patch.object(tg_ws_proxy, 'ws_blacklist', set()), \
+                patch.object(tg_ws_proxy, 'ip_fail_until', {}), \
+                patch.object(tg_ws_proxy, 'set_sock_opts'), \
+                patch.object(tg_ws_proxy, '_read_client_init', AsyncMock(
+                    return_value=(b'', reader, writer, 'test'))), \
+                patch.object(tg_ws_proxy, '_try_handshake', return_value=(
+                    2, True, PROTO_TAG_INTERMEDIATE, b'')), \
+                patch.object(tg_ws_proxy, '_build_crypto_ctx'), \
+                patch.object(tg_ws_proxy.ws_pool, 'get', AsyncMock(return_value=ws)), \
+                patch.object(tg_ws_proxy.ws_pool, 'report_success'), \
+                patch.object(tg_ws_proxy, 'bridge_ws_reencrypt', AsyncMock()) as direct, \
+                patch.object(tg_ws_proxy, 'do_fallback', AsyncMock()) as fallback:
+            await tg_ws_proxy._handle_client(reader, writer, b'')
+        direct.assert_awaited_once()
+        fallback.assert_not_awaited()
+        pool.open.assert_not_awaited()
+
+    async def test_cf_fallback_order_and_route_guards(self):
+        cases = [
+            # media, test DC, CF enabled, Worker succeeds, H2 opens, expected
+            (True, False, True, False, True, ['h2', 'bridge']),
+            (True, False, True, False, False, ['h2', 'ws']),
+            (False, False, True, False, True, ['ws']),
+            (True, True, True, False, True, ['tcp']),
+            (True, False, False, False, True, ['tcp']),
+            (True, False, True, True, True, ['worker']),
+        ]
+        for media, test_dc, cf, worker, opens, expected in cases:
+            with self.subTest(expected=expected, media=media, test_dc=test_dc, cf=cf):
+                calls = []
+
+                async def open_h2(*args):
+                    calls.append('h2')
+                    return object() if opens else None
+
+                def route(name):
+                    async def run(*args, **kwargs):
+                        calls.append(name)
+                        return True
+                    return run
+
+                pool = SimpleNamespace(open=AsyncMock(side_effect=open_h2))
+                with patch.object(proxy_config, 'fallback_cfproxy', cf), \
+                        patch.object(proxy_config, 'cfproxy_worker_domains',
+                                     ['worker.example'] if worker else []), \
+                        patch.object(bridge, 'bridge_h2', side_effect=route('bridge')), \
+                        patch.object(bridge, '_cfproxy_fallback', side_effect=route('ws')), \
+                        patch.object(bridge, '_cfproxy_worker_fallback', side_effect=route('worker')), \
+                        patch.object(bridge, '_tcp_fallback', side_effect=route('tcp')):
+                    result = await bridge.do_fallback(
+                        Mock(), Mock(), b'', 'test', 2, test_dc, media, '', Mock(),
+                        h2_pool=pool, proto_tag=PROTO_TAG_INTERMEDIATE)
+                self.assertTrue(result)
+                self.assertEqual(calls, expected)