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: