Optimizations

proxy/bridge.py

@@ -12,7 +12,6 @@ from .balancer import balancer
 from .config import proxy_config
 from .raw_websocket import RawWebSocket
 from .pool import cf_worker_pool
-from ._aes import Cipher, algorithms, modes
 from .network_debug import WsActivity
 from .cf_h2 import bridge_h2
 
@@ -32,8 +31,6 @@ def reset_tcp_backoff() -> None:
     _tcp_retry_after.clear()
     _tcp_connecting.clear()
 
-ZERO_64 = b'\x00' * 64
-
 
 class CryptoCtx:
     __slots__ = ('clt_dec', 'clt_enc', 'tg_enc', 'tg_dec')
@@ -50,34 +47,26 @@ class MsgSplitter:
     Splits TCP stream data into individual MTProto transport packets
     so each can be sent as a separate WS frame.
     """
-    __slots__ = ('_dec', '_proto', '_cipher_buf', '_plain_buf', '_disabled')
+    __slots__ = ('_proto', '_cipher_buf', '_plain_buf', '_disabled')
 
-    def __init__(self, relay_init: bytes, proto_int: int):
-        cipher = Cipher(algorithms.AES(relay_init[8:40]),
-                        modes.CTR(relay_init[40:56]))
-        self._dec = cipher.encryptor()
-        self._dec.update(ZERO_64)
+    def __init__(self, proto_int: int):
         self._proto = proto_int
         self._cipher_buf = bytearray()
         self._plain_buf = bytearray()
         self._disabled = False
 
-    def split(self, chunk: bytes) -> List[bytes]:
+    def split(self, chunk: bytes, plain: bytes) -> List[bytes]:
         if not chunk:
             return []
         if self._disabled:
             return [chunk]
 
         self._cipher_buf.extend(chunk)
-        self._plain_buf.extend(self._dec.update(chunk))
+        self._plain_buf.extend(plain)
 
         parts = []
         offset = 0
         buf_len = len(self._cipher_buf)
-        # Walk the buffer with an offset instead of deleting each packet from
-        # the front. Front-deletion on a bytearray shifts the remaining bytes,
-        # so a chunk holding many small packets degrades to O(N^2); a single
-        # trailing del keeps splitting O(N).
         while offset < buf_len:
             packet_len = self._next_packet_len(offset, buf_len - offset)
             if packet_len is None:
@@ -365,7 +354,7 @@ async def bridge_ws_reencrypt(reader, writer, ws: RawWebSocket, label,
                 plain = ctx.clt_dec.update(chunk)
                 chunk = ctx.tg_enc.update(plain)
                 if splitter:
-                    parts = splitter.split(chunk)
+                    parts = splitter.split(chunk, plain)
                     if not parts:
                         continue
                     if activity is not None:

proxy/cf_h2.py

@@ -16,7 +16,7 @@ from .h2_transport import H2Transport
 from .stats import stats
 from .utils import (
     DC_DEFAULT_IPS, PROTO_TAG_ABRIDGED,
-    PROTO_TAG_SECURE, create_ssl_context,
+    PROTO_TAG_SECURE,
 )
 
 log = logging.getLogger('tg-mtproto-proxy')
@@ -616,7 +616,7 @@ class _HttpLane:
         self.reply_buffer_bytes = 0
         self.capacity = asyncio.Condition()
         self.client = httpx.AsyncClient(
-            transport=H2Transport(create_ssl_context(), max_streams=MAX_LANE_REQUESTS),
+            transport=H2Transport(max_streams=MAX_LANE_REQUESTS),
             trust_env=False,
             timeout=httpx.Timeout(REQUEST_TIMEOUT, connect=8, pool=REQUEST_TIMEOUT),
         )

proxy/h2_transport.py

@@ -3,6 +3,7 @@ from __future__ import annotations
 import asyncio
 import time
 from collections import deque
+from functools import lru_cache
 
 import httpx
 from h2.config import H2Configuration
@@ -15,10 +16,19 @@ from h2.events import (
 from h2.exceptions import H2Error, StreamClosedError
 from h2.settings import SettingCodes, Settings
 
+from .utils import create_ssl_context
+
 STREAM_RECEIVE_WINDOW = 256 * 1024
 CONNECTION_RECEIVE_WINDOW = 4 * 1024 * 1024
 
 
+@lru_cache(maxsize=1)
+def _default_ssl_context():
+    context = create_ssl_context()
+    context.set_alpn_protocols(['h2'])
+    return context
+
+
 class _ResponseStream(httpx.AsyncByteStream):
     def __init__(self, connection, stream_id, read_timeout):
         self.connection = connection
@@ -278,9 +288,12 @@ class _Connection:
 
 
 class H2Transport(httpx.AsyncBaseTransport):
-    def __init__(self, ssl_context, max_streams=64):
-        self.ssl_context = ssl_context
-        self.ssl_context.set_alpn_protocols(['h2'])
+    def __init__(self, ssl_context=None, max_streams=64):
+        if ssl_context is None:
+            self.ssl_context = _default_ssl_context()
+        else:
+            self.ssl_context = ssl_context
+            self.ssl_context.set_alpn_protocols(['h2'])
         self.max_streams = max_streams
         self.connection = None
         self.origin = None

proxy/raw_websocket.py

@@ -21,6 +21,8 @@ _st_BBQ4s = struct.Struct('>BBQ4s')
 _st_H = struct.Struct('>H')
 _st_Q = struct.Struct('>Q')
 
+_XOR_TABLES = tuple(bytes(value ^ key for value in range(256)) for key in range(256))
+
 _ssl_ctx = create_ssl_context()
 _ssl_ctx_fronting = create_ssl_context(check_hostname=False)
 
@@ -42,6 +44,11 @@ def _xor_mask(data: bytes, mask: bytes) -> bytes:
     if not data:
         return data
     n = len(data)
+    if n >= 512:
+        masked = bytearray(data)
+        for offset in range(4):
+            masked[offset::4] = masked[offset::4].translate(_XOR_TABLES[mask[offset]])
+        return bytes(masked)
     mask_rep = (mask * (n // 4 + 1))[:n]
     return (int.from_bytes(data, 'big') ^
             int.from_bytes(mask_rep, 'big')).to_bytes(n, 'big')

proxy/tg_ws_proxy.py

@@ -299,7 +299,7 @@ async def _handle_client(reader, writer, secret: bytes):
                      label, dc, media_tag)
             splitter_fb = None
             try:
-                splitter_fb = MsgSplitter(relay_init, proto_int)
+                splitter_fb = MsgSplitter(proto_int)
             except Exception:
                 pass
             ok = await do_fallback(
@@ -317,7 +317,7 @@ async def _handle_client(reader, writer, secret: bytes):
 
         splitter = None
         try:
-            splitter = MsgSplitter(relay_init, proto_int)
+            splitter = MsgSplitter(proto_int)
             log.debug("[%s] MsgSplitter activated for proto 0x%08X",
                       label, proto_int)
         except Exception:

tests/test_bridge.py

@@ -3,6 +3,7 @@ import asyncio
 import time
 import unittest
 from types import SimpleNamespace
+from unittest.mock import AsyncMock, Mock
 
 from proxy._aes import Cipher, algorithms, modes
 from proxy.bridge import MsgSplitter, bridge_ws_reencrypt
@@ -40,9 +41,10 @@ def _intermediate(payload: bytes) -> bytes:
 class MsgSplitterTest(unittest.TestCase):
     def _split(self, proto_int, packets, chunk_sizes=None):
         relay_init = _relay_init()
-        splitter = MsgSplitter(relay_init, proto_int)
+        splitter = MsgSplitter(proto_int)
         enc = _encryptor(relay_init)
-        stream = enc.update(b''.join(packets))
+        plain = b''.join(packets)
+        stream = enc.update(plain)
 
         chunks = []
         if chunk_sizes is None:
@@ -56,8 +58,10 @@ class MsgSplitterTest(unittest.TestCase):
                 chunks.append(stream[offset:])
 
         parts = []
+        offset = 0
         for chunk in chunks:
-            parts.extend(splitter.split(chunk))
+            parts.extend(splitter.split(chunk, plain[offset:offset + len(chunk)]))
+            offset += len(chunk)
         return splitter, stream, parts
 
     def test_abridged_stream_splits_into_packets(self):
@@ -94,27 +98,77 @@ class MsgSplitterTest(unittest.TestCase):
         self.assertEqual(len(parts), 8)
 
     def test_empty_chunk_yields_nothing(self):
-        splitter = MsgSplitter(_relay_init(), PROTO_INTERMEDIATE_INT)
-        self.assertEqual(splitter.split(b''), [])
+        splitter = MsgSplitter(PROTO_INTERMEDIATE_INT)
+        self.assertEqual(splitter.split(b'', b''), [])
 
     def test_zero_length_packet_disables_splitting(self):
         relay_init = _relay_init()
-        splitter = MsgSplitter(relay_init, PROTO_INTERMEDIATE_INT)
+        splitter = MsgSplitter(PROTO_INTERMEDIATE_INT)
         enc = _encryptor(relay_init)
-        stream = enc.update((0).to_bytes(4, 'little') + b'tail')
-        parts = splitter.split(stream)
+        plain = (0).to_bytes(4, 'little') + b'tail'
+        stream = enc.update(plain)
+        parts = splitter.split(stream, plain)
         self.assertEqual(parts, [stream])
-        self.assertEqual(splitter.split(b'raw'), [b'raw'])
+        self.assertEqual(splitter.split(b'raw', b'raw'), [b'raw'])
 
     def test_flush_returns_buffered_tail_once(self):
         relay_init = _relay_init()
-        splitter = MsgSplitter(relay_init, PROTO_INTERMEDIATE_INT)
+        splitter = MsgSplitter(PROTO_INTERMEDIATE_INT)
         enc = _encryptor(relay_init)
-        partial = enc.update(_intermediate(b'x' * 32)[:10])
-        self.assertEqual(splitter.split(partial), [])
+        plain = _intermediate(b'x' * 32)[:10]
+        partial = enc.update(plain)
+        self.assertEqual(splitter.split(partial, plain), [])
         self.assertEqual(splitter.flush(), [partial])
         self.assertEqual(splitter.flush(), [])
 
+    def test_extended_and_quick_ack_headers_can_cross_chunk_boundaries(self):
+        for proto, frame in ((PROTO_ABRIDGED_INT, _abridged),
+                             (PROTO_INTERMEDIATE_INT, _intermediate),
+                             (PROTO_PADDED_INTERMEDIATE_INT, _intermediate)):
+            for size in (16, 512, 65536):
+                packet = bytearray(frame(b'p' * size))
+                packet[0 if proto == PROTO_ABRIDGED_INT else 3] |= 0x80
+                packets = [bytes(packet), frame(b't' * 8)]
+                for header_split in (1, 2, 3):
+                    with self.subTest(proto=proto, size=size, header_split=header_split):
+                        _, stream, parts = self._split(
+                            proto, packets, chunk_sizes=[header_split, 1, 11, size - 8])
+                        self.assertEqual(parts, [stream[:len(packet)], stream[len(packet):]])
+
+
+class WsReencryptTest(unittest.IsolatedAsyncioTestCase):
+    async def test_relay_preserves_encrypted_packets_and_incomplete_tail(self):
+        for proto, frame in ((PROTO_ABRIDGED_INT, _abridged),
+                             (PROTO_INTERMEDIATE_INT, _intermediate),
+                             (PROTO_PADDED_INTERMEDIATE_INT, _intermediate)):
+            with self.subTest(proto=proto):
+                packets = [frame(b'a' * 8), frame(b'b' * 16),
+                           frame(b'c' * 70000), frame(b'd' * 32)[:7]]
+                client_init, relay_init = _relay_init(), _relay_init()
+                reader = asyncio.StreamReader()
+                reader.feed_data(_encryptor(client_init).update(b''.join(packets)))
+                reader.feed_eof()
+                sent = []
+
+                async def send(data):
+                    sent.append(data)
+
+                async def send_batch(parts):
+                    sent.extend(parts)
+
+                async def recv():
+                    await asyncio.Future()
+
+                ws = SimpleNamespace(send=send, send_batch=send_batch, recv=recv,
+                                     close=AsyncMock())
+                writer = Mock(wait_closed=AsyncMock())
+                ctx = SimpleNamespace(clt_dec=_encryptor(client_init), tg_enc=_encryptor(relay_init))
+                await asyncio.wait_for(bridge_ws_reencrypt(
+                    reader, writer, ws, 'framing-test', ctx, splitter=MsgSplitter(proto)), 1)
+                decoder = _encryptor(relay_init)
+                self.assertEqual([decoder.update(part) for part in sent], packets)
+                ws.close.assert_awaited_once()
+
 
 class WsDiagnosticsTest(unittest.IsolatedAsyncioTestCase):
     async def test_flow_distinguishes_upstream_send_from_native_backpressure(self):

tests/test_h2_transport.py

@@ -27,6 +27,40 @@ from proxy.config import proxy_config
 from proxy.utils import PROTO_TAG_ABRIDGED, PROTO_TAG_INTERMEDIATE, PROTO_TAG_SECURE
 
 
+class H2ContextTest(unittest.IsolatedAsyncioTestCase):
+    async def test_lanes_reuse_verified_context_without_changing_websocket_contexts(self):
+        from proxy.h2_transport import _default_ssl_context
+        from proxy.raw_websocket import _ssl_ctx, _ssl_ctx_fronting
+        from proxy.utils import create_ssl_context
+
+        _default_ssl_context.cache_clear()
+        self.addCleanup(_default_ssl_context.cache_clear)
+        lanes = []
+        try:
+            with patch('proxy.h2_transport.create_ssl_context', wraps=create_ssl_context) as create:
+                lanes = [_HttpLane('kws4.example.org', 1), _HttpLane('kws2.example.net', 2)]
+                create.assert_called_once_with()
+            first, second = (lane.client._transport.ssl_context for lane in lanes)
+            self.assertIs(first, second)
+            self.assertIsNot(first, _ssl_ctx)
+            self.assertIsNot(first, _ssl_ctx_fronting)
+            self.assertTrue(first.check_hostname)
+            self.assertEqual(first.verify_mode, ssl.CERT_REQUIRED)
+            self.assertTrue(_ssl_ctx.check_hostname)
+            self.assertFalse(_ssl_ctx_fronting.check_hostname)
+            self.assertEqual(set(first.get_ca_certs(binary_form=True)),
+                             set(_ssl_ctx.get_ca_certs(binary_form=True)))
+
+            custom = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT)
+            custom.check_hostname = False
+            transport = H2Transport(custom)
+            self.assertIs(transport.ssl_context, custom)
+            self.assertTrue(first.check_hostname)
+            await transport.aclose()
+        finally:
+            await asyncio.gather(*(lane.close() for lane in lanes))
+
+
 class H2WireTest(unittest.IsolatedAsyncioTestCase):
     @classmethod
     def setUpClass(cls):
@@ -163,6 +197,14 @@ class H2WireTest(unittest.IsolatedAsyncioTestCase):
         await asyncio.gather(slow1, slow2)
         self.assertEqual(len(self.connections), 1)
 
+    async def test_shared_default_context_rejects_untrusted_certificate(self):
+        async with httpx.AsyncClient(transport=H2Transport(), trust_env=False) as client:
+            with self.assertRaises(httpx.ConnectError):
+                await client.get('https://' + self.lane.host + '/api')
+        task, request = await self.post()
+        self.respond(request)
+        self.assertEqual(await asyncio.wait_for(task, 1), b'r' * 40)
+
     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

tests/test_raw_websocket.py

@@ -36,6 +36,15 @@ def _recv(chunks, cls=RawWebSocket):
 
 
 class XorMaskTest(unittest.TestCase):
+    def test_matches_bytewise_xor_at_fast_path_and_frame_boundaries(self):
+        for length in (0, 1, 2, 3, 4, 5, 125, 126, 127, 255, 511, 512, 513, 514, 515,
+                       16384, 65535, 65536, 65537, 1024 * 1024 + 3):
+            data = (bytes(range(256)) * (length // 256 + 1))[:length]
+            for mask in (b'\x00' * 4, b'\xff' * 4, b'\x01\x7f\x80\xff'):
+                with self.subTest(length=length, mask=mask):
+                    expected = bytes(value ^ mask[index % 4] for index, value in enumerate(data))
+                    self.assertEqual(_xor_mask(data, mask), expected)
+
     def test_roundtrip(self):
         data = bytes(range(256)) * 3
         mask = b'\x01\x02\x03\x04'
@@ -46,6 +55,21 @@ class XorMaskTest(unittest.TestCase):
 
 
 class BuildFrameTest(unittest.TestCase):
+    def test_masked_frame_wire_format_across_all_length_encodings(self):
+        mask = b'\x01\x7f\x80\xff'
+        for length in (1, 125, 126, 511, 512, 513, 65535, 65536, 65537):
+            payload = (bytes(range(256)) * (length // 256 + 1))[:length]
+            with self.subTest(length=length), patch('proxy.raw_websocket.os.urandom', return_value=mask):
+                frame = RawWebSocket._build_frame(RawWebSocket.OP_BINARY, payload, mask=True)
+                if length < 126:
+                    header = bytes([0x82, 0x80 | length])
+                elif length < 65536:
+                    header = b'\x82\xfe' + length.to_bytes(2, 'big')
+                else:
+                    header = b'\x82\xff' + length.to_bytes(8, 'big')
+                masked = bytes(value ^ mask[index % 4] for index, value in enumerate(payload))
+                self.assertEqual(frame, header + mask + masked)
+
     def test_short_unmasked_frame(self):
         self.assertEqual(
             RawWebSocket._build_frame(RawWebSocket.OP_BINARY, b'abc'),