import asyncio import time import unittest from collections import deque from types import SimpleNamespace from unittest import mock from proxy.config import proxy_config from proxy.pool import _CfWorkerPool, _WsPool from proxy.raw_websocket import RawWebSocket, WsHandshakeError from proxy.utils import DC_DEFAULT_IPS, ws_domains class _StopRotation(Exception): pass def _open_ws(): transport = SimpleNamespace(is_closing=lambda: False) writer = mock.Mock(transport=transport, drain=mock.AsyncMock(), wait_closed=mock.AsyncMock()) return RawWebSocket(asyncio.StreamReader(), writer) class WsPoolRotationTest(unittest.IsolatedAsyncioTestCase): async def test_warmup_preserves_media_ws_with_or_without_h2(self): for cf, secure, opted_out, media in [(True, True, False, True), (False, True, False, True), (True, False, False, True), (True, True, True, True)]: with self.subTest(cf=cf, secure=secure, opted_out=opted_out), \ mock.patch.object(proxy_config, 'dc_redirects', {2: '149.154.167.220'}), \ mock.patch.object(proxy_config, 'fallback_cfproxy', cf), \ mock.patch.object(proxy_config, 'disable_secure', not secure), \ mock.patch.object(proxy_config, 'cfproxy_h2_media', not opted_out), \ mock.patch.object(proxy_config, 'force_test_dc', False): pool = _WsPool() with mock.patch.object(pool, '_schedule_refill') as refill: await pool.warmup() keys = [call.args[0] for call in refill.call_args_list] self.assertIn((2, False, False), keys) self.assertEqual((2, True, False) in keys, media) async def test_refills_partially_populated_bucket(self): pool = _WsPool() key = (2, False, False) pool._idle[key] = deque([ (_open_ws(), time.monotonic()), (_open_ws(), time.monotonic()), ]) async def stop_after_one_iteration(_delay): raise _StopRotation with mock.patch.object(proxy_config, 'pool_size', 4): with mock.patch( 'proxy.pool.asyncio.sleep', side_effect=stop_after_one_iteration): with mock.patch.object( pool, '_schedule_refill') as schedule_refill: with self.assertRaises(_StopRotation): await pool._rotate( key, '149.154.167.220', ['example.com']) schedule_refill.assert_called_once_with( key, '149.154.167.220', ['example.com']) class WsPoolTest(unittest.IsolatedAsyncioTestCase): async def asyncSetUp(self): self.pool = _WsPool() self.key = (2, False, False) self.pool.WS_POOL_CHECK_INTERVAL = .01 for name, value in [('dc_redirects', {2: '192.0.2.1'}), ('pool_size', 1), ('force_test_dc', False)]: patcher = mock.patch.object(proxy_config, name, value) patcher.start() self.addCleanup(patcher.stop) self.addAsyncCleanup(self.pool.close) async def wait_for(self, predicate): async def wait(): while not predicate(): await asyncio.sleep(.001) await asyncio.wait_for(wait(), 1) async def test_regular_sessions_never_use_media_domain(self): for dc in (2, 4): with self.subTest(dc=dc), mock.patch( 'proxy.pool.RawWebSocket.connect', side_effect=WsHandshakeError(302, 'redirect')) as connect: self.assertIsNone(await self.pool._connect_one('192.0.2.1', ws_domains(dc, False))) self.assertEqual([call.args[1] for call in connect.call_args_list], [f'kws{dc}.web.telegram.org'] * 2) async def test_media_tries_regular_domain_after_both_modes_fail(self): for error in (asyncio.TimeoutError(), ConnectionResetError(), WsHandshakeError(302, 'redirect'), WsHandshakeError(500, 'error')): with self.subTest(error=type(error)), mock.patch( 'proxy.pool.RawWebSocket.connect', side_effect=[error, error, _open_ws()]) as connect: self.assertIsNotNone(await self.pool._connect_one('192.0.2.1', ws_domains(2, True))) self.assertEqual([call.args[1] for call in connect.call_args_list], ['kws2-1.web.telegram.org'] * 2 + ['kws2.web.telegram.org']) async def test_concurrent_preference_change_does_not_skip_fronting(self): ws = _open_ws() async def connect(*args, **kwargs): if kwargs['sni'] is None: self.pool.try_fronting_first = True raise asyncio.TimeoutError() return ws with mock.patch('proxy.pool.RawWebSocket.connect', side_effect=connect) as dial: self.assertIs(await self.pool._connect_one('192.0.2.1', ws_domains(2, False)), ws) self.assertEqual(dial.await_count, 2) async def test_miss_returns_without_waiting_and_starts_only_one_refill(self): started = asyncio.Event() async def connect(*args): started.set() await asyncio.Future() with mock.patch.object(self.pool, '_connect_one', side_effect=connect) as dial: for _ in range(10): self.assertIsNone(await self.pool.get(2, False)) await asyncio.wait_for(started.wait(), 1) self.assertEqual(dial.await_count, 1) await self.pool.close() async def test_empty_pool_recovers_in_background_after_backoff(self): ws = _open_ws() self.pool.REFILL_BACKOFF_INITIAL = .02 with mock.patch.object(self.pool, '_connect_one', side_effect=[None, ws]) as dial: self.assertIsNone(await self.pool.get(2, False)) await self.wait_for(lambda: self.key in self.pool._refill_after) self.assertEqual(dial.await_count, 1) # No new client requests are needed to resume refilling. await self.wait_for(lambda: bool(self.pool._idle[self.key])) self.assertEqual(dial.await_count, 2) self.assertNotIn(self.key, self.pool._refill_failures) await self.pool.close() async def test_ready_connection_is_available_before_slower_attempt(self): ready = _open_ws() slow_started = asyncio.Event() calls = 0 async def connect(*args): nonlocal calls calls += 1 if calls == 1: slow_started.set() await asyncio.Future() return ready with mock.patch.object(proxy_config, 'pool_size', 2), \ mock.patch.object(self.pool, '_connect_one', side_effect=connect): await self.pool.get(2, False) await asyncio.wait_for(slow_started.wait(), 1) await self.wait_for(lambda: bool(self.pool._idle[self.key])) self.assertIs(await self.pool.get(2, False), ready) await self.pool.close() async def test_stale_connections_are_skipped(self): for state in ('expired', 'eof', 'exception', 'closed', 'closing'): with self.subTest(state=state): stale, good = _open_ws(), _open_ws() created = time.monotonic() if state == 'expired': created -= self.pool.WS_POOL_MAX_AGE elif state == 'eof': stale.reader.feed_eof() elif state == 'exception': stale.reader.set_exception(ConnectionResetError()) elif state == 'closed': stale._closed = True else: stale.writer.transport.is_closing = lambda: True self.pool._idle[self.key] = deque([(stale, created), (good, time.monotonic())]) with mock.patch.object(self.pool, '_schedule_refill'): self.assertIs(await self.pool.get(2, False), good) await asyncio.sleep(0) async def test_hit_does_not_reset_failed_refill_backoff(self): self.pool._idle[self.key] = deque([(_open_ws(), time.monotonic())]) self.pool._refill_failures[self.key] = 3 self.pool._refill_after[self.key] = time.monotonic() + 60 with mock.patch.object(self.pool, '_schedule_rotation'): self.assertIsNotNone(await self.pool.get(2, False)) self.assertEqual(self.pool._refill_failures[self.key], 3) self.assertFalse(self.pool._refilling) async def test_zero_size_and_unconfigured_dc_do_not_connect(self): with mock.patch.object(self.pool, '_connect_one') as dial: self.assertIsNone(await self.pool.get(4, False)) with mock.patch.object(proxy_config, 'pool_size', 0): await self.pool.warmup() self.assertIsNone(await self.pool.get(2, False)) self.assertFalse(self.pool._rotating) self.assertFalse(self.pool._refilling) dial.assert_not_called() async def test_test_dc_has_separate_pool_and_path(self): prod = _open_ws() self.pool._idle[self.key] = deque([(prod, time.monotonic())]) with mock.patch.object(self.pool, '_connect_one', return_value=_open_ws()) as dial: self.assertIsNone(await self.pool.get(2, False, is_test_dc=True)) test_key = (2, False, True) await self.wait_for(lambda: bool(self.pool._idle.get(test_key))) dial.assert_awaited_once_with('192.0.2.1', ['kws2.web.telegram.org'], '/apiws_test') self.assertIs(self.pool._idle[self.key][0][0], prod) await self.pool.close() async def test_close_cancels_refill_and_closes_unclaimed_connections(self): ready = _open_ws() idle = _open_ws() started = asyncio.Event() self.pool._idle[self.key] = deque([(idle, time.monotonic())]) async def connect(*args): started.set() try: await asyncio.Future() except asyncio.CancelledError: # Completion raced with shutdown; it must not leak into the pool. return ready with mock.patch.object(proxy_config, 'pool_size', 2), \ mock.patch.object(self.pool, '_connect_one', side_effect=connect): self.pool._schedule_refill(self.key, '192.0.2.1', ws_domains(2, False)) await asyncio.wait_for(started.wait(), 1) await self.pool.close() self.assertFalse(self.pool._idle) self.assertFalse(self.pool._rotating) self.assertFalse(self.pool._refilling) ready.writer.close.assert_called_once() idle.writer.close.assert_called_once() class CfWorkerPoolTest(unittest.IsolatedAsyncioTestCase): async def asyncSetUp(self): self.pool = _CfWorkerPool() self.pool.WS_POOL_CHECK_INTERVAL = .01 for name, value in [ ('pool_size', 4), ('dc_redirects', {dc: ip for dc, ip in DC_DEFAULT_IPS.items() if dc != 2}), ('cfproxy_worker_domains', ['worker.example'])]: patcher = mock.patch.object(proxy_config, name, value) patcher.start() self.addCleanup(patcher.stop) self.addAsyncCleanup(self.pool.close) async def wait_for(self, predicate): async def wait(): while not predicate(): await asyncio.sleep(.001) await asyncio.wait_for(wait(), 1) async def test_failed_warmup_recovers_without_new_client_requests(self): ready = _open_ws() self.pool.REFILL_BACKOFF_INITIAL = .03 with mock.patch.object(self.pool, '_connect_one', side_effect=[None, (ready, 'worker.example')]) as dial: await self.pool.warmup() await self.wait_for(lambda: 2 in self.pool._refill_after) self.assertFalse(self.pool._refilling) await self.wait_for(lambda: bool(self.pool._idle.get(2))) self.assertEqual(dial.await_count, 2) self.assertIs(self.pool._idle[2][0][0], ready) self.assertNotIn(2, self.pool._refill_failures) self.assertNotIn(2, self.pool._refill_after) async def test_stale_connections_are_replaced(self): for state in ('expired', 'eof', 'exception', 'closed', 'closing'): with self.subTest(state=state): stale, ready = _open_ws(), _open_ws() created = time.monotonic() if state == 'expired': created -= self.pool.WS_POOL_MAX_AGE elif state == 'eof': stale.reader.feed_eof() elif state == 'exception': stale.reader.set_exception(ConnectionResetError()) elif state == 'closed': stale._closed = True else: stale.writer.transport.is_closing = lambda: True self.pool._idle[2] = deque([(stale, created, 'worker.example')]) with mock.patch.object(self.pool, '_connect_one', return_value=(ready, 'worker.example')): self.assertIsNone(await self.pool.get(2, '192.0.2.2', ['worker.example'])) await self.wait_for(lambda: bool(self.pool._idle[2])) self.assertIs(self.pool._idle[2][0][0], ready) stale.writer.close.assert_called_once() await self.pool.close() async def test_rotation_replaces_idle_connection_without_client_requests(self): stale, ready = _open_ws(), _open_ws() self.pool.WS_POOL_MAX_AGE = .05 self.pool._idle[2] = deque([(stale, time.monotonic(), 'worker.example')]) with mock.patch.object(self.pool, '_connect_one', return_value=(ready, 'worker.example')) as dial: await self.pool.warmup() await self.wait_for(lambda: stale.writer.close.called) self.assertIs(self.pool._idle[2][0][0], ready) stale.writer.close.assert_called_once() dial.assert_awaited_once() async def test_concurrent_misses_start_only_one_background_connection(self): started = asyncio.Event() async def connect(*args): started.set() await asyncio.Future() with mock.patch.object(self.pool, '_connect_one', side_effect=connect) as dial: for _ in range(10): self.assertIsNone(await self.pool.get(2, '192.0.2.2', ['worker.example'])) await asyncio.wait_for(started.wait(), 1) self.assertEqual(dial.await_count, 1) async def test_failed_refills_use_capped_exponential_backoff(self): ready = _open_ws() self.pool.REFILL_BACKOFF_INITIAL = 1 self.pool.REFILL_BACKOFF_MAX = 2 with mock.patch.object(self.pool, '_connect_one', side_effect=[None, None, None, (ready, 'worker.example')]), \ mock.patch.object(self.pool, '_schedule_rotation'): for failures, delay in enumerate((1, 2, 2), 1): self.pool._refill_after.pop(2, None) started = time.monotonic() self.pool._schedule_refill(2, '192.0.2.2', ['worker.example']) await asyncio.gather(*self.pool._refilling.values()) self.assertFalse(self.pool._refilling) self.assertEqual(self.pool._refill_failures[2], failures) self.assertGreaterEqual(self.pool._refill_after[2], started + delay) self.assertLessEqual(self.pool._refill_after[2], time.monotonic() + delay) self.pool._refill_after.pop(2) self.pool._schedule_refill(2, '192.0.2.2', ['worker.example']) await asyncio.gather(*self.pool._refilling.values()) self.assertIs(self.pool._idle[2][0][0], ready) self.assertNotIn(2, self.pool._refill_failures) self.assertNotIn(2, self.pool._refill_after) async def test_misses_during_backoff_do_not_retry_and_shutdown_cancels_retry(self): self.pool.REFILL_BACKOFF_INITIAL = 60 with mock.patch.object(self.pool, '_connect_one', return_value=None) as dial: await self.pool.warmup() await self.wait_for(lambda: 2 in self.pool._refill_after) for _ in range(10): self.assertIsNone(await self.pool.get(2, '192.0.2.2', ['worker.example'])) await asyncio.sleep(.01) self.assertEqual(dial.await_count, 1) self.assertFalse(self.pool._refilling) self.assertTrue(self.pool._rotating) tasks = list(self.pool._rotating.values()) await self.pool.close() self.assertFalse(self.pool._refilling) self.assertFalse(self.pool._rotating) self.assertFalse(self.pool._refill_after) self.assertFalse(self.pool._refill_failures) self.assertTrue(all(task.done() for task in tasks)) async def test_hit_does_not_reset_failed_refill_backoff(self): ready = _open_ws() self.pool._idle[2] = deque([(ready, time.monotonic(), 'worker.example')]) self.pool._refill_failures[2] = 3 retry_at = time.monotonic() + 60 self.pool._refill_after[2] = retry_at with mock.patch.object(self.pool, '_schedule_rotation'): self.assertEqual(await self.pool.get(2, '192.0.2.2', ['worker.example']), (ready, 'worker.example')) self.assertEqual(self.pool._refill_failures[2], 3) self.assertEqual(self.pool._refill_after[2], retry_at) self.assertFalse(self.pool._refilling) async def test_zero_size_and_missing_worker_domains_do_not_start_refills(self): with mock.patch.object(self.pool, '_connect_one') as dial: with mock.patch.object(proxy_config, 'pool_size', 0): await self.pool.warmup() self.assertIsNone(await self.pool.get(2, '192.0.2.2', ['worker.example'])) self.assertFalse(self.pool._refilling) self.assertFalse(self.pool._rotating) self.assertIsNone(await self.pool.get(2, '192.0.2.2', [])) self.assertFalse(self.pool._refilling) self.assertFalse(self.pool._rotating) dial.assert_not_called() async def test_close_cancels_refill_and_closes_unclaimed_connections(self): pool = _CfWorkerPool() ready = _open_ws() idle = _open_ws() started = asyncio.Event() pool._idle[2] = deque([(idle, time.monotonic(), 'a.example')]) async def connect(*args): started.set() try: await asyncio.Future() except asyncio.CancelledError: return ready, 'b.example' with mock.patch.object(proxy_config, 'pool_size', 1), \ mock.patch.object(pool, '_connect_one', side_effect=connect): pool._schedule_refill(4, '192.0.2.4', ['b.example']) await asyncio.wait_for(started.wait(), 1) await pool.close() self.assertFalse(pool._idle) self.assertFalse(pool._refilling) self.assertFalse(pool._rotating) self.assertFalse(pool._refill_after) self.assertFalse(pool._refill_failures) ready.writer.close.assert_called_once() idle.writer.close.assert_called_once() if __name__ == '__main__': unittest.main()