multiplexing media through cfproxy

tests/test_bridge.py

@@ -1,8 +1,12 @@
 import os
+import asyncio
+import time
 import unittest
+from types import SimpleNamespace
 
 from proxy._aes import Cipher, algorithms, modes
-from proxy.bridge import MsgSplitter
+from proxy.bridge import MsgSplitter, bridge_ws_reencrypt
+from proxy.network_debug import WsActivity, log_ws_flow
 from proxy.utils import (
     PROTO_ABRIDGED_INT,
     PROTO_INTERMEDIATE_INT,
@@ -112,5 +116,90 @@ class MsgSplitterTest(unittest.TestCase):
         self.assertEqual(splitter.flush(), [])
 
 
+class WsDiagnosticsTest(unittest.IsolatedAsyncioTestCase):
+    async def test_flow_distinguishes_upstream_send_from_native_backpressure(self):
+        class Cipher:
+            def update(self, data):
+                return data
+
+        class Socket:
+            def __init__(self):
+                self.queue = asyncio.Queue()
+                self.sending = asyncio.Event()
+                self.release = asyncio.Event()
+                self.sent = []
+
+            async def send(self, data):
+                self.sent.append(data)
+                self.sending.set()
+                await self.release.wait()
+
+            async def recv(self):
+                return await self.queue.get()
+
+            async def close(self):
+                pass
+
+        class Writer:
+            def __init__(self):
+                self.data = bytearray()
+                self.written = asyncio.Event()
+                self.release = asyncio.Event()
+
+            def write(self, data):
+                self.data.extend(data)
+                self.written.set()
+
+            async def drain(self):
+                await self.release.wait()
+
+            def close(self):
+                pass
+
+            async def wait_closed(self):
+                pass
+
+        ws, writer = Socket(), Writer()
+        reader = asyncio.StreamReader()
+        ctx = SimpleNamespace(clt_dec=Cipher(), clt_enc=Cipher(), tg_dec=Cipher(), tg_enc=Cipher())
+        with self.assertLogs('tg-mtproto-proxy', level='DEBUG') as captured:
+            bridge = asyncio.create_task(bridge_ws_reencrypt(reader, writer, ws, 'owned-test', ctx, dc=203))
+            try:
+                reader.feed_data(b'private-native-payload')
+                await asyncio.wait_for(ws.sending.wait(), 1)
+                activity = next(item for item in WsActivity.active if item.label == 'owned-test')
+                self.assertIsNotNone(activity.sending)
+                self.assertEqual(activity.ws_up, 0)
+                log_ws_flow(time.monotonic())
+                # A reply arriving before send drain finishes must not create
+                # a spurious wait for another response when drain resumes.
+                await ws.queue.put(b'private-upstream-payload')
+                await asyncio.wait_for(writer.written.wait(), 1)
+                self.assertEqual(activity.ws_down, len(writer.data))
+                self.assertEqual(activity.native_down, 0)
+                self.assertIsNotNone(activity.writing_native)
+                log_ws_flow(time.monotonic())
+                ws.release.set()
+                writer.release.set()
+                for _ in range(100):
+                    if activity.native_down and activity.ws_up:
+                        break
+                    await asyncio.sleep(.001)
+                self.assertEqual(activity.native_down, len(writer.data))
+                self.assertIsNone(activity.awaiting_rx)
+                self.assertEqual(ws.sent, [b'private-native-payload'])
+                self.assertEqual(bytes(writer.data), b'private-upstream-payload')
+                reader.feed_eof()
+                await asyncio.wait_for(bridge, 1)
+            finally:
+                bridge.cancel()
+                await asyncio.gather(bridge, return_exceptions=True)
+        self.assertFalse(WsActivity.active)
+        output = '\n'.join(captured.output)
+        self.assertIn('WS FLOW DC203', output)
+        self.assertIn('WS END DC203', output)
+        self.assertNotIn('private-', output)
+
+
 if __name__ == '__main__':
     unittest.main()

tests/test_cf_h2.py

@@ -0,0 +1,912 @@
+import asyncio
+import inspect
+import struct
+import time
+import unittest
+from types import SimpleNamespace
+from unittest.mock import patch
+
+import httpx
+
+from proxy.balancer import _Balancer
+from proxy.cf_h2 import (
+    CfH2Pool, MAX_CHANNEL_REQUESTS, MAX_PACKET, _HttpChannel, _HttpLane,
+    _ReplayPacket, _encode_reply, _read_packet, bridge_h2,
+)
+from proxy.utils import PROTO_TAG_ABRIDGED, PROTO_TAG_INTERMEDIATE, PROTO_TAG_SECURE
+
+class _IdentityCipher:
+    def update(self, value):
+        return value
+
+
+class NativeFramingTest(unittest.IsolatedAsyncioTestCase):
+    async def _packet(self, wire, tag):
+        reader = asyncio.StreamReader()
+        async def feed():
+            for byte in wire:
+                reader.feed_data(bytes([byte]))
+                await asyncio.sleep(0)
+            reader.feed_eof()
+        task = asyncio.create_task(feed())
+        try:
+            return await _read_packet(reader, _IdentityCipher(), tag)
+        finally:
+            task.cancel()
+            await asyncio.gather(task, return_exceptions=True)
+
+    async def test_fragmented_abridged_extended_and_quick_ack(self):
+        body = b'a' * 512
+        decoded, quick = await self._packet(b'\xff\x80\x00\x00' + body, PROTO_TAG_ABRIDGED)
+        self.assertEqual(decoded, body)
+        self.assertTrue(quick)
+
+    async def test_intermediate_quick_ack_does_not_change_body(self):
+        body = b'a' * 40
+        decoded, quick = await self._packet(struct.pack('<I', len(body) | 0x80000000) + body,
+                                            PROTO_TAG_INTERMEDIATE)
+        self.assertEqual(decoded, body)
+        self.assertTrue(quick)
+
+    async def test_padded_plaintext_removes_all_fifteen_padding_lengths(self):
+        body = b'\x00' * 8 + b'm' * 8 + struct.pack('<I', 20) + b'p' * 20
+        for padding in range(16):
+            wire = body + b'z' * padding
+            decoded, quick = await self._packet(struct.pack('<I', len(wire)) + wire, PROTO_TAG_SECURE)
+            self.assertEqual(decoded, body, padding)
+            self.assertFalse(quick)
+
+    async def test_padded_ciphertext_removes_full_padding_blocks(self):
+        body = b'k' * 8 + b'm' * 16 + b'c' * 32
+        for padding in range(16):
+            wire = body + b'z' * padding
+            decoded, _ = await self._packet(struct.pack('<I', len(wire)) + wire, PROTO_TAG_SECURE)
+            self.assertEqual(decoded, body, padding)
+
+    async def test_oversize_header_rejected_before_body_read(self):
+        reader = asyncio.StreamReader()
+        reader.feed_data(struct.pack('<I', MAX_PACKET + 4))
+        with self.assertRaisesRegex(ValueError, 'length'):
+            await _read_packet(reader, _IdentityCipher(), PROTO_TAG_INTERMEDIATE)
+
+    def test_negative_transport_error_keeps_native_framing(self):
+        error = struct.pack('<i', -404)
+        self.assertEqual(_encode_reply(error, PROTO_TAG_ABRIDGED), b'\x01' + error)
+        self.assertEqual(_encode_reply(error, PROTO_TAG_INTERMEDIATE), b'\x04\x00\x00\x00' + error)
+
+
+class HttpMultiplexTest(unittest.IsolatedAsyncioTestCase):
+    async def asyncSetUp(self):
+        self.lane = _HttpLane('kws4.example.org', 1)
+        await self.lane.client.aclose()
+        self.channels = []
+
+    async def asyncTearDown(self):
+        await self.lane.close()
+
+    def _channel(self):
+        channel = _HttpChannel(self.lane, len(self.channels) + 1, 'owned-test')
+        self.channels.append(channel)
+        return channel
+
+    def _transport(self, handler, *, trace_send=True):
+        async def handle(request):
+            if trace_send and 'trace' in request.extensions:
+                # MockTransport receives the complete body before calling us.
+                await request.extensions['trace']('http2.send_request_body.complete', {'stream_id': 3})
+            result = handler(request)
+            return await result if inspect.isawaitable(result) else result
+        self.lane.client = httpx.AsyncClient(http2=True, transport=httpx.MockTransport(handle))
+
+    async def test_same_auth_prefix_out_of_order_replies_stay_with_source_channel(self):
+        first_started, release_first = asyncio.Event(), asyncio.Event()
+        body_a, body_b = b'samekey!' + b'a' * 32, b'samekey!' + b'b' * 32
+        async def handler(request):
+            if request.content == body_a:
+                first_started.set()
+                await release_first.wait()
+            return httpx.Response(200, content=request.content, extensions={
+                'http_version': b'HTTP/2', 'stream_id': 3 if request.content == body_a else 5})
+        self._transport(handler)
+        first, second = self._channel(), self._channel()
+        await first.send(body_a, False)
+        await first_started.wait()
+        await second.send(body_b, False)
+        self.assertEqual(await asyncio.wait_for(second.receive(), 1), body_b)
+        release_first.set()
+        self.assertEqual(await asyncio.wait_for(first.receive(), 1), body_a)
+        self.assertEqual(self.lane.max_inflight, 2)
+
+    async def test_close_one_channel_does_not_cancel_another(self):
+        async def handler(request):
+            if request.content[:1] == b'a':
+                await asyncio.Event().wait()
+            return httpx.Response(200, content=request.content, extensions={'http_version': b'HTTP/2'})
+        self._transport(handler)
+        first, second = self._channel(), self._channel()
+        await first.send(b'a' * 40, False)
+        await first.close()
+        await second.send(b'b' * 40, False)
+        self.assertEqual(await asyncio.wait_for(second.receive(), 1), b'b' * 40)
+        await second.close()
+        self.assertEqual(self.lane.inflight, 0)
+        self.assertEqual(self.lane.queued_bytes, 0)
+        self.assertFalse(self.lane.closed)
+
+    async def test_backpressure_waits_instead_of_dropping_requests(self):
+        release = asyncio.Event()
+        async def handler(request):
+            await release.wait()
+            return httpx.Response(200, content=request.content, extensions={'http_version': b'HTTP/2'})
+        self._transport(handler)
+        channel = self._channel()
+        for _ in range(MAX_CHANNEL_REQUESTS):
+            await channel.send(b'x' * 40, False)
+        extra = asyncio.create_task(channel.send(b'y' * 40, False))
+        await asyncio.sleep(0.01)
+        self.assertFalse(extra.done())
+        self.assertEqual(len(channel.pending), MAX_CHANNEL_REQUESTS)
+        release.set()
+        await asyncio.wait_for(extra, 1)
+        received = [await asyncio.wait_for(channel.receive(), 1) for _ in range(MAX_CHANNEL_REQUESTS + 1)]
+        self.assertEqual(received.count(b'y' * 40), 1)
+
+    async def test_http_error_closes_only_its_channel_without_replaying_packet(self):
+        calls = []
+        async def handler(request):
+            calls.append(request.content)
+            return httpx.Response(403, extensions={'http_version': b'HTTP/2'})
+        self._transport(handler)
+        channel = self._channel()
+        await channel.send(b'x' * 40, False)
+        with self.assertRaisesRegex(ConnectionError, '403'):
+            await asyncio.wait_for(channel.receive(), 1)
+        self.assertEqual(calls, [b'x' * 40])
+        self.assertTrue(channel.closed)
+        self.assertFalse(self.lane.closed)
+
+    async def test_long_polls_cannot_block_fresh_request_and_retired_packet_recovers(self):
+        bodies = [b'samekey!' + bytes([index]) * 32 for index in range(MAX_CHANNEL_REQUESTS)]
+        fresh = b'samekey!' + b'n' * 32
+        calls = []
+        async def handler(request):
+            calls.append(request.content)
+            await request.extensions['trace']('http2.send_request_body.complete', {'stream_id': len(calls) * 2 + 1})
+            if request.content == fresh:
+                return httpx.Response(200, content=b'n' * 40, extensions={'http_version': b'HTTP/2'})
+            if request.content == bodies[0] and calls.count(bodies[0]) == 2:
+                return httpx.Response(200, content=b'r' * 40, extensions={'http_version': b'HTTP/2'})
+            await asyncio.Event().wait()
+        self._transport(handler)
+        channel = self._channel()
+        for body in bodies:
+            await channel.send(body, False)
+        await asyncio.sleep(0)
+        original = next(iter(channel.pending))
+        with patch('proxy.cf_h2.REPLAY_REQUEST_SECONDS', .02), \
+                patch('proxy.cf_h2.REPLAY_IDLE_SECONDS', .01), \
+                patch('proxy.cf_h2.REPLAY_CHECK_SECONDS', .005):
+            recovery = asyncio.create_task(channel._recover())
+            try:
+                await asyncio.wait_for(channel.send(fresh, False), 1)
+                received = [await asyncio.wait_for(channel.receive(), 1) for _ in range(2)]
+                self.assertCountEqual(received, [b'n' * 40, b'r' * 40])
+                self.assertTrue(original.cancelled())
+                self.assertTrue(channel.replay_history[0].retired)
+                self.assertEqual(calls.count(bodies[0]), 2)
+                self.assertEqual(channel.capacity_rotations, 1)
+                self.assertEqual(self.lane.max_inflight, MAX_CHANNEL_REQUESTS)
+                self.assertEqual(self.lane.errors, 0)
+                self.assertFalse(channel.closed)
+            finally:
+                recovery.cancel()
+                await asyncio.gather(recovery, return_exceptions=True)
+
+    async def test_shared_lane_long_polls_cannot_block_new_channel(self):
+        fresh = b'samekey!' + b'n' * 32
+        async def handler(request):
+            await request.extensions['trace']('http2.send_request_body.complete', {'stream_id': 3})
+            if request.content == fresh:
+                return httpx.Response(200, content=b'n' * 40, extensions={'http_version': b'HTTP/2'})
+            await asyncio.Event().wait()
+        self._transport(handler)
+        first, second, waiting = self._channel(), self._channel(), self._channel()
+        with patch('proxy.cf_h2.MAX_LANE_REQUESTS', 2), \
+                patch('proxy.cf_h2.REPLAY_REQUEST_SECONDS', .02), \
+                patch('proxy.cf_h2.REPLAY_CHECK_SECONDS', .005):
+            await first.send(b'samekey!' + b'a' * 32, False)
+            await second.send(b'samekey!' + b'b' * 32, False)
+            await asyncio.sleep(0)
+            originals = [next(iter(channel.pending)) for channel in (first, second)]
+            recovery = asyncio.create_task(waiting._recover())
+            try:
+                await asyncio.wait_for(waiting.send(fresh, False), 1)
+                self.assertEqual(await asyncio.wait_for(waiting.receive(), 1), b'n' * 40)
+                self.assertEqual(sum(task.cancelled() for task in originals), 1)
+                self.assertEqual(sum(channel.capacity_rotations for channel in (first, second)), 1)
+                self.assertEqual(self.lane.capacity_rotations, 1)
+                self.assertEqual(self.lane.max_inflight, 2)
+            finally:
+                recovery.cancel()
+                await asyncio.gather(recovery, return_exceptions=True)
+
+    async def test_upload_room_accounts_for_waiting_packet_size(self):
+        for limit in ('MAX_CHANNEL_BYTES', 'MAX_LANE_BYTES'):
+            with self.subTest(limit=limit):
+                started = asyncio.Event()
+                async def handler(request):
+                    await request.extensions['trace']('http2.send_request_body.complete', {'stream_id': 3})
+                    if request.content == b'k' * 40:
+                        started.set()
+                        await asyncio.Event().wait()
+                    return httpx.Response(200, content=b'n' * 40, extensions={'http_version': b'HTTP/2'})
+                self._transport(handler)
+                channel = self._channel()
+                with patch('proxy.cf_h2.' + limit, 60), \
+                        patch('proxy.cf_h2.REPLAY_REQUEST_SECONDS', .02), \
+                        patch('proxy.cf_h2.REPLAY_CHECK_SECONDS', .005):
+                    await channel.send(b'k' * 40, False)
+                    await started.wait()
+                    recovery = asyncio.create_task(channel._recover())
+                    try:
+                        # 20 bytes are free, but the pending upload needs 40.
+                        await asyncio.wait_for(channel.send(b'n' * 40, False), .5)
+                        self.assertEqual(await asyncio.wait_for(channel.receive(), .5), b'n' * 40)
+                        self.assertEqual(channel.capacity_rotations, 1)
+                    finally:
+                        recovery.cancel()
+                        await asyncio.gather(recovery, return_exceptions=True)
+                        await channel.close()
+
+    async def test_close_releases_request_cancelled_before_coroutine_start_once(self):
+        self._transport(lambda request: self.fail('cancelled request reached HTTP'))
+        channel = self._channel()
+        create_task = asyncio.create_task
+        def cancelled_task(coroutine):
+            task = create_task(coroutine)
+            task.cancel()
+            return task
+        with patch('proxy.cf_h2.asyncio.create_task', cancelled_task):
+            await channel.send(b'k' * 40, False)
+        packet = channel.replay_history[0]
+        await channel.close()
+        await channel.close()
+        self.assertEqual((self.lane.inflight, self.lane.queued_bytes, channel.pending_bytes), (0, 0, 0))
+        self.assertFalse(channel.pending)
+        self.assertFalse(channel.pending_since)
+        self.assertFalse(channel.sent_since)
+        self.assertFalse(packet.originals)
+
+    async def test_capacity_rotation_keeps_upload_and_started_response_intact(self):
+        uploading, receiving, waiting = [b'samekey!' + bytes([index]) * 32 for index in range(3)]
+        fresh = b'samekey!' + b'n' * 32
+        body_started = asyncio.Event()
+        class SlowBody(httpx.AsyncByteStream):
+            async def __aiter__(self):
+                yield b'r' * 20
+                body_started.set()
+                await asyncio.Event().wait()
+        async def handler(request):
+            if request.content == uploading:
+                await asyncio.Event().wait()  # No send_request_body.complete yet.
+            await request.extensions['trace']('http2.send_request_body.complete', {'stream_id': 3})
+            if request.content == receiving:
+                return httpx.Response(200, stream=SlowBody(), extensions={'http_version': b'HTTP/2'})
+            if request.content == waiting:
+                await asyncio.Event().wait()
+            return httpx.Response(200, content=b'n' * 40, extensions={'http_version': b'HTTP/2'})
+        self._transport(handler, trace_send=False)
+        channel = self._channel()
+        with patch('proxy.cf_h2.MAX_CHANNEL_REQUESTS', 3), \
+                patch('proxy.cf_h2.REPLAY_REQUEST_SECONDS', .02), \
+                patch('proxy.cf_h2.REPLAY_CHECK_SECONDS', .005):
+            for body in (uploading, receiving, waiting):
+                await channel.send(body, False)
+            await asyncio.wait_for(body_started.wait(), 1)
+            tasks = list(channel.pending)
+            recovery = asyncio.create_task(channel._recover())
+            try:
+                await asyncio.wait_for(channel.send(fresh, False), 1)
+                self.assertEqual(await asyncio.wait_for(channel.receive(), 1), b'n' * 40)
+                self.assertFalse(tasks[0].done())
+                self.assertFalse(tasks[1].done())
+                self.assertTrue(tasks[2].cancelled())
+            finally:
+                recovery.cancel()
+                await asyncio.gather(recovery, return_exceptions=True)
+
+    async def test_capacity_rotation_never_retires_unretained_or_plaintext_request(self):
+        release = asyncio.Event()
+        async def handler(request):
+            await request.extensions['trace']('http2.send_request_body.complete', {'stream_id': 3})
+            await release.wait()
+            return httpx.Response(200, content=b'r' * 40, extensions={'http_version': b'HTTP/2'})
+        self._transport(handler)
+        channel = self._channel()
+        with patch('proxy.cf_h2.MAX_CHANNEL_REQUESTS', 2), \
+                patch('proxy.cf_h2.REPLAY_REQUEST_SECONDS', .005), \
+                patch('proxy.cf_h2.REPLAY_CHECK_SECONDS', .002):
+            await channel.send(b'\0' * 8 + b'p' * 32, False)
+            await channel.send(b'samekey!' + b'u' * (64 * 1024), False)
+            extra = asyncio.create_task(channel.send(b'samekey!' + b'n' * 32, False))
+            recovery = asyncio.create_task(channel._recover())
+            try:
+                await asyncio.sleep(.04)
+                self.assertFalse(extra.done())
+                self.assertEqual(channel.capacity_rotations, 0)
+                self.assertTrue(all(not task.done() for task in channel.pending))
+                release.set()
+                await asyncio.wait_for(extra, 1)
+            finally:
+                release.set()
+                extra.cancel()
+                recovery.cancel()
+                await asyncio.gather(extra, recovery, return_exceptions=True)
+
+    async def test_http1_rejected_by_preflight(self):
+        self._transport(lambda request: httpx.Response(200, extensions={'http_version': b'HTTP/1.1'}))
+        with self.assertRaisesRegex(ValueError, 'HTTP/1.1'):
+            await self.lane._preflight()
+
+    async def test_telegram_head_501_still_allows_http2_preflight(self):
+        self._transport(lambda request: httpx.Response(501, extensions={'http_version': b'HTTP/2'}))
+        await self.lane._preflight()
+
+    async def _packet_bridge(self, *, content_length=True, truncate=False, concurrent_small=False):
+        first_chunk_read, release_tail = asyncio.Event(), asyncio.Event()
+        body = b'x' * (64 * 1024 + 40)
+        first_chunk = body[:16384]
+        class ResponseStream(httpx.AsyncByteStream):
+            async def __aiter__(self):
+                yield first_chunk
+                first_chunk_read.set()
+                await release_tail.wait()
+                if not truncate:
+                    yield body[len(first_chunk):]
+        async def handler(request):
+            if request.content == b's' * 40:
+                return httpx.Response(200, content=b'z' * 40,
+                                      extensions={'http_version': b'HTTP/2'})
+            headers = {'content-length': str(len(body))} if content_length else {}
+            return httpx.Response(200, headers=headers, stream=ResponseStream(),
+                                  extensions={'http_version': b'HTTP/2'})
+        self._transport(handler)
+        class Writer:
+            def __init__(self):
+                self.body = bytearray()
+                self.changed = asyncio.Event()
+            def write(self, data):
+                self.body.extend(data)
+                self.changed.set()
+            async def drain(self):
+                await asyncio.sleep(0)
+        reader, writer = asyncio.StreamReader(), Writer()
+        channel = self._channel()
+        ctx = SimpleNamespace(clt_enc=_IdentityCipher(), clt_dec=_IdentityCipher())
+        task = asyncio.create_task(bridge_h2(reader, writer, channel, ctx, PROTO_TAG_INTERMEDIATE))
+        reader.feed_data(struct.pack('<I', 40) + b'l' * 40)
+        try:
+            await asyncio.wait_for(first_chunk_read.wait(), 1)
+            self.assertEqual(writer.body, b'')
+            if concurrent_small:
+                reader.feed_data(struct.pack('<I', 40) + b's' * 40)
+                await asyncio.wait_for(writer.changed.wait(), 1)
+                self.assertEqual(writer.body, struct.pack('<I', 40) + b'z' * 40)
+            release_tail.set()
+            if truncate:
+                await asyncio.wait_for(task, 1)
+                self.assertTrue(channel.closed)
+                self.assertEqual(writer.body, b'')
+            else:
+                expected = struct.pack('<I', len(body)) + body
+                if concurrent_small:
+                    expected = struct.pack('<I', 40) + b'z' * 40 + expected
+                for _ in range(100):
+                    if len(writer.body) >= len(expected):
+                        break
+                    await asyncio.sleep(0)
+                self.assertEqual(writer.body, expected)
+        finally:
+            release_tail.set()
+            reader.feed_eof()
+            if not task.done():
+                task.cancel()
+            await asyncio.gather(task, return_exceptions=True)
+
+    async def test_native_frame_starts_only_after_entire_body_is_available(self):
+        await self._packet_bridge()
+
+    async def test_ready_reply_bypasses_incomplete_large_response_on_same_channel(self):
+        await self._packet_bridge(concurrent_small=True)
+
+    async def test_response_without_length_uses_buffered_framing(self):
+        await self._packet_bridge(content_length=False)
+
+    async def test_truncated_stream_closes_native_channel(self):
+        await self._packet_bridge(truncate=True)
+
+    async def test_recovery_pauses_during_body_receive_and_resumes_afterwards(self):
+        started, release = asyncio.Event(), asyncio.Event()
+        calls = []
+
+        class SlowBody(httpx.AsyncByteStream):
+            async def __aiter__(self):
+                yield b'r' * 20
+                started.set()
+                await release.wait()
+                yield b'r' * 20
+
+        def handler(request):
+            calls.append(request.content)
+            if len(calls) == 1:
+                return httpx.Response(200, stream=SlowBody(), extensions={'http_version': b'HTTP/2'})
+            return httpx.Response(200, content=b's' * 40, extensions={'http_version': b'HTTP/2'})
+
+        self._transport(handler)
+        channel = self._channel()
+        await channel.send(b'x' * 40, False)
+        await asyncio.wait_for(started.wait(), 1)
+        with patch('proxy.cf_h2.REPLAY_IDLE_SECONDS', .01), patch('proxy.cf_h2.REPLAY_CHECK_SECONDS', .005):
+            recovery = asyncio.create_task(channel._recover())
+            try:
+                await asyncio.sleep(.06)
+                self.assertEqual(len(calls), 1)
+                self.assertEqual(channel.receiving_http, 1)
+                release.set()
+                self.assertEqual(await asyncio.wait_for(channel.receive(), 1), b'r' * 40)
+                self.assertEqual(await asyncio.wait_for(channel.receive(), 1), b's' * 40)
+                self.assertEqual(calls, [b'x' * 40, b'x' * 40])
+                self.assertEqual(channel.receiving_http, 0)
+                self.assertEqual(self.lane.reply_buffer_bytes, 0)
+            finally:
+                recovery.cancel()
+                await asyncio.gather(recovery, return_exceptions=True)
+
+    async def test_lane_reply_memory_bound_and_cancellation_cleanup(self):
+        started = asyncio.Event()
+
+        class SlowBody(httpx.AsyncByteStream):
+            async def __aiter__(self):
+                yield b'r' * 32
+                started.set()
+                await asyncio.Event().wait()
+
+        def handler(request):
+            if request.content == b'a' * 40:
+                return httpx.Response(200, stream=SlowBody(), extensions={'http_version': b'HTTP/2'})
+            return httpx.Response(200, content=b'r' * 16, extensions={'http_version': b'HTTP/2'})
+
+        self._transport(handler)
+        slow = self._channel()
+        with patch('proxy.cf_h2.MAX_LANE_REPLY_BYTES', 40):
+            await slow.send(b'a' * 40, False)
+            await asyncio.wait_for(started.wait(), 1)
+            self.assertEqual(self.lane.reply_buffer_bytes, 32)
+            with self.assertRaisesRegex(BufferError, 'response buffers'):
+                await self.lane._post(b'b' * 40, 2)
+            self.assertEqual(self.lane.reply_buffer_bytes, 32)
+            self.assertEqual(self.lane.failed_until, 0)
+            await slow.close()
+            self.assertEqual(slow.receiving_http, 0)
+            self.assertEqual(self.lane.reply_buffer_bytes, 0)
+            self.assertEqual(await self.lane._post(b'b' * 40, 2), b'r' * 16)
+
+    async def test_native_disconnect_does_not_cool_down_healthy_lane(self):
+        self._transport(lambda request: httpx.Response(
+            200, content=b'r' * 65536, extensions={'http_version': b'HTTP/2'}))
+
+        class DisconnectedWriter:
+            def write(self, data):
+                pass
+
+            async def drain(self):
+                raise ConnectionResetError('Connection lost')
+
+        channel = self._channel()
+        reader = asyncio.StreamReader()
+        reader.feed_data(struct.pack('<I', 40) + b'x' * 40)
+        ctx = SimpleNamespace(clt_enc=_IdentityCipher(), clt_dec=_IdentityCipher())
+        await asyncio.wait_for(bridge_h2(reader, DisconnectedWriter(), channel, ctx,
+                                         PROTO_TAG_INTERMEDIATE), 1)
+        self.assertTrue(channel.closed)
+        self.assertEqual(self.lane.failed_until, 0)
+        self.assertEqual(self.lane.errors, 0)
+        healthy = self._channel()
+        await healthy.send(b'h' * 40, False)
+        self.assertEqual(await asyncio.wait_for(healthy.receive(), 1), b'r' * 65536)
+
+    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'})
+        self._transport(handler)
+        slow, healthy = self._channel(), self._channel()
+        with patch('proxy.cf_h2.MAX_CHANNEL_BYTES', 80):
+            for _ in range(3):
+                await slow.send(b's' * 40, False)
+                await asyncio.gather(*list(slow.pending))
+            with self.assertRaises(BufferError):
+                await asyncio.wait_for(slow.receive(), 1)
+            await healthy.send(b'h' * 40, False)
+            self.assertEqual(await asyncio.wait_for(healthy.receive(), 1), b'h' * 40)
+        self.assertFalse(self.lane.closed)
+
+    async def test_simultaneous_pool_opens_preflight_once_and_share_client(self):
+        pool = CfH2Pool()
+        balancer = _Balancer()
+        balancer.update_domains_list(['example.org'])
+        started = []
+        async def preflight(lane):
+            started.append(lane)
+            await asyncio.sleep(0)
+        try:
+            with patch('proxy.cf_h2.balancer', balancer), patch.object(_HttpLane, '_preflight', preflight):
+                channels = await asyncio.gather(*(pool.open(4, 'owned-test') for _ in range(6)))
+            self.assertEqual(len(started), 1)
+            self.assertEqual(len({id(channel.lane.client) for channel in channels}), 1)
+            self.assertEqual(len({channel.channel_id for channel in channels}), 6)
+        finally:
+            await pool.close()
+
+    async def test_recovery_reuses_exact_ciphertext_when_original_http_is_still_pending(self):
+        calls = []
+        original_waiting = asyncio.Event()
+        body = b'opaque-key-and-message-id' + b'x' * 16
+        async def handler(request):
+            calls.append(request.content)
+            if len(calls) == 1:
+                original_waiting.set()
+                await asyncio.Event().wait()
+            return httpx.Response(200, content=b'r' * 40, extensions={'http_version': b'HTTP/2'})
+        self._transport(handler)
+        channel = self._channel()
+        await channel.send(body, False)
+        await original_waiting.wait()
+        with patch('proxy.cf_h2.REPLAY_REQUEST_SECONDS', .01), patch('proxy.cf_h2.REPLAY_CHECK_SECONDS', .005):
+            recovery = asyncio.create_task(channel._recover())
+            try:
+                self.assertEqual(await asyncio.wait_for(channel.receive(), 1), b'r' * 40)
+                self.assertEqual(calls, [body, body])
+                self.assertEqual(channel.replay_requests, 1)
+                self.assertEqual(self.lane.max_inflight, 2)
+            finally:
+                recovery.cancel()
+                await asyncio.gather(recovery, return_exceptions=True)
+                await channel.close()
+        self.assertFalse(channel.pending)
+        self.assertEqual(self.lane.inflight, 0)
+
+    async def test_recovery_after_http_finished_without_native_reply(self):
+        calls = []
+        async def handler(request):
+            calls.append(request.content)
+            return httpx.Response(200, content=b'' if len(calls) == 1 else b'r' * 40,
+                                  extensions={'http_version': b'HTTP/2'})
+        self._transport(handler)
+        channel = self._channel()
+        body = b'k' * 40
+        await channel.send(body, False)
+        await asyncio.gather(*list(channel.pending))
+        self.assertTrue(channel.queue.empty())
+        with patch('proxy.cf_h2.REPLAY_IDLE_SECONDS', .01), patch('proxy.cf_h2.REPLAY_CHECK_SECONDS', .005):
+            recovery = asyncio.create_task(channel._recover())
+            try:
+                self.assertEqual(await asyncio.wait_for(channel.receive(), 1), b'r' * 40)
+                await asyncio.sleep(.03)
+                self.assertEqual(calls, [body, body])
+            finally:
+                recovery.cancel()
+                await asyncio.gather(recovery, return_exceptions=True)
+
+    async def test_idle_receiver_preserves_capacity_for_new_native_request(self):
+        calls = []
+        probes_waiting = asyncio.Event()
+        async def handler(request):
+            calls.append(request.content)
+            if calls.count(request.content) > 1:
+                if len(calls) == 4:
+                    probes_waiting.set()
+                await asyncio.Event().wait()
+            return httpx.Response(200, content=b'r' * 40, extensions={'http_version': b'HTTP/2'})
+        self._transport(handler)
+        channel = self._channel()
+        bodies = [b'samekey!' + bytes([index]) * 32 for index in range(4)]
+        for body in bodies[:3]:
+            await channel.send(body, False)
+            await asyncio.gather(*list(channel.pending))
+            await channel.receive()
+        with patch('proxy.cf_h2.REPLAY_IDLE_SECONDS', .01), patch('proxy.cf_h2.REPLAY_CHECK_SECONDS', .005):
+            recovery = asyncio.create_task(channel._recover())
+            try:
+                await asyncio.wait_for(probes_waiting.wait(), 1)
+                await asyncio.sleep(.03)
+                self.assertEqual(channel.replay_requests, 1)
+                self.assertEqual(len(channel.replay_pending), 1)
+                await asyncio.wait_for(channel.send(bodies[3], False), 1)
+                self.assertEqual(await asyncio.wait_for(channel.receive(), 1), b'r' * 40)
+                self.assertEqual(calls[3:4], [bodies[2]])
+            finally:
+                recovery.cancel()
+                await asyncio.gather(recovery, return_exceptions=True)
+                await asyncio.gather(channel.close(), channel.close())
+        self.assertEqual(self.lane.inflight, 0)
+        self.assertEqual(channel.pending_bytes, 0)
+        self.assertFalse(channel.replay_pending)
+
+    async def test_handshake_and_expired_ciphertext_are_never_replayed(self):
+        calls = []
+        async def handler(request):
+            calls.append(request.content)
+            return httpx.Response(200, content=b'r' * 40, extensions={'http_version': b'HTTP/2'})
+        self._transport(handler)
+        channel = self._channel()
+        for body in (b'\x00' * 8 + b'h' * 32, b'k' * 40):
+            await channel.send(body, False)
+            await asyncio.gather(*list(channel.pending))
+            await channel.receive()
+        # Windows may give consecutive monotonic() calls the same timestamp;
+        # an event wakeup need not advance the clock like the old sleep did.
+        channel.replay_history[-1].sent_at -= 1
+        with patch('proxy.cf_h2.REPLAY_IDLE_SECONDS', .01), patch('proxy.cf_h2.REPLAY_CHECK_SECONDS', .005), \
+                patch('proxy.cf_h2.REPLAY_MAX_AGE', 0):
+            recovery = asyncio.create_task(channel._recover())
+            try:
+                await asyncio.sleep(.04)
+                self.assertEqual(len(calls), 2)
+            finally:
+                recovery.cancel()
+                await asyncio.gather(recovery, return_exceptions=True)
+
+    async def test_unhelpful_replies_do_not_cause_unbounded_recovery(self):
+        calls = []
+        retries_finished = asyncio.Event()
+        async def handler(request):
+            calls.append((request.content, time.monotonic()))
+            if len(calls) == 5:
+                retries_finished.set()
+            return httpx.Response(200, content=b'', extensions={'http_version': b'HTTP/2'})
+        self._transport(handler)
+        channel = self._channel()
+        bodies = [b'samekey!' + bytes([index]) * 32 for index in range(2)]
+        for body in bodies:
+            await channel.send(body, False)
+        await asyncio.gather(*list(channel.pending))
+        with patch('proxy.cf_h2.REPLAY_IDLE_SECONDS', .005), \
+                patch('proxy.cf_h2.REPLAY_CHECK_SECONDS', .002), \
+                patch('proxy.cf_h2.REPLAY_RETRY_SECONDS', .02):
+            recovery = asyncio.create_task(channel._recover())
+            try:
+                await asyncio.wait_for(retries_finished.wait(), 1)
+                await asyncio.sleep(.06)
+                older = [when for body, when in calls if body == bodies[0]]
+                latest = [when for body, when in calls if body == bodies[1]]
+                # Completed history is not scanned; only the latest is polled.
+                self.assertEqual(len(older), 1)
+                self.assertEqual(len(latest), 4)
+                for first, second in zip(latest[1:], latest[2:]):
+                    self.assertGreaterEqual(second - first, .015)
+                self.assertFalse(channel.pending)
+                self.assertFalse(channel.replay_pending)
+                self.assertEqual(self.lane.inflight, 0)
+            finally:
+                recovery.cancel()
+                await asyncio.gather(recovery, return_exceptions=True)
+
+    async def test_old_pending_request_has_bounded_retries_after_newer_requests(self):
+        body = b'samekey!' + b't' * 32
+        original_started = asyncio.Event()
+        retries_finished = asyncio.Event()
+        attempts = 0
+        async def handler(request):
+            nonlocal attempts
+            if request.content == body:
+                attempts += 1
+                if attempts == 1:
+                    original_started.set()
+                    await asyncio.Event().wait()
+                if attempts == 4:
+                    retries_finished.set()
+            return httpx.Response(200, content=b'', extensions={'http_version': b'HTTP/2'})
+        self._transport(handler)
+        channel = self._channel()
+        await channel.send(body, False)
+        await original_started.wait()
+        original = next(iter(channel.pending))
+        await channel.send(b'samekey!' + b'n' * 32, False)
+        with patch('proxy.cf_h2.REPLAY_IDLE_SECONDS', .01), \
+                patch('proxy.cf_h2.REPLAY_REQUEST_SECONDS', .02), \
+                patch('proxy.cf_h2.REPLAY_CHECK_SECONDS', .005), \
+                patch('proxy.cf_h2.REPLAY_RETRY_SECONDS', .02):
+            recovery = asyncio.create_task(channel._recover())
+            try:
+                await asyncio.wait_for(retries_finished.wait(), 1)
+                await asyncio.sleep(.06)
+                self.assertEqual(attempts, 4)
+                self.assertFalse(original.done())
+                self.assertFalse(channel.replay_pending)
+            finally:
+                recovery.cancel()
+                await asyncio.gather(recovery, return_exceptions=True)
+
+    async def test_closed_pool_cannot_open_new_channels(self):
+        pool = CfH2Pool()
+        await pool.close()
+        with self.assertRaisesRegex(ConnectionError, 'pool closed'):
+            await pool.open(4, 'owned-test')
+
+    async def test_concurrent_upload_and_recovery_cannot_exceed_channel_request_bound(self):
+        release = asyncio.Event()
+        async def handler(request):
+            await release.wait()
+            return httpx.Response(200, content=request.content, extensions={'http_version': b'HTTP/2'})
+        self._transport(handler)
+        channel = self._channel()
+        for _ in range(MAX_CHANNEL_REQUESTS - 1):
+            await channel.send(b'k' * 40, False)
+        first = asyncio.create_task(channel.send(b'a' * 40, False))
+        second = asyncio.create_task(channel.send(b'b' * 40, False))
+        try:
+            await asyncio.wait_for(first, 1)
+            await asyncio.sleep(.02)
+            self.assertEqual(len(channel.pending), MAX_CHANNEL_REQUESTS)
+            self.assertFalse(second.done())
+            release.set()
+            await asyncio.wait_for(second, 1)
+        finally:
+            release.set()
+            first.cancel()
+            second.cancel()
+            await asyncio.gather(first, second, return_exceptions=True)
+
+    async def test_failing_preflight_returns_to_ws_without_consuming_native_data(self):
+        pool = CfH2Pool()
+        balancer = _Balancer()
+        balancer.update_domains_list(['example.org'])
+        async def preflight(lane):
+            raise OSError('owned connection failure')
+        try:
+            with patch('proxy.cf_h2.balancer', balancer), patch.object(_HttpLane, '_preflight', preflight):
+                self.assertIsNone(await pool.open(4, 'owned-test'))
+            self.assertFalse(pool.lanes)
+        finally:
+            await pool.close()
+
+    async def test_hanging_setup_returns_to_ws_and_cleans_up_lane(self):
+        pool = CfH2Pool()
+        balancer = _Balancer()
+        balancer.update_domains_list(['example.org'])
+        lanes = []
+
+        async def preflight(lane):
+            lanes.append(lane)
+            await asyncio.Event().wait()
+
+        try:
+            with patch('proxy.cf_h2.balancer', balancer), \
+                    patch('proxy.cf_h2.SETUP_TIMEOUT', .02), \
+                    patch.object(_HttpLane, '_preflight', preflight):
+                self.assertIsNone(await pool.open(4, 'owned-test'))
+                self.assertIsNone(await pool.open(4, 'owned-test'))
+            self.assertEqual(len(lanes), 1)
+            self.assertTrue(lanes[0].closed)
+            self.assertTrue(lanes[0].client.is_closed)
+            self.assertFalse(pool.lanes)
+        finally:
+            await pool.close()
+
+    async def test_recovery_does_not_wait_for_channel_or_lane_capacity(self):
+        channel = self._channel()
+        body = b'k' * 40
+        channel.replay_history.append(_ReplayPacket(0, body, 0))
+        with patch('proxy.cf_h2.MAX_CHANNEL_REQUESTS', 0):
+            self.assertFalse(await asyncio.wait_for(channel.send(body, False, replay=True), .1))
+        with patch('proxy.cf_h2.MAX_LANE_REQUESTS', 0):
+            self.assertFalse(await asyncio.wait_for(channel.send(body, False, replay=True), .1))
+        self.assertFalse(channel.send_lock.locked())
+        self.assertFalse(channel.replay_history[0].replayed)
+        self.assertEqual(channel.requests, 0)
+        self.assertEqual(self.lane.inflight, 0)
+
+    async def test_recovery_does_not_queue_behind_native_upload(self):
+        channel = self._channel()
+        async with channel.send_lock:
+            self.assertFalse(await asyncio.wait_for(channel.send(b'k' * 40, False, replay=True), .1))
+        self.assertEqual(channel.requests, 0)
+
+    async def test_wait_diagnostics_combine_channels_and_are_rate_limited(self):
+        async def handler(request):
+            await asyncio.Event().wait()
+        self._transport(handler)
+        now = time.monotonic()
+        for _ in range(2):
+            channel = self._channel()
+            await channel.send(b'k' * 40, False)
+            channel.last_progress = now - 4
+            channel.pending_since = dict.fromkeys(channel.pending, now - 4)
+        with self.assertLogs('tg-mtproto-proxy', level='DEBUG') as captured:
+            for _ in range(20):
+                self.lane.log_wait(now)
+        self.assertEqual(len(captured.records), 1)
+        self.assertIn('H2 WAIT', captured.output[0])
+        self.assertIn('affected=2', captured.output[0])
+        self.assertIn('channel=1', captured.output[0])
+        self.assertIn('channel=2', captured.output[0])
+        self.assertIn('recovery_left=', captured.output[0])
+
+    async def test_new_request_after_idle_does_not_log_a_stall(self):
+        async def handler(request):
+            await asyncio.Event().wait()
+        self._transport(handler)
+        channel = self._channel()
+        await channel.send(b'k' * 40, False)
+        now = time.monotonic()
+        channel.last_progress = now - 20
+        with patch('proxy.cf_h2.log.isEnabledFor', return_value=True), \
+                patch('proxy.cf_h2.log.debug') as debug:
+            self.lane.log_wait(now)
+        debug.assert_not_called()
+
+    async def test_wait_diagnostics_include_old_request_while_other_replies_progress(self):
+        async def handler(request):
+            await asyncio.Event().wait()
+        self._transport(handler)
+        channel = self._channel()
+        await channel.send(b'k' * 40, False)
+        now = time.monotonic()
+        channel.pending_since = dict.fromkeys(channel.pending, now - 4)
+        channel.last_progress = now
+        with self.assertLogs('tg-mtproto-proxy', level='DEBUG') as captured:
+            self.lane.log_wait(now)
+        self.assertEqual(len(captured.records), 1)
+        self.assertIn('waiting_headers=1', captured.output[0])
+        self.assertIn('idle_ms=0', captured.output[0])
+
+    async def test_slow_response_diagnostics_are_rate_limited(self):
+        self._transport(lambda request: httpx.Response(
+            200, content=b'r' * 40, extensions={'http_version': b'HTTP/2'}))
+        with patch('proxy.cf_h2.SLOW_RESPONSE_SECONDS', 0), \
+                self.assertLogs('tg-mtproto-proxy', level='DEBUG') as captured:
+            for _ in range(10):
+                await self.lane._post(b'k' * 40, 1)
+        self.assertEqual(len(captured.records), 1)
+        self.assertIn('H2 SLOW', captured.output[0])
+        self.assertIn('headers_ms=', captured.output[0])
+        self.assertIn('body_ms=', captured.output[0])
+
+    async def test_pending_details_link_replay_to_original_and_bound_history(self):
+        calls = 0
+        original_waiting = asyncio.Event()
+        async def handler(request):
+            nonlocal calls
+            calls += 1
+            if calls == 1:
+                original_waiting.set()
+                await asyncio.Event().wait()
+            return httpx.Response(200, content=b'', extensions={'http_version': b'HTTP/2'})
+        self._transport(handler)
+        channel = self._channel()
+        with self.assertLogs('tg-mtproto-proxy', level='DEBUG') as captured:
+            await channel.send(b'samekey!' + b'p' * 32, False)
+            await original_waiting.wait()
+            original = next(iter(channel.pending))
+            with patch('proxy.cf_h2.REPLAY_REQUEST_SECONDS', 0):
+                await channel.send(b'samekey!' + b'p' * 32, False, replay=True)
+            await asyncio.gather(*(task for task in channel.pending if task is not original))
+            trace = channel.request_traces[original]
+            replay = channel.recent_traces[-1]
+            self.assertEqual(replay.packet, trace.packet)
+            self.assertNotEqual(replay.request, trace.request)
+            self.assertTrue(replay.replay)
+            now = time.monotonic()
+            trace.started = now - 4
+            channel.pending_since[original] = now - 4
+            self.lane.log_wait(now)
+            for index in range(24):
+                await channel.send(b'samekey!' + bytes([index]) * 32, False)
+                await asyncio.gather(*(task for task in channel.pending if task is not original))
+            self.assertEqual(len(channel.recent_traces), 16)
+            self.assertEqual(len(channel.request_traces), 1)
+        output = '\n'.join(captured.output)
+        self.assertIn('H2 DETAIL', output)
+        self.assertIn('pkt=1', output)
+        self.assertIn('replay=1 phase=complete', output)
+        self.assertNotIn('samekey!', output)

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)

tests/test_h2_settings.py

@@ -0,0 +1,127 @@
+"""Default/CLI/GUI configuration must select the same effective H2 route."""
+import asyncio
+import contextlib
+import io
+import json
+import logging
+import tempfile
+import unittest
+from pathlib import Path
+from unittest.mock import AsyncMock, patch
+
+from proxy.config import ProxyConfig
+from proxy import tg_ws_proxy
+from utils import tray_common
+from utils.default_config import default_tray_config
+from ui.settings import prepare_settings
+
+
+class H2SettingsTest(unittest.TestCase):
+    def test_cli_default_opt_out_and_existing_launch_script(self):
+        cases = [([], True), (['--no-h2'], False),
+                 (['--cfproxy-h2-media'], True), (['--no-cfproxy'], False),
+                 (['--no-secure'], False), (['--force-test-dc'], False)]
+        for flags, enabled in cases:
+            with self.subTest(flags=flags):
+                config = ProxyConfig()
+                observed = []
+
+                async def run():
+                    observed.append(config.h2_enabled)
+
+                root = logging.getLogger()
+                previous_handlers, previous_level = root.handlers[:], root.level
+                try:
+                    with patch.object(tg_ws_proxy, 'proxy_config', config), \
+                            patch.object(tg_ws_proxy, '_run', run), \
+                            patch('sys.argv', ['proxy', '--secret', 'ab' * 16] + flags):
+                        tg_ws_proxy.main()
+                    self.assertEqual(observed, [enabled])
+                finally:
+                    for handler in root.handlers[:]:
+                        if handler not in previous_handlers:
+                            root.removeHandler(handler)
+                            handler.close()
+                    root.setLevel(previous_level)
+
+    def test_help_advertises_opt_out_only(self):
+        output = io.StringIO()
+        with patch('sys.argv', ['proxy', '--help']), contextlib.redirect_stdout(output):
+            with self.assertRaises(SystemExit) as raised:
+                tg_ws_proxy.main()
+        self.assertEqual(raised.exception.code, 0)
+        self.assertIn('--no-h2', output.getvalue())
+        self.assertNotIn('--cfproxy-h2-media', output.getvalue())
+
+    def test_saved_setting_migration_and_runtime_restart(self):
+        with tempfile.TemporaryDirectory() as folder:
+            path = Path(folder) / 'config.json'
+            path.write_text(json.dumps({'port': 1444}), encoding='utf-8')
+            with patch.object(tray_common, 'CONFIG_FILE', path), \
+                    patch.object(tray_common, 'ensure_dirs'), \
+                    patch.object(tray_common, '_apply_ui_language'):
+                config = tray_common.load_config()
+                self.assertTrue(config['h2'])
+                change = prepare_settings(config, {'h2': False}, default_tray_config())
+                self.assertTrue(change.requires_restart)
+                tray_common.save_config(change.config)
+                saved = tray_common.load_config()
+                self.assertFalse(saved['h2'])
+                runtime = ProxyConfig()
+                with patch.object(tray_common, 'proxy_config', runtime):
+                    self.assertTrue(tray_common.apply_proxy_config(saved))
+                    self.assertFalse(runtime.h2_enabled)
+                    saved['h2'] = True
+                    self.assertTrue(tray_common.apply_proxy_config(saved))
+                    self.assertTrue(runtime.h2_enabled)
+                    for key, value in [('cfproxy', False), ('no_secure', True),
+                                       ('force_test_dc', True)]:
+                        self.assertTrue(tray_common.apply_proxy_config({**saved, key: value}))
+                        self.assertFalse(runtime.h2_enabled)
+                        self.assertTrue(runtime.cfproxy_h2_media)
+
+
+class H2ServerLifecycleTest(unittest.IsolatedAsyncioTestCase):
+    async def test_stop_and_cancellation_close_default_h2_and_listener_tasks(self):
+        for cancel in (False, True):
+            with self.subTest(cancel=cancel):
+                config = ProxyConfig(port=0, dc_redirects={}, pool_size=0)
+                stop = asyncio.Event()
+                before = asyncio.all_tasks()
+                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()):
+                    server = asyncio.create_task(tg_ws_proxy._run(stop))
+                    writer = None
+                    try:
+                        async def listening():
+                            while tg_ws_proxy._server_instance is None:
+                                await asyncio.sleep(.005)
+                        await asyncio.wait_for(listening(), 1)
+                        pool = tg_ws_proxy.cf_h2_pool
+                        self.assertIsNotNone(pool)
+                        port = tg_ws_proxy._server_instance.sockets[0].getsockname()[1]
+                        _, writer = await asyncio.open_connection('127.0.0.1', port)
+                        if cancel:
+                            server.cancel()
+                            with self.assertRaises(asyncio.CancelledError):
+                                await asyncio.wait_for(server, 2)
+                        else:
+                            stop.set()
+                            await asyncio.wait_for(server, 2)
+                        self.assertTrue(pool.closed)
+                        self.assertIsNone(tg_ws_proxy.cf_h2_pool)
+                        self.assertIsNone(tg_ws_proxy._server_instance)
+                        self.assertFalse(tg_ws_proxy._client_tasks)
+                        self.assertFalse(asyncio.all_tasks() - before)
+                    finally:
+                        server.cancel()
+                        await asyncio.gather(server, return_exceptions=True)
+                        if writer is not None:
+                            writer.close()
+                            await writer.wait_closed()
+
+
+if __name__ == '__main__':
+    unittest.main()

tests/test_h2_transport.py

@@ -0,0 +1,858 @@
+"""Wire-level regressions: real TLS sockets, ALPN and HTTP/2 flow control."""
+import asyncio
+import datetime
+import ipaddress
+import ssl
+import struct
+import tempfile
+import time
+import unittest
+from pathlib import Path
+from types import SimpleNamespace
+from unittest.mock import patch
+
+import httpx
+from h2.config import H2Configuration
+from h2.connection import H2Connection
+from h2.events import DataReceived, RequestReceived, StreamEnded, StreamReset
+from h2.settings import SettingCodes
+
+from proxy.h2_transport import H2Transport, STREAM_RECEIVE_WINDOW, _Connection
+from proxy.cf_h2 import _HttpChannel, _HttpLane, bridge_h2
+from proxy.utils import PROTO_TAG_INTERMEDIATE
+
+
+class H2WireTest(unittest.IsolatedAsyncioTestCase):
+    @classmethod
+    def setUpClass(cls):
+        from cryptography import x509
+        from cryptography.hazmat.primitives import hashes, serialization
+        from cryptography.hazmat.primitives.asymmetric import rsa
+        from cryptography.x509.oid import NameOID
+
+        cls.folder = tempfile.TemporaryDirectory(prefix='tg-h2-tests-')
+        cls.addClassCleanup(cls.folder.cleanup)
+        key = rsa.generate_private_key(public_exponent=65537, key_size=2048)
+        name = x509.Name([x509.NameAttribute(NameOID.COMMON_NAME, 'localhost')])
+        now = datetime.datetime.now(datetime.timezone.utc)
+        cert = (x509.CertificateBuilder().subject_name(name).issuer_name(name)
+                .public_key(key.public_key()).serial_number(x509.random_serial_number())
+                .not_valid_before(now - datetime.timedelta(minutes=1))
+                .not_valid_after(now + datetime.timedelta(hours=1))
+                .add_extension(x509.SubjectAlternativeName([
+                    x509.IPAddress(ipaddress.ip_address('127.0.0.1'))]), critical=False)
+                .sign(key, hashes.SHA256()))
+        cls.cert_path = Path(cls.folder.name) / 'cert.pem'
+        key_path = Path(cls.folder.name) / 'key.pem'
+        cls.cert_path.write_bytes(cert.public_bytes(serialization.Encoding.PEM))
+        key_path.write_bytes(key.private_bytes(serialization.Encoding.PEM,
+                                              serialization.PrivateFormat.PKCS8,
+                                              serialization.NoEncryption()))
+        cls.server_ssl = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER)
+        cls.server_ssl.load_cert_chain(str(cls.cert_path), str(key_path))
+        cls.server_ssl.set_alpn_protocols(['h2'])
+
+    async def asyncSetUp(self):
+        self.requests = asyncio.Queue()
+        self.resets = asyncio.Queue()
+        self.connections = []
+        self.handlers = []
+        self.tasks = []
+        self.server_max_streams = 64
+        self.server = await asyncio.start_server(self.serve, '127.0.0.1', 0, ssl=self.server_ssl)
+        port = self.server.sockets[0].getsockname()[1]
+        self.lane = _HttpLane('127.0.0.1:%d' % port, 1)
+        await self.lane.client.aclose()
+        self.transport = H2Transport(ssl.create_default_context(cafile=str(self.cert_path)))
+        self.lane.client = httpx.AsyncClient(transport=self.transport, trust_env=False,
+                                           timeout=httpx.Timeout(2, read=1))
+
+    async def asyncTearDown(self):
+        for task in self.tasks:
+            task.cancel()
+        await asyncio.gather(*self.tasks, return_exceptions=True)
+        await self.lane.close()
+        self.server.close()
+        await self.server.wait_closed()
+        for connection in self.connections:
+            connection.writer.close()
+        await asyncio.wait_for(asyncio.gather(*self.handlers), 3)
+
+    async def serve(self, reader, writer):
+        self.handlers.append(asyncio.current_task())
+        conn = SimpleNamespace(h2=H2Connection(config=H2Configuration(client_side=False)),
+                               writer=writer, bodies={}, pending={})
+        self.connections.append(conn)
+        conn.h2.initiate_connection()
+        conn.h2.update_settings({SettingCodes.MAX_CONCURRENT_STREAMS: self.server_max_streams})
+        writer.write(conn.h2.data_to_send())
+        try:
+            while True:
+                data = await reader.read(65536)
+                if not data:
+                    break
+                for event in conn.h2.receive_data(data):
+                    if isinstance(event, RequestReceived):
+                        conn.bodies[event.stream_id] = bytearray()
+                    elif isinstance(event, DataReceived):
+                        conn.bodies[event.stream_id].extend(event.data)
+                        conn.h2.acknowledge_received_data(event.flow_controlled_length, event.stream_id)
+                    elif isinstance(event, StreamEnded):
+                        await self.requests.put((conn, event.stream_id, bytes(conn.bodies[event.stream_id])))
+                    elif isinstance(event, StreamReset):
+                        conn.pending.pop(event.stream_id, None)
+                        await self.resets.put((event.stream_id, event.error_code))
+                self.flush(conn)
+                await writer.drain()
+        except (ConnectionError, ssl.SSLError):
+            pass
+        finally:
+            writer.close()
+            try:
+                await writer.wait_closed()
+            except (ConnectionError, ssl.SSLError):
+                pass
+
+    def flush(self, conn):
+        for sid, (body, end) in list(conn.pending.items()):
+            while body:
+                count = min(len(body), conn.h2.max_outbound_frame_size,
+                            conn.h2.local_flow_control_window(sid))
+                if count <= 0:
+                    break
+                chunk, body = body[:count], body[count:]
+                conn.h2.send_data(sid, chunk, end_stream=end and not body)
+            if body:
+                conn.pending[sid] = (body, end)
+            else:
+                del conn.pending[sid]
+        conn.writer.write(conn.h2.data_to_send())
+
+    def respond(self, request, body=b'r' * 40, *, end=True, length=None, status=200):
+        conn, sid, _ = request
+        conn.h2.send_headers(sid, [(b':status', str(status).encode()),
+                                   (b'content-type', b'application/octet-stream'),
+                                   (b'content-length', str(len(body) if length is None else length).encode())],
+                             end_stream=end and not body)
+        if body:
+            conn.pending[sid] = (body, end)
+        self.flush(conn)
+
+    async def post(self, body=b'x' * 40):
+        task = asyncio.create_task(self.lane._post(body, len(self.tasks) + 1))
+        self.tasks.append(task)
+        request = await asyncio.wait_for(self.requests.get(), 2)
+        self.assertEqual(request[2], body)
+        return task, request
+
+    async def test_fast_response_does_not_wait_for_unrelated_long_polls(self):
+        slow1, request1 = await self.post()
+        slow2, request2 = await self.post()
+        fast, request3 = await self.post()
+        self.respond(request3)
+        self.assertEqual(await asyncio.wait_for(fast, .5), b'r' * 40)
+        self.assertFalse(slow1.done())
+        self.assertFalse(slow2.done())
+        self.respond(request1)
+        self.respond(request2)
+        await asyncio.gather(slow1, slow2)
+        self.assertEqual(len(self.connections), 1)
+
+    async def test_debug_trace_follows_real_stream_without_logging_payload(self):
+        channel = _HttpChannel(self.lane, 1, 'trace-test')
+        body = b'private-payload-' + b'x' * 24
+        with self.assertLogs('tg-mtproto-proxy', level='DEBUG') as captured:
+            await channel.send(body, False)
+            request = await asyncio.wait_for(self.requests.get(), 1)
+            trace = next(iter(channel.request_traces.values()))
+
+            async def phase(expected):
+                while trace.phase != expected:
+                    await asyncio.sleep(.001)
+
+            await asyncio.wait_for(phase('headers'), .5)
+            self.assertEqual(trace.stream, request[1])
+            self.assertEqual(trace.packet, 1)
+            self.assertEqual(trace.request, 1)
+            self.assertIsNotNone(trace.sent_at)
+            self.assertIsNone(trace.headers_at)
+            self.respond(request, b'r' * 20, end=False, length=40)
+            await asyncio.wait_for(phase('body'), .5)
+            self.assertIsNotNone(trace.headers_at)
+            conn, sid, _ = request
+            conn.h2.send_data(sid, b'r' * 20, end_stream=True)
+            self.flush(conn)
+            self.assertEqual(await asyncio.wait_for(channel.receive(), .5), b'r' * 40)
+            await asyncio.gather(*list(channel.pending))
+            self.assertEqual(trace.received, 40)
+            self.assertEqual(trace.phase, 'complete')
+            self.assertFalse(channel.request_traces)
+            self.assertIs(channel.recent_traces[-1], trace)
+            self.lane.log_flow(trace.finished_at)
+            await channel.close()
+        output = '\n'.join(captured.output)
+        self.assertIn('http_down=40', output)
+        self.assertIn('small=1 large=0', output)
+        self.assertNotIn('private-payload', output)
+        self.assertFalse(channel.recent_traces)
+
+    async def test_stale_recovery_slot_replaced_without_resetting_original_request(self):
+        channel = _HttpChannel(self.lane, 1, 'recovery-wire-test')
+        bodies = [b'samekey!' + bytes([index]) * 32 for index in range(3)]
+        originals = []
+        for body in bodies[:2]:
+            await channel.send(body, False)
+            originals.append(await asyncio.wait_for(self.requests.get(), 1))
+
+        with patch('proxy.cf_h2.REPLAY_REQUEST_SECONDS', .02), \
+                patch('proxy.cf_h2.REPLAY_CHECK_SECONDS', .005), \
+                patch('proxy.cf_h2.REPLAY_SLOT_SECONDS', .1):
+            recovery = asyncio.create_task(channel._recover())
+            self.tasks.append(recovery)
+            old_poll = await asyncio.wait_for(self.requests.get(), 1)
+            other_poll = await asyncio.wait_for(self.requests.get(), 1)
+            self.assertEqual([old_poll[2], other_poll[2]], bodies[:2])
+
+            # The two duplicate requests remain silent. A new native request
+            # also gets no reply until its own duplicate solicits one.
+            await channel.send(bodies[2], False)
+            original = await asyncio.wait_for(self.requests.get(), 1)
+            probe = await asyncio.wait_for(self.requests.get(), 1)
+            self.assertEqual(probe[2], bodies[2])
+            reset_sid, reset_code = await asyncio.wait_for(self.resets.get(), 1)
+            self.assertEqual((reset_sid, reset_code), (old_poll[1], 8))
+            self.assertIn(original[1], self.transport.connection.streams)
+            for request in originals:
+                self.assertIn(request[1], self.transport.connection.streams)
+            self.assertIn(other_poll[1], self.transport.connection.streams)
+            self.assertEqual(len(channel.replay_pending), 2)
+            self.respond(probe, b'r' * 40)
+            self.assertEqual(await asyncio.wait_for(channel.receive(), 1), b'r' * 40)
+            # The original stream remains usable, even after recovery.
+            self.respond(original, b'o' * 40)
+            self.assertEqual(await asyncio.wait_for(channel.receive(), 1), b'o' * 40)
+            self.assertEqual(channel.replay_rotations, 1)
+            self.assertEqual(self.lane.failed_until, 0)
+            self.assertEqual(len(self.connections), 1)
+
+    async def test_normal_long_poll_survives_six_seconds_and_delivers_late_reply(self):
+        channel = _HttpChannel(self.lane, 1, 'normal-long-poll-test')
+        bodies = [b'samekey!' + bytes([index]) * 32 for index in range(3)]
+        for body in bodies[:2]:
+            await channel.send(body, False)
+            self.respond(await asyncio.wait_for(self.requests.get(), 1))
+            await asyncio.wait_for(channel.receive(), 1)
+        with patch('proxy.cf_h2.REPLAY_IDLE_SECONDS', .01), \
+                patch('proxy.cf_h2.REPLAY_RETRY_SECONDS', .01), \
+                patch('proxy.cf_h2.REPLAY_CHECK_SECONDS', .005):
+            recovery = asyncio.create_task(channel._recover())
+            self.tasks.append(recovery)
+            poll = await asyncio.wait_for(self.requests.get(), 1)
+            # Advance request ages, keeping the real TLS/H2 streams open.
+            # Six seconds is still inside Telegram's default 25s HTTP wait.
+            for task in channel.replay_pending:
+                channel.pending_since[task] = time.monotonic() - 6
+                channel.sent_since[task] = time.monotonic() - 6
+            await asyncio.sleep(.03)
+            self.assertTrue(self.resets.empty(), 'normal HTTP wait was reset before its delayed result')
+            await channel.send(bodies[2], False)
+            fresh = await asyncio.wait_for(self.requests.get(), 1)
+            self.assertEqual(fresh[2], bodies[2])
+            self.respond(fresh, b'n' * 40)
+            self.assertEqual(await asyncio.wait_for(channel.receive(), 1), b'n' * 40)
+            self.respond(poll, b'd' * (128 * 1024))
+            self.assertEqual(await asyncio.wait_for(channel.receive(), 1), b'd' * (128 * 1024))
+            self.assertEqual(channel.replay_rotations, 0)
+            self.assertEqual(channel.replay_requests, 1)
+
+    async def test_replay_rotation_does_not_interrupt_upload(self):
+        channel = _HttpChannel(self.lane, 1, 'replay-upload-test')
+        body = b'samekey!' + b'p' * 32
+        await channel.send(body, False)
+        self.respond(await asyncio.wait_for(self.requests.get(), 1))
+        await asyncio.wait_for(channel.receive(), 1)
+        upload_started, resume = asyncio.Event(), asyncio.Event()
+        send_body = _Connection.send_body
+
+        async def paused_send(connection, request, stream):
+            upload_started.set()
+            await resume.wait()
+            await send_body(connection, request, stream)
+
+        with patch.object(_Connection, 'send_body', paused_send), \
+                patch('proxy.cf_h2.REPLAY_IDLE_SECONDS', .01), \
+                patch('proxy.cf_h2.REPLAY_RETRY_SECONDS', .01), \
+                patch('proxy.cf_h2.REPLAY_CHECK_SECONDS', .005):
+            recovery = asyncio.create_task(channel._recover())
+            self.tasks.append(recovery)
+            await asyncio.wait_for(upload_started.wait(), 1)
+            probe = next(iter(channel.replay_pending))
+            channel.pending_since[probe] = time.monotonic() - 40
+            await asyncio.sleep(.03)
+            self.assertFalse(probe.done())
+            self.assertTrue(self.resets.empty())
+            self.assertEqual(channel.replay_rotations, 0)
+            resume.set()
+            request = await asyncio.wait_for(self.requests.get(), 1)
+            self.respond(request)
+            self.assertEqual(await asyncio.wait_for(channel.receive(), 1), b'r' * 40)
+
+    async def test_recovery_deadline_starts_after_original_upload_finishes(self):
+        channel = _HttpChannel(self.lane, 1, 'unfinished-upload-test')
+        body = b'samekey!' + b'q' * 32
+        started, resume = asyncio.Event(), asyncio.Event()
+        send_body = _Connection.send_body
+
+        async def paused_send(connection, request, stream):
+            started.set()
+            await resume.wait()
+            await send_body(connection, request, stream)
+
+        with patch.object(_Connection, 'send_body', paused_send), \
+                patch('proxy.cf_h2.REPLAY_REQUEST_SECONDS', .02), \
+                patch('proxy.cf_h2.REPLAY_CHECK_SECONDS', .005):
+            await channel.send(body, False)
+            await asyncio.wait_for(started.wait(), .5)
+            self.tasks.append(asyncio.create_task(channel._recover()))
+            await asyncio.sleep(.06)
+            self.assertEqual(channel.requests, 1, 'unsent original was duplicated')
+            self.assertTrue(self.requests.empty())
+            resume.set()
+            original = await asyncio.wait_for(self.requests.get(), .5)
+            probe = await asyncio.wait_for(self.requests.get(), .5)
+            self.assertEqual(original[2], body)
+            self.assertEqual(probe[2], body)
+            self.respond(probe)
+            self.assertEqual(await asyncio.wait_for(channel.receive(), .5), b'r' * 40)
+            self.assertTrue(self.resets.empty())
+
+    async def test_body_tail_does_not_wait_for_unrelated_long_polls(self):
+        slow1, _ = await self.post()
+        slow2, _ = await self.post()
+        fast, request = await self.post()
+        self.respond(request, b'a' * 20, end=False, length=40)
+        await asyncio.sleep(.03)
+        self.assertFalse(fast.done())
+        conn, sid, _ = request
+        conn.h2.send_data(sid, b'b' * 20, end_stream=True)
+        self.flush(conn)
+        self.assertEqual(await asyncio.wait_for(fast, .5), b'a' * 20 + b'b' * 20)
+        self.assertFalse(slow1.done())
+        self.assertFalse(slow2.done())
+
+    async def test_receiving_body_is_marked_before_request_coroutine_resumes(self):
+        channel = _HttpChannel(self.lane, 1, 'early-headers-test')
+        resume = asyncio.Event()
+        send_body = _Connection.send_body
+
+        async def paused_send(connection, request, stream):
+            await send_body(connection, request, stream)
+            await resume.wait()
+
+        with patch.object(_Connection, 'send_body', paused_send):
+            await channel.send(b'k' * 40, False)
+            request = await asyncio.wait_for(self.requests.get(), 1)
+            self.respond(request, b'a' * 20, end=False, length=40)
+
+            async def headers_received():
+                while channel.receiving_http == 0:
+                    await asyncio.sleep(.005)
+            try:
+                await asyncio.wait_for(headers_received(), .5)
+                self.assertEqual(channel.receiving_http, 1)
+                self.assertTrue(channel.queue.empty())
+            finally:
+                resume.set()
+            conn, sid, _ = request
+            conn.h2.send_data(sid, b'b' * 20, end_stream=True)
+            self.flush(conn)
+            self.assertEqual(await asyncio.wait_for(channel.receive(), 1), b'a' * 20 + b'b' * 20)
+            self.assertEqual(channel.receiving_http, 0)
+
+    async def test_native_channel_delivers_ready_reply_while_large_body_is_pending(self):
+        class IdentityCipher:
+            def update(self, data):
+                return data
+
+        class Writer:
+            def __init__(self):
+                self.data = bytearray()
+                self.changed = asyncio.Event()
+
+            def write(self, data):
+                self.data.extend(data)
+                self.changed.set()
+
+            async def drain(self):
+                pass
+
+        reader, writer = asyncio.StreamReader(), Writer()
+        channel = _HttpChannel(self.lane, 1, 'wire-test')
+        ctx = SimpleNamespace(clt_enc=IdentityCipher(), clt_dec=IdentityCipher())
+        bridge = asyncio.create_task(bridge_h2(reader, writer, channel, ctx, PROTO_TAG_INTERMEDIATE))
+        self.tasks.append(bridge)
+        reader.feed_data(struct.pack('<I', 40) + b'x' * 40)
+        slow = await asyncio.wait_for(self.requests.get(), 2)
+        self.respond(slow, b'a' * 65536, end=False, length=65576)
+        reader.feed_data(struct.pack('<I', 40) + b'y' * 40)
+        fast = await asyncio.wait_for(self.requests.get(), 2)
+        self.respond(fast, b'f' * 40)
+        await asyncio.wait_for(writer.changed.wait(), .5)
+        self.assertEqual(writer.data, struct.pack('<I', 40) + b'f' * 40)
+        writer.changed.clear()
+        conn, sid, _ = slow
+        # Finish after the first 64 KiB have passed through flow control.
+        for _ in range(100):
+            if sid not in conn.pending:
+                break
+            await asyncio.sleep(.001)
+        self.assertNotIn(sid, conn.pending)
+        conn.pending[sid] = (b'b' * 40, True)
+        self.flush(conn)
+        await asyncio.wait_for(writer.changed.wait(), .5)
+        self.assertEqual(writer.data, struct.pack('<I', 40) + b'f' * 40
+                         + struct.pack('<I', 65576) + b'a' * 65536 + b'b' * 40)
+
+    async def test_parallel_first_windows_do_not_wait_for_connection_credit(self):
+        first, request1 = await self.post()
+        second, request2 = await self.post()
+        self.respond(request1, b'a' * 65532)
+        self.respond(request2, b'b' * 65532)
+        # No event-loop yield between responses: the server has not received
+        # any WINDOW_UPDATE credit for either body yet.
+        self.assertFalse(request1[0].pending)
+        self.assertEqual(await first, b'a' * 65532)
+        self.assertEqual(await second, b'b' * 65532)
+
+    async def test_file_chunk_fits_first_window_without_extra_network_round_trip(self):
+        # A 128 KiB file part plus its MTProto envelope must fit before any
+        # WINDOW_UPDATE arrives. This matters even when the consumer is fast.
+        body = b'f' * (128 * 1024 + 168)
+        task, request = await self.post()
+        self.respond(request, body)
+        self.assertNotIn(request[1], list(request[0].pending))
+        self.assertEqual(await asyncio.wait_for(task, 1), body)
+
+    async def test_stalled_body_does_not_block_recovery_of_another_native_request(self):
+        channel = _HttpChannel(self.lane, 1, 'stalled-body-test')
+        await channel.send(b'k' * 8 + b'a' * 32, False)
+        slow = await asyncio.wait_for(self.requests.get(), 1)
+        self.respond(slow, b'a' * 20, end=False, length=40)
+        fresh_body = b'k' * 8 + b'b' * 32
+        await channel.send(fresh_body, False)
+        original = await asyncio.wait_for(self.requests.get(), 1)
+        with patch('proxy.cf_h2.REPLAY_IDLE_SECONDS', .01), \
+                patch('proxy.cf_h2.REPLAY_CHECK_SECONDS', .005), \
+                patch('proxy.cf_h2.REPLAY_REQUEST_SECONDS', .05):
+            recovery = asyncio.create_task(channel._recover())
+            self.tasks.append(recovery)
+            probe = await asyncio.wait_for(self.requests.get(), .5)
+            self.assertEqual(probe[2], fresh_body)
+            self.respond(probe, b'r' * 40)
+            self.assertEqual(await asyncio.wait_for(channel.receive(), .5), b'r' * 40)
+            self.assertIn(slow[1], self.transport.connection.streams)
+            self.assertIn(original[1], self.transport.connection.streams)
+            self.assertTrue(self.resets.empty())
+            conn, sid, _ = slow
+            conn.h2.send_data(sid, b'a' * 20, end_stream=True)
+            self.flush(conn)
+            self.assertEqual(await asyncio.wait_for(channel.receive(), .5), b'a' * 40)
+            self.assertEqual(self.lane.failed_until, 0)
+
+    async def test_first_recovery_without_result_does_not_end_recovery(self):
+        channel = _HttpChannel(self.lane, 1, 'second-recovery-test')
+        body = b'k' * 40
+        await channel.send(body, False)
+        self.respond(await asyncio.wait_for(self.requests.get(), 1), b'')
+        with patch('proxy.cf_h2.REPLAY_IDLE_SECONDS', .01), \
+                patch('proxy.cf_h2.REPLAY_CHECK_SECONDS', .005), \
+                patch('proxy.cf_h2.REPLAY_RETRY_SECONDS', .02):
+            recovery = asyncio.create_task(channel._recover())
+            self.tasks.append(recovery)
+            first = await asyncio.wait_for(self.requests.get(), .5)
+            self.respond(first, b'')
+            second = await asyncio.wait_for(self.requests.get(), .5)
+            self.assertEqual(first[2], body)
+            self.assertEqual(second[2], body)
+            self.respond(second, b'r' * 40)
+            self.assertEqual(await asyncio.wait_for(channel.receive(), .5), b'r' * 40)
+            self.assertEqual(channel.replay_requests, 2)
+            self.assertEqual(len(self.connections), 1)
+
+    async def test_completed_http_ack_does_not_add_idle_delay_to_each_chunk(self):
+        channel = _HttpChannel(self.lane, 1, 'chunked-download-poll-test')
+        # An HTTP reply can be only an MTProto acknowledgement. The queued
+        # file chunk needs another HTTP receiver, with no new native upload.
+        # Keep the production 1s idle threshold; shorten only scheduler ticks.
+        with patch('proxy.cf_h2.REPLAY_CHECK_SECONDS', .005):
+            self.tasks.append(asyncio.create_task(channel._recover()))
+            for index in range(4):
+                body = b'samekey!' + bytes([index]) * 32
+                await channel.send(body, False)
+                original = await asyncio.wait_for(self.requests.get(), 1)
+                self.assertEqual(original[2], body)
+                self.respond(original, b'a' * 40)
+                self.assertEqual(await asyncio.wait_for(channel.receive(), 1), b'a' * 40)
+                poll = await asyncio.wait_for(self.requests.get(), .5)
+                self.assertEqual(poll[2], body)
+                reply = bytes([index]) * (128 * 1024)
+                self.respond(poll, reply)
+                self.assertEqual(await asyncio.wait_for(channel.receive(), 1), reply)
+            self.assertTrue(self.resets.empty())
+            self.assertEqual(len(self.connections), 1)
+
+    async def test_drained_receiver_wakes_first_poll_without_waiting_for_timer(self):
+        channel = _HttpChannel(self.lane, 1, 'receiver-wakeup-test')
+        body = b'samekey!' + b'q' * 32
+        with patch('proxy.cf_h2.REPLAY_CHECK_SECONDS', 10):
+            self.tasks.append(asyncio.create_task(channel._recover()))
+            await channel.send(body, False)
+            original = await asyncio.wait_for(self.requests.get(), 1)
+            self.respond(original, b'a' * 40)
+            self.assertEqual(await asyncio.wait_for(channel.receive(), 1), b'a' * 40)
+            poll = await asyncio.wait_for(self.requests.get(), .5)
+            self.assertEqual(poll[2], body)
+            self.respond(poll, b'f' * 40)
+            self.assertEqual(await asyncio.wait_for(channel.receive(), .5), b'f' * 40)
+            await asyncio.sleep(.03)
+            self.assertTrue(self.requests.empty(), 'a wakeup bypassed the retry backoff')
+            self.assertEqual(channel.replay_requests, 1)
+
+    async def test_native_write_completion_wakes_poll_but_backpressure_does_not(self):
+        class IdentityCipher:
+            def update(self, data):
+                return data
+
+        class Writer:
+            def __init__(self):
+                self.started, self.resume = asyncio.Event(), asyncio.Event()
+
+            def write(self, data):
+                self.started.set()
+
+            async def drain(self):
+                await self.resume.wait()
+
+        channel = _HttpChannel(self.lane, 1, 'native-wakeup-test')
+        reader, writer = asyncio.StreamReader(), Writer()
+        ctx = SimpleNamespace(clt_enc=IdentityCipher(), clt_dec=IdentityCipher())
+        body = b'samekey!' + b'q' * 32
+        with patch('proxy.cf_h2.REPLAY_CHECK_SECONDS', 10):
+            self.tasks.append(asyncio.create_task(
+                bridge_h2(reader, writer, channel, ctx, PROTO_TAG_INTERMEDIATE)))
+            reader.feed_data(struct.pack('<I', len(body)) + body)
+            original = await asyncio.wait_for(self.requests.get(), 1)
+            self.respond(original)
+            await asyncio.wait_for(writer.started.wait(), .5)
+            await asyncio.sleep(.03)
+            self.assertTrue(self.requests.empty())
+            writer.resume.set()
+            poll = await asyncio.wait_for(self.requests.get(), .5)
+            self.assertEqual(poll[2], body)
+            self.assertIsNone(channel.delivering_since)
+            self.assertEqual(channel.down, 40)
+
+    async def test_idle_poll_preserves_slot_for_new_overdue_request(self):
+        channel = _HttpChannel(self.lane, 1, 'idle-poll-capacity-test')
+        bodies = [b'samekey!' + bytes([index]) * 32 for index in range(4)]
+        for body in bodies[:3]:
+            await channel.send(body, False)
+            self.respond(await asyncio.wait_for(self.requests.get(), 1))
+            await asyncio.wait_for(channel.receive(), 1)
+        with patch('proxy.cf_h2.REPLAY_IDLE_SECONDS', .01), \
+                patch('proxy.cf_h2.REPLAY_REQUEST_SECONDS', .06), \
+                patch('proxy.cf_h2.REPLAY_CHECK_SECONDS', .005):
+            self.tasks.append(asyncio.create_task(channel._recover()))
+            poll = await asyncio.wait_for(self.requests.get(), .5)
+            self.assertEqual(poll[2], bodies[2])
+            await asyncio.sleep(.04)
+            self.assertTrue(self.requests.empty(), 'completed history occupied the spare recovery slot')
+            await channel.send(bodies[3], False)
+            original = await asyncio.wait_for(self.requests.get(), .5)
+            recovery = await asyncio.wait_for(self.requests.get(), .5)
+            self.assertEqual(original[2], bodies[3])
+            self.assertEqual(recovery[2], bodies[3])
+            self.respond(recovery, b'n' * 40)
+            self.assertEqual(await asyncio.wait_for(channel.receive(), .5), b'n' * 40)
+            # Neither the old receiver nor the new original was cancelled.
+            self.assertTrue(self.resets.empty())
+            self.respond(poll, b'p' * 40)
+            self.assertEqual(await asyncio.wait_for(channel.receive(), .5), b'p' * 40)
+            self.respond(original, b'o' * 40)
+            self.assertEqual(await asyncio.wait_for(channel.receive(), .5), b'o' * 40)
+            self.assertEqual(channel.replay_rotations, 0)
+
+    async def test_busy_channel_recovers_old_request_despite_newer_replies(self):
+        channel = _HttpChannel(self.lane, 1, 'busy-recovery-test')
+        body = b'samekey!' + b't' * 32
+        original = None
+        recovered = asyncio.Event()
+        chatter_reply = asyncio.Event()
+        chatter_requests = 0
+
+        async def server():
+            nonlocal original, chatter_requests
+            while True:
+                request = await self.requests.get()
+                if request[2] == body:
+                    if original is None:
+                        original = request
+                    else:
+                        self.respond(request, b't' * 40)
+                else:
+                    chatter_requests += 1
+                    self.respond(request, b'c' * 40)
+
+        async def chatter():
+            index = 0
+            while True:
+                chatter_reply.clear()
+                await channel.send(b'samekey!' + struct.pack('<I', index) + b'c' * 28, False)
+                index += 1
+                await chatter_reply.wait()
+                await asyncio.sleep(.005)
+
+        async def receiver():
+            while True:
+                if await channel.receive() == b't' * 40:
+                    recovered.set()
+                else:
+                    chatter_reply.set()
+
+        with patch('proxy.cf_h2.REPLAY_IDLE_SECONDS', .05), \
+                patch('proxy.cf_h2.REPLAY_CHECK_SECONDS', .005), \
+                patch('proxy.cf_h2.REPLAY_REQUEST_SECONDS', .12), \
+                patch('proxy.cf_h2.REPLAY_HISTORY_PACKETS', 3):
+            await channel.send(body, False)
+            self.tasks.extend(asyncio.create_task(coro) for coro in (
+                server(), chatter(), receiver(), channel._recover()))
+            await asyncio.wait_for(recovered.wait(), .6)
+            self.assertGreater(chatter_requests, 3)
+            self.assertLessEqual(len(channel.replay_history), 3)
+            self.assertIn(original[1], self.transport.connection.streams)
+            self.assertTrue(self.resets.empty())
+            self.assertEqual(len(self.connections), 1)
+
+    async def test_receiving_other_body_does_not_mask_overdue_request(self):
+        channel = _HttpChannel(self.lane, 1, 'parallel-body-recovery-test')
+        body = b'samekey!' + b't' * 32
+        await channel.send(body, False)
+        original = await asyncio.wait_for(self.requests.get(), 1)
+        await channel.send(b'samekey!' + b'b' * 32, False)
+        other = await asyncio.wait_for(self.requests.get(), 1)
+        self.respond(other, b'b' * 4, end=False, length=400)
+
+        async def body_progress():
+            for _ in range(90):
+                await asyncio.sleep(.01)
+                conn, sid, _ = other
+                conn.h2.send_data(sid, b'b' * 4)
+                self.flush(conn)
+
+        with patch('proxy.cf_h2.REPLAY_IDLE_SECONDS', .03), \
+                patch('proxy.cf_h2.REPLAY_CHECK_SECONDS', .005), \
+                patch('proxy.cf_h2.REPLAY_REQUEST_SECONDS', .07):
+            self.tasks.extend(asyncio.create_task(coro) for coro in (body_progress(), channel._recover()))
+            probe = await asyncio.wait_for(self.requests.get(), .5)
+            self.assertEqual(probe[2], body)
+            self.assertEqual(channel.receiving_http, 1)
+            self.respond(probe, b't' * 40)
+            self.assertEqual(await asyncio.wait_for(channel.receive(), .5), b't' * 40)
+            self.assertIn(original[1], self.transport.connection.streams)
+            self.assertIn(other[1], self.transport.connection.streams)
+            self.assertTrue(self.resets.empty())
+
+    async def test_silent_recovery_polls_refresh_without_new_native_input(self):
+        channel = _HttpChannel(self.lane, 1, 'refresh-recovery-test')
+        bodies = [b'samekey!' + bytes([index]) * 32 for index in range(2)]
+        for body in bodies:
+            await channel.send(body, False)
+            await asyncio.wait_for(self.requests.get(), 1)
+        with patch('proxy.cf_h2.REPLAY_REQUEST_SECONDS', .01), \
+                patch('proxy.cf_h2.REPLAY_CHECK_SECONDS', .005), \
+                patch('proxy.cf_h2.REPLAY_SLOT_SECONDS', .2), \
+                patch('proxy.cf_h2.REPLAY_RETRY_SECONDS', .02):
+            recovery = asyncio.create_task(channel._recover())
+            self.tasks.append(recovery)
+            latest = await asyncio.wait_for(self.requests.get(), .5)
+            older = await asyncio.wait_for(self.requests.get(), .5)
+            channel.sent_since[channel.replay_history[0].pending] = time.monotonic() - 1
+            refreshed = await asyncio.wait_for(self.requests.get(), .5)
+            self.assertEqual(refreshed[2], bodies[0])
+            self.assertEqual(latest[2], bodies[0])
+            self.assertEqual(await asyncio.wait_for(self.resets.get(), .5), (latest[1], 8))
+            self.assertIn(older[1], self.transport.connection.streams)
+            self.assertEqual(len(channel.replay_pending), 2)
+            self.respond(refreshed)
+            self.assertEqual(await asyncio.wait_for(channel.receive(), .5), b'r' * 40)
+            self.assertEqual(channel.replay_rotations, 1)
+
+    async def test_recovery_never_resets_replays_with_partially_received_bodies(self):
+        channel = _HttpChannel(self.lane, 1, 'protected-body-test')
+        bodies = [b'k' * 8 + bytes([index]) * 32 for index in range(3)]
+        for body in bodies[:2]:
+            await channel.send(body, False)
+            await asyncio.wait_for(self.requests.get(), 1)
+        with patch('proxy.cf_h2.REPLAY_REQUEST_SECONDS', .01), \
+                patch('proxy.cf_h2.REPLAY_CHECK_SECONDS', .005), \
+                patch('proxy.cf_h2.REPLAY_SLOT_SECONDS', .04):
+            recovery = asyncio.create_task(channel._recover())
+            self.tasks.append(recovery)
+            probes = [await asyncio.wait_for(self.requests.get(), 1) for _ in range(2)]
+            for probe in probes:
+                self.respond(probe, b'p' * 20, end=False, length=40)
+            await channel.send(bodies[2], False)
+            original = await asyncio.wait_for(self.requests.get(), 1)
+            await asyncio.sleep(.1)
+            self.assertEqual(channel.receiving_http, 2)
+            self.assertTrue(self.requests.empty())
+            self.assertTrue(self.resets.empty())
+            self.assertIn(original[1], self.transport.connection.streams)
+            # Finishing a partial response frees a slot without throwing its
+            # MTProto result away, and recovery can then use that slot.
+            conn, sid, _ = probes[0]
+            conn.h2.send_data(sid, b'p' * 20, end_stream=True)
+            self.flush(conn)
+            self.assertEqual(await asyncio.wait_for(channel.receive(), 1), b'p' * 40)
+            fresh_probe = await asyncio.wait_for(self.requests.get(), .5)
+            self.assertEqual(fresh_probe[2], bodies[2])
+            self.respond(fresh_probe)
+            self.assertEqual(await asyncio.wait_for(channel.receive(), 1), b'r' * 40)
+            self.assertEqual(channel.replay_rotations, 0)
+
+    async def test_preflight_and_posts_use_same_tls_connection(self):
+        preflight = asyncio.create_task(self.lane._preflight())
+        self.tasks.append(preflight)
+        request = await asyncio.wait_for(self.requests.get(), 2)
+        self.respond(request, b'', status=501)
+        await preflight
+        task, request = await self.post()
+        self.respond(request)
+        await task
+        self.assertEqual(self.lane.tcp_connections, 1)
+        self.assertEqual(len(self.connections), 1)
+
+    async def test_concurrent_waiters_respect_remote_stream_limit(self):
+        self.server_max_streams = 2
+        pending = [asyncio.create_task(self.lane._post(b'x' * 40, i)) for i in range(12)]
+        self.tasks.extend(pending)
+        for _ in range(6):
+            batch = [await asyncio.wait_for(self.requests.get(), 2) for _ in range(2)]
+            await asyncio.sleep(.01)
+            self.assertTrue(self.requests.empty())
+            for request in batch:
+                self.respond(request)
+        self.assertEqual(await asyncio.gather(*pending), [b'r' * 40] * 12)
+        self.assertEqual(len(self.connections), 1)
+
+    async def test_cancellation_resets_stream_and_reuses_server_capacity(self):
+        self.server_max_streams = 2
+        healthy, healthy_request = await self.post()
+        for _ in range(5):
+            task, request = await self.post()
+            task.cancel()
+            await asyncio.gather(task, return_exceptions=True)
+            self.assertEqual(await asyncio.wait_for(self.resets.get(), .5), (request[1], 8))
+        self.respond(healthy_request)
+        self.assertEqual(await asyncio.wait_for(healthy, .5), b'r' * 40)
+        self.assertEqual(len(self.connections), 1)
+
+    async def test_cancel_while_waiting_for_stream_slot_does_not_leak_capacity(self):
+        self.server_max_streams = 1
+        first, request = await self.post()
+        waiting = asyncio.create_task(self.lane._post(b'w' * 40, 2))
+        self.tasks.append(waiting)
+        await asyncio.sleep(.03)
+        self.assertTrue(self.requests.empty())
+        waiting.cancel()
+        await asyncio.gather(waiting, return_exceptions=True)
+        self.respond(request)
+        await first
+        next_task, next_request = await self.post()
+        self.respond(next_request)
+        await next_task
+        self.assertFalse(self.transport.connection.streams)
+
+    async def test_large_upload_and_download_cross_flow_control_windows(self):
+        body = b'x' * (256 * 1024)
+        task, request = await self.post(body)
+        self.respond(request, body)
+        self.assertEqual(await asyncio.wait_for(task, 2), body)
+        self.assertEqual(len(self.connections), 1)
+
+    async def test_stalled_body_consumer_does_not_block_other_streams(self):
+        # Keep a streaming response open without reading it. Its window must
+        # bound buffering while other responses can exceed a connection window.
+        opening = asyncio.create_task(self.lane.client.send(
+            self.lane.client.build_request('POST', 'https://' + self.lane.host + '/api', content=b'x' * 40),
+            stream=True))
+        self.tasks.append(opening)
+        request = await asyncio.wait_for(self.requests.get(), 2)
+        self.respond(request, b's' * (1024 * 1024))
+        response = await opening
+        fast, fast_request = await self.post()
+        self.respond(fast_request, b'f' * (128 * 1024))
+        self.assertEqual(await asyncio.wait_for(fast, .5), b'f' * (128 * 1024))
+        stream = self.transport.connection.streams[request[1]]
+        self.assertLessEqual(sum(len(data) for data, _ in stream.chunks), STREAM_RECEIVE_WINDOW)
+        await response.aclose()
+        self.assertEqual(await asyncio.wait_for(self.resets.get(), .5), (request[1], 8))
+
+    async def test_read_timeout_resets_only_its_stream(self):
+        self.lane.client.timeout = httpx.Timeout(2, read=.15)
+        stalled, request = await self.post()
+        healthy, healthy_request = await self.post()
+        self.respond(healthy_request)
+        self.assertEqual(await healthy, b'r' * 40)
+        with self.assertRaises(httpx.ReadTimeout):
+            await stalled
+        self.assertEqual(await asyncio.wait_for(self.resets.get(), .5), (request[1], 8))
+        self.assertIsNone(self.transport.connection.error)
+
+    async def test_peer_reset_does_not_break_other_streams(self):
+        bad, request = await self.post()
+        good, good_request = await self.post()
+        conn, sid, _ = request
+        conn.h2.reset_stream(sid, error_code=8)
+        self.flush(conn)
+        with self.assertRaises(httpx.RemoteProtocolError):
+            await bad
+        self.respond(good_request)
+        self.assertEqual(await good, b'r' * 40)
+
+    async def test_connection_loss_fails_waiters_and_next_request_reconnects(self):
+        tasks = [await self.post() for _ in range(3)]
+        self.connections[0].writer.close()
+        for task, _ in tasks:
+            with self.assertRaises(httpx.HTTPError):
+                await asyncio.wait_for(task, .5)
+        task, request = await self.post()
+        self.respond(request)
+        self.assertEqual(await task, b'r' * 40)
+        self.assertEqual(len(self.connections), 2)
+
+    async def test_goaway_fails_pending_without_replaying_and_reconnects(self):
+        task, request = await self.post()
+        conn, _, _ = request
+        conn.h2.close_connection(error_code=0)
+        self.flush(conn)
+        with self.assertRaisesRegex(httpx.RemoteProtocolError, 'GOAWAY'):
+            await task
+        self.assertTrue(self.requests.empty())
+        next_task, next_request = await self.post()
+        self.respond(next_request)
+        await next_task
+        self.assertEqual(len(self.connections), 2)
+
+    async def test_close_stops_reader_and_wakes_pending_requests(self):
+        task, _ = await self.post()
+        connection = self.transport.connection
+        await self.lane.close()
+        with self.assertRaises(httpx.HTTPError):
+            await asyncio.wait_for(task, .5)
+        self.assertTrue(connection.reader_task.done())
+        self.assertTrue(connection.writer.is_closing())
+
+
+if __name__ == '__main__':
+    unittest.main()

tests/test_pool.py

@@ -20,6 +20,24 @@ def _open_ws():
 
 
 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), keys)
+                self.assertEqual((2, True) in keys, media)
+
     async def test_refills_partially_populated_bucket(self):
         pool = _WsPool()
         key = (2, False)