From 3c54c54d36ae59e5e09885fad3ddf6aa60f9cdea Mon Sep 17 00:00:00 2001 From: Flowseal Date: Tue, 6 Oct 2026 18:44:23 +0300 Subject: [PATCH] Optimizations --- proxy/bridge.py | 21 +++------- proxy/cf_h2.py | 4 +- proxy/h2_transport.py | 19 +++++++-- proxy/raw_websocket.py | 7 ++++ proxy/tg_ws_proxy.py | 4 +- tests/test_bridge.py | 78 +++++++++++++++++++++++++++++++------ tests/test_h2_transport.py | 42 ++++++++++++++++++++ tests/test_raw_websocket.py | 24 ++++++++++++ 8 files changed, 164 insertions(+), 35 deletions(-) diff --git a/proxy/bridge.py b/proxy/bridge.py index da68607..c2db7f6 100644 --- a/proxy/bridge.py +++ b/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: diff --git a/proxy/cf_h2.py b/proxy/cf_h2.py index f299710..96da4f6 100644 --- a/proxy/cf_h2.py +++ b/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), ) diff --git a/proxy/h2_transport.py b/proxy/h2_transport.py index 17bb571..7bcfa50 100644 --- a/proxy/h2_transport.py +++ b/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 diff --git a/proxy/raw_websocket.py b/proxy/raw_websocket.py index fc0ad53..0f66db2 100644 --- a/proxy/raw_websocket.py +++ b/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') diff --git a/proxy/tg_ws_proxy.py b/proxy/tg_ws_proxy.py index e6d16af..a181983 100644 --- a/proxy/tg_ws_proxy.py +++ b/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: diff --git a/tests/test_bridge.py b/tests/test_bridge.py index 3386450..70fcaf9 100644 --- a/tests/test_bridge.py +++ b/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): diff --git a/tests/test_h2_transport.py b/tests/test_h2_transport.py index 41538c2..96e616f 100644 --- a/tests/test_h2_transport.py +++ b/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 diff --git a/tests/test_raw_websocket.py b/tests/test_raw_websocket.py index c1bcb1f..7b22d31 100644 --- a/tests/test_raw_websocket.py +++ b/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'),