Correct handle of the 404 mtproto error

proxy/cf_h2.py

@@ -79,6 +79,14 @@ class _ReplyBufferFull(BufferError):
     """A local memory limit"""
 
 
+class _MTProtoTransportError(Exception):
+    """An endpoint error that the native client must receive, not just EOF."""
+
+    def __init__(self, code: int, source: str):
+        super().__init__('%d (%s)' % (code, source))
+        self.code = code
+
+
 @dataclass
 class _ReplayPacket:
     sent_at: float
@@ -139,6 +147,7 @@ class _HttpChannel:
     def __init__(self, lane, channel_id: int, label: str):
         self.lane, self.channel_id, self.label = lane, channel_id, label
         self.closed = False
+        self.transport_error: Optional[int] = None
         self.pending: Dict[asyncio.Task, int] = {}
         self.pending_since: Dict[asyncio.Task, float] = {}
         self.sent_since: Dict[asyncio.Task, float] = {}
@@ -336,6 +345,8 @@ class _HttpChannel:
     def _fail(self, error: Exception) -> None:
         if self.closed:
             return
+        if isinstance(error, _MTProtoTransportError):
+            self.transport_error = error.code
         self.closed = True
         self.recovery_wakeup.set()
         while not self.queue.empty():
@@ -460,14 +471,27 @@ class _HttpChannel:
             if trace is not None:
                 trace.phase = 'cancelled-' + trace.phase
             raise
+        except _MTProtoTransportError as exc:
+            if trace is not None:
+                trace.phase = 'transport-error'
+            stats.h2_errors += 1
+            if not self.closed:
+                log.warning('[%s] H2 lane=%d channel=%d host=%s MTProto transport error=%s '
+                            'request=%d encrypted=%d bytes=%d replay=%d; reporting to client',
+                            self.label, self.lane.lane_id, self.channel_id, self.lane.host, exc,
+                            trace.request if trace is not None else 0,
+                            body[:8] != b'\x00' * 8, len(body), replay)
+            self._fail(exc)
         except Exception as exc:
             if trace is not None:
                 trace.phase = 'error-' + type(exc).__name__
             stats.h2_errors += 1
             if isinstance(exc, (httpx.HTTPError, OSError, ValueError, BufferError)):
                 log.warning('[%s] H2 lane=%d channel=%d request failed: %s: %s; '
+                            'detail=[%s]; '
                             'closing native channel, reconnect may select another domain',
-                            self.label, self.lane.lane_id, self.channel_id, type(exc).__name__, exc)
+                            self.label, self.lane.lane_id, self.channel_id, type(exc).__name__, exc,
+                            trace.summary(time.monotonic()) if trace is not None else '-')
             else:
                 log.exception('[%s] H2 lane=%d channel=%d unexpected request failure',
                               self.label, self.lane.lane_id, self.channel_id)
@@ -610,8 +634,7 @@ class _HttpLane:
                                      extensions={'trace': self._trace}) as response:
             if response.http_version != 'HTTP/2':
                 raise ValueError('CF /api preflight negotiated ' + response.http_version)
-            if response.status_code in (403, 429) or (response.status_code >= 500
+            if not (200 <= response.status_code < 300 or response.status_code in (405, 501)):
-                                                     and response.status_code != 501):
                 raise ConnectionError('CF /api preflight HTTP %d cf_ray=%s' % (
                     response.status_code, response.headers.get('cf-ray', '-')))
             log.debug('H2 lane=%d ready host=%s status=%d ms=%.0f',
@@ -696,6 +719,9 @@ class _HttpLane:
                 self.max_headers_seconds = max(self.max_headers_seconds, headers_at - started)
                 if response.http_version != 'HTTP/2':
                     raise ValueError('CF /api downgraded to ' + response.http_version)
+                if response.status_code in (403, 404, 429, 444):
+                    raise _MTProtoTransportError(-response.status_code, 'HTTP %d cf_ray=%s' % (
+                        response.status_code, response.headers.get('cf-ray', '-')))
                 if response.status_code != 200:
                     raise ConnectionError('CF /api HTTP %d cf_ray=%s' % (
                         response.status_code, response.headers.get('cf-ray', '-')))
@@ -727,8 +753,9 @@ class _HttpLane:
                 if content and len(content) % 4:
                     raise ValueError('CF /api response is not an aligned MTProto packet')
                 if len(content) == 4:
-                    log.warning('H2 lane=%d channel=%d request=%d MTProto transport error=%d',
+                    code, = struct.unpack('<i', content)
-                                self.lane_id, channel_id, request_id, struct.unpack('<i', content)[0])
+                    if code < 0:
+                        raise _MTProtoTransportError(code, 'HTTP 200 payload')
                 self.responses += 1
                 if channel is not None:
                     channel.http_bytes += received
@@ -750,6 +777,9 @@ class _HttpLane:
                               (headers_at - started) * 1000, body_seconds * 1000,
                               trace.summary(finished) if trace is not None else '-')
                 return bytes(content)
+        except _MTProtoTransportError:
+            self.errors += 1
+            raise
         except _ReplyBufferFull:
             raise
         except (httpx.HTTPError, OSError, ValueError, BufferError):
@@ -929,20 +959,27 @@ async def bridge_h2(reader, writer, channel: _HttpChannel, ctx, tag: bytes) -> N
 
     tasks = [asyncio.create_task(upload()), asyncio.create_task(download()),
              asyncio.create_task(channel._recover())]
+    try:
         try:
             done, _ = await asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED)
+        finally:
+            for task in tasks:
+                task.cancel()
+            await asyncio.gather(*tasks, return_exceptions=True)
+        if channel.transport_error is not None:
+            await write_native(_encode_reply(struct.pack('<i', channel.transport_error), tag))
+            channel.down += 4
+            stats.bytes_down += 4
+        else:
             for task in done:
                 task.result()
     except asyncio.IncompleteReadError:
         log.debug('[%s] H2 channel=%d native EOF', channel.label, channel.channel_id)
     except _NativeDeliveryError as exc:
         log.debug('[%s] H2 channel=%d native closed: %s', channel.label, channel.channel_id, exc)
-    except (OSError, ValueError, BufferError, asyncio.TimeoutError) as exc:
+    except (httpx.HTTPError, OSError, ValueError, BufferError, asyncio.TimeoutError) as exc:
         if not channel.closed:
             log.warning('[%s] H2 channel=%d bridge ended: %s: %s',
                         channel.label, channel.channel_id, type(exc).__name__, exc)
     finally:
-        for task in tasks:
-            task.cancel()
-        await asyncio.gather(*tasks, return_exceptions=True)
         await channel.close()

tests/test_cf_h2.py

@@ -155,16 +155,75 @@ class HttpMultiplexTest(unittest.IsolatedAsyncioTestCase):
         calls = []
         async def handler(request):
             calls.append(request.content)
-            return httpx.Response(403, extensions={'http_version': b'HTTP/2'})
+            return httpx.Response(503, extensions={'http_version': b'HTTP/2'})
         self._transport(handler)
         channel = self._channel()
         await channel.send(b'x' * 40, False)
-        with self.assertRaisesRegex(ConnectionError, '403'):
+        with self.assertRaisesRegex(ConnectionError, '503'):
             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_transport_error_reaches_client_before_close_without_cooling_domain(self):
+        class ErrorBody(httpx.AsyncByteStream):
+            async def __aiter__(self):
+                raise AssertionError('HTTP transport error bodies must be ignored')
+                yield b''
+
+        for tag in (PROTO_TAG_ABRIDGED, PROTO_TAG_INTERMEDIATE, PROTO_TAG_SECURE):
+            for status in (403, 404, 429, 444):
+                with self.subTest(tag=tag, status=status):
+                    calls = []
+
+                    def handler(request):
+                        calls.append(request.content)
+                        if request.content == b'h' * 40:
+                            return httpx.Response(200, content=b'r' * 40,
+                                                  extensions={'http_version': b'HTTP/2'})
+                        return httpx.Response(status, stream=ErrorBody(),
+                                              headers={'content-type': 'text/html'},
+                                              extensions={'http_version': b'HTTP/2'})
+
+                    self._transport(handler)
+                    channel = self._channel()
+                    output = bytearray()
+                    draining, release = asyncio.Event(), asyncio.Event()
+
+                    async def drain():
+                        draining.set()
+                        await release.wait()
+
+                    writer = SimpleNamespace(write=output.extend, drain=drain)
+                    reader = asyncio.StreamReader()
+                    packet = b'k' * 40
+                    reader.feed_data(_encode_reply(packet, tag))
+                    ctx = SimpleNamespace(clt_enc=_IdentityCipher(), clt_dec=_IdentityCipher())
+                    task = asyncio.create_task(bridge_h2(reader, writer, channel, ctx, tag))
+                    try:
+                        await asyncio.wait_for(draining.wait(), 1)
+                        self.assertFalse(task.done())
+                        self.assertFalse(channel.close_done)
+                        if tag == PROTO_TAG_ABRIDGED:
+                            self.assertEqual(output, b'\x01' + struct.pack('<i', -status))
+                        else:
+                            length, = struct.unpack('<I', output[:4])
+                            self.assertEqual(length, len(output) - 4)
+                            self.assertEqual(output[4:8], struct.pack('<i', -status))
+                            self.assertLessEqual(length, 7)
+                        self.assertEqual(self.lane.failed_until, 0)
+                        healthy = self._channel()
+                        await healthy.send(b'h' * 40, False)
+                        self.assertEqual(await asyncio.wait_for(healthy.receive(), 1), b'r' * 40)
+                        self.assertEqual(calls, [packet, b'h' * 40])
+                        self.assertFalse(self.lane.closed)
+                    finally:
+                        release.set()
+                        await asyncio.wait_for(task, 1)
+                    self.assertTrue(channel.close_done)
+                    self.assertFalse(channel.pending)
+                    await healthy.close()
+
     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
@@ -352,6 +411,12 @@ class HttpMultiplexTest(unittest.IsolatedAsyncioTestCase):
         self._transport(lambda request: httpx.Response(501, extensions={'http_version': b'HTTP/2'}))
         await self.lane._preflight()
 
+    async def test_missing_http_endpoint_is_rejected_before_native_requests(self):
+        self._transport(lambda request: httpx.Response(
+            404, extensions={'http_version': b'HTTP/2'}))
+        with self.assertRaisesRegex(ConnectionError, 'preflight HTTP 404'):
+            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)

tests/test_h2_transport.py

@@ -1,7 +1,9 @@
 """Wire-level regressions: real TLS sockets, ALPN and HTTP/2 flow control."""
 import asyncio
 import datetime
+import hashlib
 import ipaddress
+import os
 import ssl
 import struct
 import tempfile
@@ -9,7 +11,7 @@ import time
 import unittest
 from pathlib import Path
 from types import SimpleNamespace
-from unittest.mock import patch
+from unittest.mock import AsyncMock, patch
 
 import httpx
 from h2.config import H2Configuration
@@ -19,7 +21,10 @@ 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
+from proxy import tg_ws_proxy
+from proxy._aes import Cipher, algorithms, modes
+from proxy.config import proxy_config
+from proxy.utils import PROTO_TAG_ABRIDGED, PROTO_TAG_INTERMEDIATE, PROTO_TAG_SECURE
 
 
 class H2WireTest(unittest.IsolatedAsyncioTestCase):
@@ -424,6 +429,184 @@ class H2WireTest(unittest.IsolatedAsyncioTestCase):
         self.assertEqual(await first, b'a' * 65532)
         self.assertEqual(await second, b'b' * 65532)
 
+    async def test_http_404_is_forwarded_while_other_native_channel_keeps_working(self):
+        class Cipher:
+            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)
+
+            async def drain(self):
+                self.changed.set()
+
+        channel = _HttpChannel(self.lane, 1, 'bad-key-test')
+        reader, writer = asyncio.StreamReader(), Writer()
+        ctx = SimpleNamespace(clt_dec=Cipher(), clt_enc=Cipher())
+        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'\x00' * 40)
+        self.respond(await asyncio.wait_for(self.requests.get(), 1), b'p' * 84)
+        await asyncio.wait_for(writer.changed.wait(), 1)
+        reader.feed_data((struct.pack('<I', 40) + b'k' * 40) * 2)
+        bad = await asyncio.wait_for(self.requests.get(), 1)
+        pending = await asyncio.wait_for(self.requests.get(), 1)
+        healthy = _HttpChannel(self.lane, 2, 'healthy-key-test')
+        await healthy.send(b'h' * 40, False)
+        good = await asyncio.wait_for(self.requests.get(), 1)
+        self.respond(bad, b'<html>not found</html>', status=404, end=False)
+        await asyncio.wait_for(bridge, 1)
+        self.assertEqual(writer.data, struct.pack('<I', 84) + b'p' * 84
+                         + struct.pack('<Ii', 4, -404))
+        self.assertTrue(channel.close_done)
+        self.assertFalse(channel.pending)
+        self.assertNotIn(pending[1], self.transport.connection.streams)
+        self.assertFalse(healthy.closed)
+        self.respond(good, b'g' * 40)
+        self.assertEqual(await asyncio.wait_for(healthy.receive(), 1), b'g' * 40)
+        recovered = _HttpChannel(self.lane, 3, 'replacement-key-test')
+        await recovered.send(b'n' * 40, False)
+        self.respond(await asyncio.wait_for(self.requests.get(), 1), b'r' * 40)
+        self.assertEqual(await asyncio.wait_for(recovered.receive(), 1), b'r' * 40)
+        self.assertEqual(len(self.connections), 1)
+        self.assertEqual(self.lane.failed_until, 0)
+
+    async def test_binary_transport_error_is_terminal_without_domain_cooldown(self):
+        channel = _HttpChannel(self.lane, 1, 'binary-error-test')
+        await channel.send(b'k' * 40, False)
+        self.respond(await asyncio.wait_for(self.requests.get(), 1), struct.pack('<i', -404))
+        await asyncio.gather(*list(channel.pending))
+        self.assertTrue(channel.closed)
+        self.assertEqual(channel.transport_error, -404)
+        self.assertEqual(self.lane.failed_until, 0)
+        self.assertEqual(self.lane.reply_buffer_bytes, 0)
+        await channel.close()
+
+    async def test_obfuscated_clients_preserve_packets_through_full_h2_fallback(self):
+        secret = os.urandom(16)
+        channels = []
+        clients = []
+        handlers = []
+
+        async def open_channel(dc, label):
+            self.assertEqual(dc, 2)
+            channel = _HttpChannel(self.lane, len(channels) + 1, label)
+            channels.append(channel)
+            return channel
+
+        def accept(reader, writer):
+            task = asyncio.create_task(tg_ws_proxy._handle_client(reader, writer, secret))
+            handlers.append(task)
+            self.tasks.append(task)
+
+        def client_init(tag):
+            raw = bytearray(os.urandom(64))
+            raw[0] = 0x42
+            raw[56:62] = tag + struct.pack('<h', -2)
+            keys = bytes(raw[8:56])
+            up = Cipher(algorithms.AES(hashlib.sha256(keys[:32] + secret).digest()),
+                        modes.CTR(keys[32:])).encryptor()
+            reverse = keys[::-1]
+            down = Cipher(algorithms.AES(hashlib.sha256(reverse[:32] + secret).digest()),
+                          modes.CTR(reverse[32:])).encryptor()
+            encrypted = up.update(bytes(raw))
+            return bytes(raw[:56]) + encrypted[56:], up, down
+
+        def frame(body, tag, padding):
+            if tag == PROTO_TAG_ABRIDGED:
+                words = len(body) // 4
+                prefix = (bytes([words | 0x80]) if words < 127
+                          else b'\xff' + words.to_bytes(3, 'little'))
+            else:
+                if tag == PROTO_TAG_SECURE:
+                    body += os.urandom(padding)
+                prefix = struct.pack('<I', len(body) | 0x80000000)
+            return prefix + body
+
+        async def check_reply(peer, expected, tag):
+            reader, _, _, down = peer
+
+            async def read(count):
+                return down.update(await asyncio.wait_for(reader.readexactly(count), 2))
+
+            if tag == PROTO_TAG_ABRIDGED:
+                words = (await read(1))[0]
+                if words == 127:
+                    words = int.from_bytes(await read(3), 'little')
+                length = words * 4
+            else:
+                length, = struct.unpack('<I', await read(4))
+            body = await read(length)
+            self.assertEqual(body[:len(expected)], expected)
+            self.assertIn(len(body) - len(expected), range(4) if tag == PROTO_TAG_SECURE else (0,))
+
+        server = await asyncio.start_server(accept, '127.0.0.1', 0)
+        pool = SimpleNamespace(open=AsyncMock(side_effect=open_channel))
+        try:
+            with patch.object(tg_ws_proxy, 'cf_h2_pool', pool), \
+                    patch.object(tg_ws_proxy.ws_pool, 'get', AsyncMock(return_value=None)) as direct, \
+                    patch.multiple(proxy_config, fallback_cfproxy=True, cfproxy_worker_domains=[],
+                                   force_test_dc=False, fake_tls_domain='', proxy_protocol=False):
+                for tag in (PROTO_TAG_ABRIDGED, PROTO_TAG_INTERMEDIATE, PROTO_TAG_SECURE):
+                    with self.subTest(tag=tag.hex()):
+                        peers = []
+                        requests = {}
+                        for index in range(2):
+                            reader, writer = await asyncio.open_connection(
+                                '127.0.0.1', server.sockets[0].getsockname()[1])
+                            clients.append(writer)
+                            init, up, down = client_init(tag)
+                            peer = (reader, writer, up, down)
+                            peers.append(peer)
+                            bodies = [b'\x00' * 8 + os.urandom(8) + struct.pack('<I', 20) + os.urandom(20)]
+                            bodies += [b'samekey!' + os.urandom(size - 8) for size in (40, 131112, 56)]
+                            wire = init + up.update(b''.join(
+                                frame(body, tag, padding) for body, padding in zip(bodies, (15, 1, 7, 0))))
+                            for start, end in ((0, 13), (13, 57), (57, 65), (65, 78), (78, len(wire))):
+                                writer.write(wire[start:end])
+                                await writer.drain()
+                                await asyncio.sleep(0)
+                            expected = {body: packet for packet, body in enumerate(bodies)}
+                            for _ in bodies:
+                                request = await asyncio.wait_for(self.requests.get(), 2)
+                                self.assertTrue(request[2] in expected, 'HTTP body differs from native input')
+                                packet = expected.pop(request[2])
+                                requests[index, packet] = request
+
+                        for index, packet in ((1, 2), (0, 2), (0, 0), (1, 0), (0, 1)):
+                            reply = bytes([0x30 + index, packet]) * 36
+                            self.respond(requests[index, packet], reply)
+                            await check_reply(peers[index], reply, tag)
+                        self.respond(requests[0, 3], b'not found', status=404)
+                        await check_reply(peers[0], struct.pack('<i', -404), tag)
+                        self.assertEqual(await asyncio.wait_for(peers[0][0].read(), 2), b'')
+                        self.respond(requests[1, 1], b'healthy!' * 9)
+                        await check_reply(peers[1], b'healthy!' * 9, tag)
+                        self.respond(requests[1, 3], struct.pack('<i', -429))
+                        await check_reply(peers[1], struct.pack('<i', -429), tag)
+                        self.assertEqual(await asyncio.wait_for(peers[1][0].read(), 2), b'')
+                await asyncio.wait_for(asyncio.gather(*handlers), 2)
+                self.assertEqual(direct.await_count, 6)
+                for call in direct.await_args_list:
+                    self.assertEqual(call.args, (2, True))
+                    self.assertEqual(call.kwargs, {'is_test_dc': False})
+                self.assertEqual(pool.open.await_count, 6)
+                self.assertTrue(all(channel.close_done and not channel.pending for channel in channels))
+                self.assertTrue(all(channel.quick_requests == 4 for channel in channels))
+                self.assertEqual(len(self.connections), 1)
+                self.assertEqual(self.lane.failed_until, 0)
+        finally:
+            for writer in clients:
+                writer.close()
+                await writer.wait_closed()
+            server.close()
+            await server.wait_closed()
+
     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.
@@ -820,6 +1003,47 @@ class H2WireTest(unittest.IsolatedAsyncioTestCase):
         self.respond(good_request)
         self.assertEqual(await good, b'r' * 40)
 
+    async def test_reset_during_response_closes_native_bridge_without_escaping(self):
+        cipher = SimpleNamespace(update=lambda data: data)
+        ctx = SimpleNamespace(clt_dec=cipher, clt_enc=cipher)
+        output = bytearray()
+        writer = SimpleNamespace(write=output.extend, drain=AsyncMock())
+        reader = asyncio.StreamReader()
+        reader.feed_data(struct.pack('<I', 40) + b'k' * 40)
+        channel = _HttpChannel(self.lane, 1, 'reset-during-body')
+
+        with patch.object(channel, '_recover', side_effect=asyncio.Event().wait), \
+                self.assertLogs('tg-mtproto-proxy', level='DEBUG') as captured:
+            bridge = asyncio.create_task(bridge_h2(reader, writer, channel, ctx, PROTO_TAG_INTERMEDIATE))
+            self.tasks.append(bridge)
+            request = await asyncio.wait_for(self.requests.get(), 1)
+            healthy, good_request = await self.post(b'h' * 40)
+            self.respond(request, b'partial-response!!!!', end=False, length=40)
+
+            async def body_received():
+                while self.lane.reply_buffer_bytes != 20:
+                    await asyncio.sleep(.001)
+
+            await asyncio.wait_for(body_received(), 1)
+            conn, sid, _ = request
+            conn.h2.reset_stream(sid, error_code=2)
+            self.flush(conn)
+            await asyncio.wait_for(bridge, 1)
+            self.assertTrue(channel.close_done)
+            self.assertFalse(channel.pending)
+            self.assertIsNone(channel.transport_error)
+            self.assertFalse(output, 'A partial HTTP response must never reach the native client')
+            self.assertEqual(self.lane.reply_buffer_bytes, 0)
+            self.assertIsNone(self.transport.connection.error)
+            self.respond(good_request)
+            self.assertEqual(await asyncio.wait_for(healthy, 1), b'r' * 40)
+            self.assertEqual(len(self.connections), 1)
+        failures = [record for record in captured.records if record.levelname in ('WARNING', 'ERROR')]
+        self.assertEqual(len(failures), 1)
+        self.assertIn('reset: 2', failures[0].getMessage())
+        self.assertIn('down=20', failures[0].getMessage())
+        self.assertIn('sid=1', failures[0].getMessage())
+
     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()