From 48954ff9ab7ab2bcf1e86555162b0f0a683066dd Mon Sep 17 00:00:00 2001 From: Flowseal Date: Sun, 4 Oct 2026 15:57:33 +0300 Subject: [PATCH] multiplexing media through cfproxy --- docs/EN/BuildFromSource.md | 1 + docs/EN/CfProxy.md | 5 + docs/RU/BuildFromSource.md | 1 + docs/RU/CfProxy.md | 6 + packaging/linux.spec | 2 + packaging/macos.spec | 2 + packaging/windows.spec | 2 + proxy/bridge.py | 36 +- proxy/cf_h2.py | 948 +++++++++++++++++++++++++++++++++++++ proxy/config.py | 6 + proxy/h2_transport.py | 332 +++++++++++++ proxy/network_debug.py | 88 ++++ proxy/stats.py | 10 +- proxy/tg_ws_proxy.py | 69 ++- pyproject.toml | 1 + tests/test_bridge.py | 91 +++- tests/test_cf_h2.py | 912 +++++++++++++++++++++++++++++++++++ tests/test_h2_routing.py | 75 +++ tests/test_h2_settings.py | 127 +++++ tests/test_h2_transport.py | 858 +++++++++++++++++++++++++++++++++ tests/test_pool.py | 18 + ui/i18n/en.json | 4 +- ui/i18n/ru.json | 4 +- ui/settings_form.py | 21 +- utils/default_config.py | 1 + utils/tray_common.py | 1 + 26 files changed, 3598 insertions(+), 23 deletions(-) create mode 100644 proxy/cf_h2.py create mode 100644 proxy/h2_transport.py create mode 100644 proxy/network_debug.py create mode 100644 tests/test_cf_h2.py create mode 100644 tests/test_h2_routing.py create mode 100644 tests/test_h2_settings.py create mode 100644 tests/test_h2_transport.py diff --git a/docs/EN/BuildFromSource.md b/docs/EN/BuildFromSource.md index 8ee690a..f439ed8 100644 --- a/docs/EN/BuildFromSource.md +++ b/docs/EN/BuildFromSource.md @@ -49,6 +49,7 @@ tg-ws-proxy [--port PORT] [--host HOST] [--dc-ip DC:IP ...] [-v] | `--secret` | `random` | 32-character hex key for client authorization | | `--dc-ip` | `2:149.154.167.220`, `4:149.154.167.220` | Target IP for DC (can be specified multiple times) | | `--no-cfproxy` | `false` | Disable [Cloudflare proxying](./CfProxy.md) attempts | +| `--no-h2` | `false` | Disable HTTP/2 media multiplexing through CF proxy | | `--cfproxy-domain` | | Specify your own domain for Cloudflare proxying [Learn more](./CfProxy.md). Can be specified multiple times. | | `--cfproxy-worker-domain` | | Cloudflare Worker domain [Learn more](./CfWorker.md). Can be specified multiple times. | | `--no-secure` | `false` | Use 80 port for CF-proxy and CF-worker connections | diff --git a/docs/EN/CfProxy.md b/docs/EN/CfProxy.md index da41b38..5fa3297 100644 --- a/docs/EN/CfProxy.md +++ b/docs/EN/CfProxy.md @@ -26,6 +26,11 @@ Cloudflare limits the number of simultaneous WebSocket (WS) connections. The def 5. In the `TgWsProxy` settings, replace the default domain with your own. +## HTTP/2 media multiplexing + +Multiplexing is enabled by default. Media connections share one TCP/TLS connection to Cloudflare per DC and domain, using `/api` instead of `/apiws`. Main connections and test DCs continue using WS. H2 requires CF proxy and TLS enabled. +Can be disabled with `--no-h2` or by unchecking the **Media multiplexing (HTTP/2)** in the **Cloudflare Proxy** settings. + ## Credits / Acknowledgments - Original Idea: https://github.com/Nekogram/WSProxy diff --git a/docs/RU/BuildFromSource.md b/docs/RU/BuildFromSource.md index 3d01b25..959a21f 100644 --- a/docs/RU/BuildFromSource.md +++ b/docs/RU/BuildFromSource.md @@ -49,6 +49,7 @@ tg-ws-proxy [--port PORT] [--host HOST] [--dc-ip DC:IP ...] [-v] | `--secret` | `random` | 32-значный hex-ключ для авторизации клиентов | | `--dc-ip` | `2:149.154.167.220`, `4:149.154.167.220` | Целевой IP для DC (параметр можно указывать несколько раз) | | `--no-cfproxy` | `false` | Отключить попытку [проксирования через Cloudflare](./CfProxy.md) | +| `--no-h2` | `false` | Отключить мультиплексирование медиа через HTTP/2 в CF-прокси | | `--cfproxy-domain` | | Указать свой домен для проксирования через Cloudflare [Подробнее](./CfProxy.md). Можно указать несколько через повторение аргумента. | | `--cfproxy-worker-domain` | | Домен Cloudflare Worker [Подробнее](./CfWorker.md). Можно указать несколько через повторение аргумента. | | `--no-secure` | `false` | Принудительно использует 80 порт для CF-proxy и CF-worker | diff --git a/docs/RU/CfProxy.md b/docs/RU/CfProxy.md index 59830a1..f5cfcce 100644 --- a/docs/RU/CfProxy.md +++ b/docs/RU/CfProxy.md @@ -26,6 +26,12 @@ Cloudflare имеет лимиты на одновременное количе 5. В настройках `TgWsProxy` замените домен на свой. +## Мультиплексирование медиа через HTTP/2 + +Мультиплексирование включено по умолчанию. Медиа-соединения используют `/api` вместо `/apiws` через одно общее TCP/TLS-соединение до Cloudflare для каждого DC и домена. Основные соединения и тестовые DC продолжают использовать WS. +Чтобы отключить мультиплексирование, добавьте `--no-h2` или снимите галочку +**Мультиплексирование медиа (HTTP/2)** в настройках **Cloudflare Proxy**. + ## Благодарности - Идея: https://github.com/Nekogram/WSProxy diff --git a/packaging/linux.spec b/packaging/linux.spec index e5237b7..f4d8d2c 100644 --- a/packaging/linux.spec +++ b/packaging/linux.spec @@ -27,6 +27,8 @@ a = Analysis( binaries=appindicator_binaries, datas=[(ctk_path, 'customtkinter/'), (_i18n_path, 'ui/i18n')] + certifi_datas, hiddenimports=[ + 'httpx', + 'h2', 'pystray._appindicator', 'PIL._tkinter_finder', 'customtkinter', diff --git a/packaging/macos.spec b/packaging/macos.spec index 711d37c..efbea71 100644 --- a/packaging/macos.spec +++ b/packaging/macos.spec @@ -17,6 +17,8 @@ a = Analysis( binaries=[], datas=[(ctk_path, 'customtkinter/'), (_i18n_path, 'ui/i18n')] + certifi_datas, hiddenimports=[ + 'httpx', + 'h2', 'tkinter', 'customtkinter', 'pystray._darwin', diff --git a/packaging/windows.spec b/packaging/windows.spec index 18e2895..e9e8800 100644 --- a/packaging/windows.spec +++ b/packaging/windows.spec @@ -20,6 +20,8 @@ a = Analysis( binaries=[], datas=[(ctk_path, 'customtkinter/'), (_i18n_path, 'ui/i18n')] + certifi_datas, hiddenimports=[ + 'httpx', + 'h2', 'pystray._win32', 'PIL._tkinter_finder', 'customtkinter', diff --git a/proxy/bridge.py b/proxy/bridge.py index 404fc3a..4ceef41 100644 --- a/proxy/bridge.py +++ b/proxy/bridge.py @@ -1,6 +1,7 @@ import asyncio import logging import struct +import time from typing import List, Optional from urllib.parse import urlencode @@ -12,6 +13,8 @@ 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 log = logging.getLogger('tg-mtproto-proxy') @@ -130,7 +133,7 @@ class MsgSplitter: async def do_fallback(reader, writer, relay_init, label, dc: int, is_test_dc: bool, is_media: bool, media_tag: str, - ctx: CryptoCtx, splitter=None): + ctx: CryptoCtx, splitter=None, *, h2_pool=None, proto_tag=None): ip_table = DC_TEST_IPS if is_test_dc else DC_DEFAULT_IPS fallback_dst = ip_table.get(dc) use_cf = proxy_config.fallback_cfproxy and not is_test_dc @@ -154,6 +157,13 @@ async def do_fallback(reader, writer, relay_init, label, if ok: return True elif method == 'cf': + if is_media and h2_pool is not None and proto_tag is not None: + channel = await h2_pool.open(dc, label) + if channel is not None: + stats.connections_h2 += 1 + stats.connections_cfproxy += 1 + await bridge_h2(reader, writer, channel, ctx, proto_tag) + return True ok = await _cfproxy_fallback( reader, writer, relay_init, label, ctx, dc=dc, is_media=is_media, @@ -231,6 +241,7 @@ async def _cfproxy_fallback(reader, writer, relay_init, label, log.info("[%s] DC%d%s -> trying CF proxy", label, dc, media_tag) + connect_started = time.monotonic() for base_domain in balancer.get_domains_for_dc(dc): domain = f'kws{dc}.{base_domain}' @@ -246,6 +257,8 @@ async def _cfproxy_fallback(reader, writer, relay_init, label, if ws is None: return False + log.debug('[%s] WS CONNECT DC%d%s route=cf setup_ms=%.0f', + label, dc, media_tag, (time.monotonic() - connect_started) * 1000) if chosen_domain and balancer.update_domain_for_dc(dc, chosen_domain): log.info("[%s] Switched active CF domain", label) @@ -290,6 +303,7 @@ async def bridge_ws_reencrypt(reader, writer, ws: RawWebSocket, label, down_packets = 0 start_time = asyncio.get_running_loop().time() close_reason = 'normal' + activity = WsActivity(label, dc_tag) if log.isEnabledFor(logging.DEBUG) else None async def tcp_to_ws(): nonlocal up_bytes, up_packets, close_reason @@ -300,9 +314,15 @@ async def bridge_ws_reencrypt(reader, writer, ws: RawWebSocket, label, if splitter: tail = splitter.flush() if tail: + if activity is not None: + activity.send_started() await ws.send(tail[0]) + if activity is not None: + activity.sent(len(tail[0])) break n = len(chunk) + if activity is not None: + activity.input(n) stats.bytes_up += n up_bytes += n up_packets += 1 @@ -312,12 +332,20 @@ async def bridge_ws_reencrypt(reader, writer, ws: RawWebSocket, label, parts = splitter.split(chunk) if not parts: continue + if activity is not None: + activity.send_started() if len(parts) > 1: await ws.send_batch(parts) else: await ws.send(parts[0]) + if activity is not None: + activity.sent(sum(map(len, parts))) else: + if activity is not None: + activity.send_started() await ws.send(chunk) + if activity is not None: + activity.sent(len(chunk)) except asyncio.CancelledError: return except (ConnectionError, OSError) as e: @@ -336,6 +364,8 @@ async def bridge_ws_reencrypt(reader, writer, ws: RawWebSocket, label, close_reason = 'upstream: ws_close' break n = len(data) + if activity is not None: + activity.received(n) stats.bytes_down += n down_bytes += n down_packets += 1 @@ -343,6 +373,8 @@ async def bridge_ws_reencrypt(reader, writer, ws: RawWebSocket, label, data = ctx.clt_enc.update(plain) writer.write(data) await writer.drain() + if activity is not None: + activity.delivered(n) except asyncio.CancelledError: return except (ConnectionError, OSError) as e: @@ -366,6 +398,8 @@ async def bridge_ws_reencrypt(reader, writer, ws: RawWebSocket, label, except BaseException: pass elapsed = asyncio.get_running_loop().time() - start_time + if activity is not None: + activity.close() log.info("[%s] %s WS session closed (%s): " "^%s (%d pkts) v%s (%d pkts) in %.1fs", label, dc_tag, close_reason, diff --git a/proxy/cf_h2.py b/proxy/cf_h2.py new file mode 100644 index 0000000..36248b3 --- /dev/null +++ b/proxy/cf_h2.py @@ -0,0 +1,948 @@ +from __future__ import annotations + +import asyncio +import logging +import os +import struct +import time +from collections import deque +from dataclasses import dataclass, field +from typing import Dict, List, Optional, Set, Tuple + +import httpx + +from .balancer import balancer +from .h2_transport import H2Transport +from .stats import stats +from .utils import ( + DC_DEFAULT_IPS, PROTO_TAG_ABRIDGED, + PROTO_TAG_SECURE, create_ssl_context, +) + +log = logging.getLogger('tg-mtproto-proxy') +MAX_PACKET = 4 * 1024 * 1024 +MAX_CHANNEL_REQUESTS = 8 +MAX_CHANNEL_BYTES = 8 * 1024 * 1024 +MAX_LANE_REQUESTS = 64 +MAX_LANE_BYTES = 32 * 1024 * 1024 +MAX_LANE_REPLY_BYTES = 32 * 1024 * 1024 +DOMAIN_COOLDOWN = 30.0 +REQUEST_TIMEOUT = 65.0 +SETUP_TIMEOUT = 8.0 +REPLAY_IDLE_SECONDS = 1.0 +REPLAY_REQUEST_SECONDS = 3.0 +REPLAY_CHECK_SECONDS = 0.25 +REPLAY_MAX_AGE = 30.0 +REPLAY_HISTORY_BYTES = 64 * 1024 +REPLAY_HISTORY_PACKETS = 16 +MAX_CHANNEL_REPLAYS = 2 +REPLAY_SLOT_SECONDS = 30.0 +REPLAY_RETRY_SECONDS = 2.0 +REPLAY_MAX_ATTEMPTS = 3 +SLOW_RESPONSE_SECONDS = 3.0 +DIAGNOSTIC_INTERVAL = 5.0 + + +@dataclass +class _RequestTrace: + packet: int + started: float + replay: bool + size: int + request: int = 0 + stream: int = 0 + phase: str = 'queued' + sent_at: Optional[float] = None + headers_at: Optional[float] = None + finished_at: Optional[float] = None + received: int = 0 + cf_ray: str = '-' + + def summary(self, now: float) -> str: + end = self.finished_at if self.finished_at is not None else now + headers = self.headers_at if self.headers_at is not None else end + send_ms = -1 if self.sent_at is None else (self.sent_at - self.started) * 1000 + headers_ms = -1 if self.sent_at is None else max(0, headers - self.sent_at) * 1000 + return ('req=%d pkt=%d sid=%d replay=%d phase=%s age_ms=%.0f ' + 'send_ms=%.0f wait_headers_ms=%.0f up=%d down=%d ray=%s' % ( + self.request, self.packet, self.stream, self.replay, self.phase, + (end - self.started) * 1000, + send_ms, headers_ms, + self.size, self.received, self.cf_ray)) + + +class _NativeDeliveryError(ConnectionError): + """The local receiver disconnected""" + + +class _ReplyBufferFull(BufferError): + """A local memory limit""" + + +@dataclass +class _ReplayPacket: + sent_at: float + body: bytes + sequence: int + attempts: int = 0 + last_attempt_at: float = 0.0 + pending: Optional[asyncio.Task] = None + originals: Set[asyncio.Task] = field(default_factory=set, repr=False) + retired: bool = False + + @property + def replayed(self) -> bool: + return self.attempts > 0 + + +async def _read_packet(reader, decryptor, tag: bytes) -> Tuple[bytes, bool]: + async def plaintext(count: int) -> bytes: + return decryptor.update(await reader.readexactly(count)) + + if tag == PROTO_TAG_ABRIDGED: + first = (await plaintext(1))[0] + quick, words = bool(first & 0x80), first & 0x7f + if words == 0x7f: + words = int.from_bytes(await plaintext(3), 'little') + length = words * 4 + else: + value, = struct.unpack(' 15: + raise ValueError('invalid padded native packet') + body = body[:packet_length] + if len(body) % 4 or len(body) > MAX_PACKET: + raise ValueError('invalid native packet alignment or size') + return body, quick + + +def _encode_reply(body: bytes, tag: bytes) -> bytes: + if not body or len(body) % 4 or len(body) > MAX_PACKET: + raise ValueError('invalid HTTP MTProto response length: %d' % len(body)) + if tag == PROTO_TAG_ABRIDGED: + words = len(body) // 4 + prefix = bytes([words]) if words < 0x7f else b'\x7f' + words.to_bytes(3, 'little') + return prefix + body + padding = os.urandom(os.urandom(1)[0] % 4) if tag == PROTO_TAG_SECURE else b'' + return struct.pack(' int: + return len(self.receiving_tasks) + + @property + def native_backpressure(self) -> bool: + return bool(self.reply_bytes) or self.delivering_since is not None + + def _has_capacity(self, length: int) -> bool: + return len(self.pending) < MAX_CHANNEL_REQUESTS and self.pending_bytes + length <= MAX_CHANNEL_BYTES + + def _wake_receiver(self) -> None: + if not self.closed and not self.pending and not self.native_backpressure: + self.recovery_wakeup.set() + + def _overdue_original(self, entry: _ReplayPacket, now: float) -> bool: + return entry.retired or any(task in self.sent_since and task not in self.receiving_tasks + and now - self.sent_since[task] >= REPLAY_REQUEST_SECONDS + for task in entry.originals) + + def _replay_limit(self, entry: _ReplayPacket) -> int: + return REPLAY_MAX_ATTEMPTS if entry.originals or entry.retired or entry is self.replay_history[-1] else 0 + + def _capacity_candidate(self, now: float): + for task, sent in sorted(self.sent_since.items(), key=lambda pair: ( + pair[0] not in self.replay_pending, pair[1])): + if (task.done() or task in self.receiving_tasks + or now - sent < REPLAY_REQUEST_SECONDS): + continue + if task in self.replay_pending: + return task, None + entry = next((entry for entry in self.replay_history if task in entry.originals), None) + if entry is not None and now - entry.sent_at <= REPLAY_MAX_AGE: + return task, entry + return None + + async def _make_upload_room(self, now: float) -> None: + if self.closed or self.upload_wait_since is None: + return + if self.upload_wait_phase == 'channel-capacity': + if self._has_capacity(self.upload_wait_bytes): + return + channels = [self] + elif self.upload_wait_phase == 'lane-capacity': + if self.lane._has_capacity(self.upload_wait_bytes): + return + channels = sorted(self.lane.channels, key=lambda channel: min( + channel.sent_since.values(), default=now)) + else: + return + for channel in channels: + candidate = channel._capacity_candidate(now) + if candidate is None: + continue + task, entry = candidate + channel.sent_since.pop(task, None) + if entry is not None: + entry.retired = True + trace = channel.request_traces.get(task) + if trace is not None: + trace.phase = 'capacity-' + trace.phase + if now >= self.lane.next_capacity_log: + self.lane.next_capacity_log = now + DIAGNOSTIC_INTERVAL + log.debug('H2 CAPACITY lane=%d waiting_channel=%d release_channel=%d ' + 'reason=%s replay=%d wait_ms=%.0f inflight=%d', + self.lane.lane_id, self.channel_id, channel.channel_id, + self.upload_wait_phase, task in channel.replay_pending, + (now - self.upload_wait_since) * 1000, self.lane.inflight) + task.cancel() + await asyncio.gather(task, return_exceptions=True) + channel.capacity_rotations += 1 + self.lane.capacity_rotations += 1 + return + + def _replay_blocker(self, entry: _ReplayPacket, now: float) -> Optional[str]: + if not 0 <= now - entry.sent_at <= REPLAY_MAX_AGE: + return 'expired' + if entry.attempts >= self._replay_limit(entry): + return 'budget-used' + if self.native_backpressure: + return 'native-backpressure' + if entry.originals.intersection(self.receiving_tasks): + return 'original-body' + if entry.pending in self.receiving_tasks: + return 'replay-body' + if entry.replayed and now - entry.last_attempt_at < REPLAY_RETRY_SECONDS: + return 'backoff-or-progress' + if self._overdue_original(entry, now): + return None + if (entry is not self.replay_history[-1] + or any(task is not entry.pending for task in self.pending)): + return 'receiver-active' + if entry.replayed and now - self.last_progress < REPLAY_IDLE_SECONDS: + return 'backoff-or-progress' + return None + + def _replay_due(self, entry: _ReplayPacket, now: float) -> bool: + return self._replay_blocker(entry, now) is None + + def _wait_state(self, now: float) -> str: + oldest = min(self.pending_since.values(), default=now) + recovery_left = sum(max(0, self._replay_limit(entry) - entry.attempts) + for entry in self.replay_history if now - entry.sent_at <= REPLAY_MAX_AGE) + waiting_headers = sum(task not in self.replay_pending and task not in self.receiving_tasks + and now - since >= REPLAY_REQUEST_SECONDS + for task, since in self.pending_since.items()) + return ('channel=%d pending=%d replays=%d receiving=%d ready_bytes=%d ' + 'native_write=%d idle_ms=%.0f oldest_ms=%.0f native_age_ms=%.0f ' + 'recovery_left=%d waiting_headers=%d' % ( + self.channel_id, len(self.pending), len(self.replay_pending), + self.receiving_http, self.reply_bytes, self.delivering_since is not None, + (now - self.last_progress) * 1000, (now - oldest) * 1000, + (now - self.last_native_send) * 1000, recovery_left, waiting_headers)) + + def _request_detail(self, trace: _RequestTrace, now: float) -> str: + entry = next((entry for entry in self.replay_history if entry.sequence == trace.packet), None) + recovery = 'not-retained' + attempts = '-' + if entry is not None: + attempts = '%d/%d' % (entry.attempts, self._replay_limit(entry)) + blocker = self._replay_blocker(entry, now) + if blocker is not None: + recovery = blocker + elif entry.pending is not None: + recovery = ('replay-body' if entry.pending in self.receiving_tasks else + 'probe-%.0fms' % ((now - self.pending_since.get(entry.pending, now)) * 1000)) + elif len(self.replay_pending) >= MAX_CHANNEL_REPLAYS: + recovery = 'replay-slots' + elif not self._has_capacity(len(entry.body)): + recovery = 'channel-capacity' + elif not self.lane._has_capacity(len(entry.body)): + recovery = 'lane-capacity' + elif self.send_lock.locked(): + recovery = 'send-lock' + else: + recovery = 'eligible' + related = [item for item in self.recent_traces if item.packet == trace.packet][-2:] + return ('channel=%d %s attempts=%s recovery=%s recent=[%s]' % ( + self.channel_id, trace.summary(now), attempts, recovery, + '; '.join(item.summary(now) for item in related))) + + def _flow_state(self, now: float) -> str: + previous = self.flow_sample + current = (now, self.requests, self.replay_requests, self.http_bytes, + self.down, self.small_responses, self.large_responses, self.poll_requests) + self.flow_sample = current + poll_gap, self.first_poll_gap = self.first_poll_gap, 0.0 + probes = [trace.summary(now) for trace in self.request_traces.values() if trace.replay][:2] + return ('channel=%d window_ms=%.0f requests=%d replays=%d http_down=%d native_down=%d ' + 'small=%d large=%d pending=%d ready=%d upload_wait=%s/%.0fms native_write_ms=%.0f ' + 'capacity_rotations=%d polls=%d first_poll_ms=%.0f probes=[%s]' % ( + self.channel_id, (now - previous[0]) * 1000, + *(current[index] - previous[index] for index in range(1, 7)), + len(self.pending), self.reply_bytes, self.upload_wait_phase, + 0 if self.upload_wait_since is None else (now - self.upload_wait_since) * 1000, + 0 if self.delivering_since is None else (now - self.delivering_since) * 1000, + self.capacity_rotations, + current[7] - previous[7], poll_gap * 1000, + '; '.join(probes))) + + def _fail(self, error: Exception) -> None: + if self.closed: + return + self.closed = True + self.recovery_wakeup.set() + while not self.queue.empty(): + self.queue.get_nowait() + self.reply_bytes = 0 + self.queue.put_nowait(error) + + async def send(self, body: bytes, quick: bool, replay: bool = False) -> bool: + if replay and self.send_lock.locked(): + return False + if not replay: + self.upload_wait_since = time.monotonic() + self.upload_wait_phase = 'send-lock' + self.upload_wait_bytes = len(body) + try: + async with self.send_lock: + return await self._send(body, quick, replay) + finally: + if not replay: + self.upload_wait_since = None + self.upload_wait_phase = '-' + self.upload_wait_bytes = 0 + + async def _send(self, body: bytes, quick: bool, replay: bool) -> bool: + async with self.capacity: + if replay and not self._has_capacity(len(body)): + return False + if not replay: + self.upload_wait_phase = 'channel-capacity' + await asyncio.wait_for(self.capacity.wait_for(lambda: self.closed or self._has_capacity(len(body))), + REQUEST_TIMEOUT) + if self.closed: + raise ConnectionError('H2 channel closed') + if not replay: + self.upload_wait_phase = 'lane-capacity' + if not await self.lane._reserve(len(body), wait=not replay): + return False + if self.closed: + await self.lane._release(len(body)) + raise ConnectionError('H2 channel closed') + replay_packet = self._find_packet(body) if replay else None + now = time.monotonic() + if replay and (replay_packet is None or replay_packet.pending is not None + or not self._replay_due(replay_packet, now)): + await self.lane._release(len(body)) + return False + self.pending_bytes += len(body) + self.requests += 1 + self.replay_requests += int(replay) + self.up += len(body) + stats.bytes_up += len(body) + if not replay: + self.last_native_send = now + native_packet = self._remember_packet(body, now) if not replay else None + if quick: + self.quick_requests += 1 + packet = replay_packet or native_packet + trace = (_RequestTrace(packet.sequence if packet is not None else self.requests, + now, replay, len(body)) + if log.isEnabledFor(logging.DEBUG) else None) + task = asyncio.create_task(self._run_request(body, replay, packet, trace)) + self.pending[task] = len(body) + self.pending_since[task] = now + if trace is not None: + self.request_traces[task] = trace + if native_packet is not None: + native_packet.originals.add(task) + self._trim_history(now) + if replay: + if not self._overdue_original(replay_packet, now): + self.poll_requests += 1 + if not replay_packet.replayed: + self.first_poll_gap = max(self.first_poll_gap, now - self.last_progress) + self.replay_pending.add(task) + replay_packet.attempts += 1 + replay_packet.last_attempt_at = now + replay_packet.pending = task + return True + + def _find_packet(self, body: bytes) -> Optional[_ReplayPacket]: + return next((entry for entry in self.replay_history if entry.body == body), None) + + def _remember_packet(self, body: bytes, now: float) -> Optional[_ReplayPacket]: + key = body[:8] + if key == b'\x00' * 8: + return None + if self.replay_key != key: + self.replay_history.clear() + self.replay_key = key + if len(body) > REPLAY_HISTORY_BYTES: + return None + packet = self._find_packet(body) + if packet is None: + packet = _ReplayPacket(now, body, self.requests) + self.replay_history.append(packet) + return packet + + def _trim_history(self, now: float) -> None: + retained_bytes = sum(len(entry.body) for entry in self.replay_history) + while len(self.replay_history) > REPLAY_HISTORY_PACKETS or retained_bytes > REPLAY_HISTORY_BYTES: + victim = next((entry for entry in self.replay_history if not entry.originals + and not (entry.retired and entry.attempts < REPLAY_MAX_ATTEMPTS + and now - entry.sent_at <= REPLAY_MAX_AGE)), + self.replay_history[0]) + self.replay_history.remove(victim) + retained_bytes -= len(victim.body) + + async def _run_request(self, body: bytes, replay: bool, packet: Optional[_ReplayPacket], + trace: Optional[_RequestTrace]) -> None: + task = asyncio.current_task() + try: + reply = await self.lane._post(body, self.channel_id, replay=replay, channel=self) + self.last_progress = time.monotonic() + if reply and not self.closed: + if self.queue.full() or self.reply_bytes + len(reply) > MAX_CHANNEL_BYTES: + raise BufferError('H2 client response queue exceeded bound') + self.reply_bytes += len(reply) + self.queue.put_nowait(reply) + if trace is not None: + trace.phase = 'complete' + except asyncio.CancelledError: + if trace is not None: + trace.phase = 'cancelled-' + trace.phase + raise + 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; ' + 'closing native channel, reconnect may select another domain', + self.label, self.lane.lane_id, self.channel_id, type(exc).__name__, exc) + else: + log.exception('[%s] H2 lane=%d channel=%d unexpected request failure', + self.label, self.lane.lane_id, self.channel_id) + self._fail(exc) + finally: + await self._finish_request(task, packet) + + async def _finish_request(self, task: asyncio.Task, packet: Optional[_ReplayPacket] = None) -> None: + length = self.pending.pop(task, None) + if length is None: + return + trace = self.request_traces.pop(task, None) + if trace is not None: + trace.finished_at = time.monotonic() + self.recent_traces.append(trace) + self.pending_bytes -= length + self.pending_since.pop(task, None) + self.sent_since.pop(task, None) + self.replay_pending.discard(task) + if packet is None: + packet = next((entry for entry in self.replay_history + if task in entry.originals or entry.pending is task), None) + if packet is not None: + packet.originals.discard(task) + if packet.pending is task: + packet.pending = None + await self.lane._release(length) + async with self.capacity: + self.capacity.notify_all() + self._wake_receiver() + + async def _recover(self) -> None: + while not self.closed: + try: + await asyncio.wait_for(self.recovery_wakeup.wait(), REPLAY_CHECK_SECONDS) + except asyncio.TimeoutError: + pass + self.recovery_wakeup.clear() + now = time.monotonic() + self.lane.log_wait(now) + await self._make_upload_room(now) + if self.closed or self.native_backpressure: + continue + def priority(entry): + overdue = self._overdue_original(entry, now) + return (not overdue, entry.sequence if overdue else -entry.sequence) + + candidates = sorted(self.replay_history, key=priority) + for entry in candidates: + if not self._replay_due(entry, now): + continue + replace = entry.pending + if replace is None and len(self.replay_pending) >= MAX_CHANNEL_REPLAYS: + replace = next((task for task in self.pending if task in self.replay_pending + and task not in self.receiving_tasks), None) + if replace is None: + continue + if replace is not None: + sent = self.sent_since.get(replace) + if (replace in self.receiving_tasks + or sent is None or now - sent < REPLAY_SLOT_SECONDS): + continue + replace.cancel() + await asyncio.gather(replace, return_exceptions=True) + self.replay_rotations += 1 + await self.send(entry.body, False, replay=True) + break + + async def receive(self) -> bytes: + value = await self.queue.get() + if isinstance(value, Exception): + raise value + self.reply_bytes -= len(value) + self._wake_receiver() + return value + + async def close(self) -> None: + async with self.close_lock: + if self.close_done: + return + waiting = self._wait_state(time.monotonic()) + if log.isEnabledFor(logging.DEBUG) and self.request_traces: + now = time.monotonic() + traces = sorted(self.request_traces.values(), key=lambda trace: trace.started)[:3] + log.debug('H2 DETAIL CLOSE lane=%d [%s]', self.lane.lane_id, + '; '.join(self._request_detail(trace, now) for trace in traces)) + self._fail(ConnectionError('H2 channel closed')) + self.closed = True + tasks = list(self.pending) + for task in tasks: + task.cancel() + await asyncio.gather(*tasks, return_exceptions=True) + for task in list(self.pending): + await self._finish_request(task) + self.replay_pending.clear() + self.replay_history.clear() + self.sent_since.clear() + self.request_traces.clear() + self.recent_traces.clear() + self.lane.channels.discard(self) + self.close_done = True + log.debug('[%s] H2 channel=%d closed lane=%d requests=%d replays=%d rotations=%d capacity_rotations=%d up=%d down=%d ' + 'quickack=%d duration=%.1fs remaining_channels=%d; %s', + self.label, self.channel_id, self.lane.lane_id, self.requests, + self.replay_requests, self.replay_rotations, self.capacity_rotations, self.up, self.down, self.quick_requests, + time.monotonic() - self.started, len(self.lane.channels), waiting) + +class _HttpLane: + def __init__(self, host: str, lane_id: int): + self.host, self.lane_id = host, lane_id + self.channels: Set[_HttpChannel] = set() + self.failed_until = 0.0 + self.closed = False + self.tcp_connections = self.requests = self.responses = self.errors = 0 + self.slow_headers = self.slow_bodies = self.replays = 0 + self.max_headers_seconds = self.max_body_seconds = 0.0 + self.next_wait_log = self.next_slow_log = 0.0 + self.next_capacity_log = 0.0 + self.capacity_rotations = 0 + self.next_large_slow_log = 0.0 + self.inflight = self.queued_bytes = self.max_inflight = 0 + self.reply_buffer_bytes = 0 + self.capacity = asyncio.Condition() + self.client = httpx.AsyncClient( + transport=H2Transport(create_ssl_context(), max_streams=MAX_LANE_REQUESTS), + trust_env=False, + timeout=httpx.Timeout(REQUEST_TIMEOUT, connect=8, pool=REQUEST_TIMEOUT), + ) + + async def _trace(self, event: str, info: Dict) -> None: + if event == 'connection.connect_tcp.started': + self.tcp_connections += 1 + stats.h2_tcp_connections += 1 + log.debug('H2 lane=%d TCP_OPEN connection=%d host=%s channels=%d', + self.lane_id, self.tcp_connections, self.host, len(self.channels)) + + async def _preflight(self) -> None: + started = time.monotonic() + async with self.client.stream('HEAD', 'https://' + self.host + '/api', + 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 + 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', + self.lane_id, self.host, response.status_code, + (time.monotonic() - started) * 1000) + + def _has_capacity(self, length: int) -> bool: + return self.inflight < MAX_LANE_REQUESTS and self.queued_bytes + length <= MAX_LANE_BYTES + + async def _reserve(self, length: int, wait: bool = True) -> bool: + async with self.capacity: + if not wait and not self._has_capacity(length): + return False + if wait: + await asyncio.wait_for(self.capacity.wait_for(lambda: self.closed or self._has_capacity(length)), + REQUEST_TIMEOUT) + if self.closed: + raise ConnectionError('H2 lane closed') + self.inflight += 1 + self.queued_bytes += length + self.max_inflight = max(self.max_inflight, self.inflight) + return True + + async def _release(self, length: int) -> None: + async with self.capacity: + self.inflight -= 1 + self.queued_bytes -= length + self.capacity.notify_all() + + async def _post(self, body: bytes, channel_id: int, replay: bool = False, + channel: Optional[_HttpChannel] = None) -> bytes: + self.requests += 1 + self.replays += int(replay) + stats.h2_requests += 1 + request_id = self.requests + started = time.monotonic() + buffered = 0 + receiving = False + request_task = asyncio.current_task() + trace = channel.request_traces.get(request_task) if channel is not None else None + if trace is not None: + trace.request = request_id + trace.phase = 'connect-or-slot' + + async def request_trace(event, info): + await self._trace(event, info) + if event == 'http2.send_request_body.complete' and channel is not None: + channel.sent_since[request_task] = time.monotonic() + if trace is not None: + if event == 'connection.connect_tcp.started': + trace.phase = 'connect' + elif event == 'http2.wait_for_stream.started': + trace.phase = 'stream-slot' + elif event == 'http2.send_request_headers.started': + trace.stream = info['stream_id'] + trace.phase = 'upload' + elif event == 'http2.send_request_body.complete': + trace.sent_at = time.monotonic() + trace.phase = 'headers' if trace.headers_at is None else 'body' + + def response_started(): + nonlocal receiving + if trace is not None and trace.headers_at is None: + trace.headers_at = time.monotonic() + trace.phase = 'body' + if channel is not None and not receiving: + channel.receiving_tasks.add(request_task) + channel.last_progress = time.monotonic() + receiving = True + + try: + async with self.client.stream( + 'POST', 'https://' + self.host + '/api', content=body, + headers={'content-type': 'application/octet-stream', 'accept-encoding': 'identity'}, + extensions={'trace': request_trace, + 'h2_response_started': response_started}) as response: + headers_at = time.monotonic() + if trace is not None: + trace.cf_ray = ''.join(c for c in response.headers.get('cf-ray', '-')[:64] + if c.isascii() and (c.isalnum() or c == '-')) + self.slow_headers += int(headers_at - started >= SLOW_RESPONSE_SECONDS) + 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 != 200: + raise ConnectionError('CF /api HTTP %d cf_ray=%s' % ( + response.status_code, response.headers.get('cf-ray', '-'))) + if 'text/html' in response.headers.get('content-type', '').lower(): + raise ValueError('CF /api returned HTML instead of MTProto') + response_started() + length_header = response.headers.get('content-length') + declared_length = int(length_header) if length_header is not None else None + if declared_length is not None and not 0 <= declared_length <= MAX_PACKET: + raise BufferError('H2 response Content-Length exceeds native packet bound') + content = bytearray() + async for chunk in response.aiter_bytes(): + if len(content) + len(chunk) > MAX_PACKET: + raise BufferError('H2 response exceeds native packet bound') + if self.reply_buffer_bytes + len(chunk) > MAX_LANE_REPLY_BYTES: + raise _ReplyBufferFull('H2 lane response buffers exceeded bound') + self.reply_buffer_bytes += len(chunk) + buffered += len(chunk) + content.extend(chunk) + if trace is not None: + trace.received += len(chunk) + if channel is not None: + channel.last_progress = time.monotonic() + received = len(content) + if (declared_length is not None + and response.headers.get('content-encoding', 'identity') == 'identity' + and received != declared_length): + raise ValueError('incomplete H2 response body') + 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', + self.lane_id, channel_id, request_id, struct.unpack(' 4096) + finished = time.monotonic() + body_seconds = finished - headers_at + self.max_body_seconds = max(self.max_body_seconds, body_seconds) + self.slow_bodies += int(body_seconds >= SLOW_RESPONSE_SECONDS) + next_slow = self.next_large_slow_log if received > 4096 else self.next_slow_log + if finished - started >= SLOW_RESPONSE_SECONDS and finished >= next_slow: + if received > 4096: + self.next_large_slow_log = finished + DIAGNOSTIC_INTERVAL + else: + self.next_slow_log = finished + DIAGNOSTIC_INTERVAL + log.debug('H2 SLOW lane=%d channel=%d request=%d bytes=%d replay=%d ' + 'headers_ms=%.0f body_ms=%.0f detail=[%s]', + self.lane_id, channel_id, request_id, received, replay, + (headers_at - started) * 1000, body_seconds * 1000, + trace.summary(finished) if trace is not None else '-') + return bytes(content) + except _ReplyBufferFull: + raise + except (httpx.HTTPError, OSError, ValueError, BufferError): + self.errors += 1 + self.failed_until = time.monotonic() + DOMAIN_COOLDOWN + raise + finally: + self.reply_buffer_bytes -= buffered + if receiving: + channel.receiving_tasks.discard(request_task) + + def log_wait(self, now: float) -> None: + if not log.isEnabledFor(logging.DEBUG) or now < self.next_wait_log: + return + self.next_wait_log = now + 1.0 + waiting = [channel for channel in self.channels + if ((channel.delivering_since is not None + and now - channel.delivering_since >= SLOW_RESPONSE_SECONDS) or any( + task not in channel.replay_pending + and task not in channel.receiving_tasks + and now - started >= SLOW_RESPONSE_SECONDS + for task, started in channel.pending_since.items())) + and now - channel.last_native_send <= REPLAY_MAX_AGE] + if waiting: + self.next_wait_log = now + DIAGNOSTIC_INTERVAL + waiting.sort(key=lambda channel: channel.channel_id) + log.debug('H2 WAIT lane=%d affected=%d inflight=%d body_buffer=%d [%s]', + self.lane_id, len(waiting), self.inflight, self.reply_buffer_bytes, + '; '.join(channel._wait_state(now) for channel in waiting[:8])) + details = sorted(((channel, trace) for channel in waiting for trace in channel.request_traces.values() + if not trace.replay and now - trace.started >= SLOW_RESPONSE_SECONDS), + key=lambda pair: pair[1].started)[:12] + if details: + log.debug('H2 DETAIL lane=%d [%s]', self.lane_id, + '; '.join(channel._request_detail(trace, now) for channel, trace in details)) + + def log_flow(self, now: float) -> None: + if log.isEnabledFor(logging.DEBUG) and self.channels: + channels = sorted(self.channels, key=lambda channel: channel.channel_id) + for offset in range(0, len(channels), 8): + log.debug('H2 FLOW lane=%d [%s]', self.lane_id, + '; '.join(channel._flow_state(now) for channel in channels[offset:offset + 8])) + + def summary(self) -> str: + return ('lane=%d host=%s channels=%d tcp_connections=%d requests=%d replies=%d ' + 'errors=%d inflight=%d max_inflight=%d queued_bytes=%d reply_buffer_bytes=%d ' + 'replays=%d capacity_rotations=%d slow_headers=%d slow_bodies=%d max_headers_ms=%.0f max_body_ms=%.0f' % ( + self.lane_id, self.host, len(self.channels), self.tcp_connections, + self.requests, self.responses, self.errors, self.inflight, + self.max_inflight, self.queued_bytes, self.reply_buffer_bytes, + self.replays, self.capacity_rotations, self.slow_headers, self.slow_bodies, + self.max_headers_seconds * 1000, self.max_body_seconds * 1000)) + + async def close(self) -> None: + self.closed = True + async with self.capacity: + self.capacity.notify_all() + await asyncio.gather(*(channel.close() for channel in list(self.channels))) + await self.client.aclose() + log.debug('H2 CLOSED %s', self.summary()) + + +class CfH2Pool: + def __init__(self): + for name in ('httpx', 'httpcore', 'hpack', 'h2'): + logging.getLogger(name).setLevel(logging.WARNING) + self.lanes: Dict[str, _HttpLane] = {} + self.locks: Dict[str, asyncio.Lock] = {} + self.failed_until: Dict[str, float] = {} + self.setup_retry_after: Dict[int, float] = {} + self.next_lane = self.next_channel = 0 + self.closed = False + + async def open(self, dc: int, label: str) -> Optional[_HttpChannel]: + if self.closed: + raise ConnectionError('H2 pool closed') + if time.monotonic() < self.setup_retry_after.get(dc, 0): + return None + try: + return await asyncio.wait_for(self._open(dc, label), SETUP_TIMEOUT) + except asyncio.TimeoutError: + self.setup_retry_after[dc] = time.monotonic() + DOMAIN_COOLDOWN + stats.h2_errors += 1 + log.warning('[%s] H2 DC%d setup timed out; using WS for %.0fs', + label, dc, DOMAIN_COOLDOWN) + return None + + async def _open(self, dc: int, label: str) -> Optional[_HttpChannel]: + if self.closed: + raise ConnectionError('H2 pool closed') + if dc not in DC_DEFAULT_IPS: + log.warning('[%s] H2 unsupported DC%d -> existing WS route', label, dc) + return None + for base_domain in balancer.get_domains_for_dc(dc): + host = 'kws%d.%s' % (dc, base_domain) + async with self.locks.setdefault(host, asyncio.Lock()): + if self.closed: + raise ConnectionError('H2 pool closed') + lane = self.lanes.get(host) + if time.monotonic() < max(self.failed_until.get(host, 0), + lane.failed_until if lane else 0): + continue + if lane is None: + self.next_lane += 1 + lane = _HttpLane(host, self.next_lane) + try: + await lane._preflight() + except (httpx.HTTPError, OSError, ValueError) as exc: + stats.h2_errors += 1 + self.failed_until[host] = time.monotonic() + DOMAIN_COOLDOWN + log.warning('[%s] H2 lane=%d preflight failed host=%s: %s: %s', + label, lane.lane_id, host, type(exc).__name__, exc) + await lane.close() + continue + except asyncio.CancelledError: + await lane.close() + raise + if self.closed: + await lane.close() + raise ConnectionError('H2 pool closed during preflight') + self.lanes[host] = lane + self.next_channel += 1 + channel = _HttpChannel(lane, self.next_channel, label) + balancer.update_domain_for_dc(dc, base_domain) + log.info('[%s] ROUTE DC%d media -> H2 lane=%d channel=%d host=%s ' + 'channels=%d tcp_connections=%d', label, dc, lane.lane_id, + channel.channel_id, host, len(lane.channels), lane.tcp_connections) + return channel + log.debug('[%s] H2 DC%d unavailable; using WS', label, dc) + return None + + def log_stats(self) -> None: + for lane in self.lanes.values(): + if lane.channels: + log.debug('H2 STATS %s', lane.summary()) + + def log_flow(self, now: float) -> None: + for lane in self.lanes.values(): + lane.log_flow(now) + + async def close(self) -> None: + self.closed = True + await asyncio.gather(*(lane.close() for lane in self.lanes.values())) + self.lanes.clear() + + +async def bridge_h2(reader, writer, channel: _HttpChannel, ctx, tag: bytes) -> None: + async def write_native(data): + try: + writer.write(ctx.clt_enc.update(data)) + await writer.drain() + except OSError as exc: + raise _NativeDeliveryError(str(exc)) from exc + + async def upload() -> None: + while True: + try: + body, quick = await _read_packet(reader, ctx.clt_dec, tag) + except OSError as exc: + raise _NativeDeliveryError(str(exc)) from exc + await channel.send(body, quick) + + async def download() -> None: + while True: + body = await channel.receive() + if channel.closed: + raise _NativeDeliveryError('H2 channel closed before native response') + channel.delivering_since = time.monotonic() + try: + await write_native(_encode_reply(body, tag)) + finally: + channel.delivering_since = None + channel.down += len(body) + stats.bytes_down += len(body) + channel.last_progress = time.monotonic() + channel._wake_receiver() + + tasks = [asyncio.create_task(upload()), asyncio.create_task(download()), + asyncio.create_task(channel._recover())] + try: + done, _ = await asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED) + 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: + 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() diff --git a/proxy/config.py b/proxy/config.py index f4834c5..b4a8405 100644 --- a/proxy/config.py +++ b/proxy/config.py @@ -70,11 +70,17 @@ class ProxyConfig: fallback_cfproxy: bool = True cfproxy_user_domains: List[str] = field(default_factory=list) cfproxy_worker_domains: List[str] = field(default_factory=list) + cfproxy_h2_media: bool = True disable_secure: bool = False fake_tls_domain: str = '' proxy_protocol: bool = False force_test_dc: bool = False + @property + def h2_enabled(self) -> bool: + return (self.cfproxy_h2_media and self.fallback_cfproxy + and not self.disable_secure and not self.force_test_dc) + proxy_config = ProxyConfig() diff --git a/proxy/h2_transport.py b/proxy/h2_transport.py new file mode 100644 index 0000000..17bb571 --- /dev/null +++ b/proxy/h2_transport.py @@ -0,0 +1,332 @@ +from __future__ import annotations + +import asyncio +import time +from collections import deque + +import httpx +from h2.config import H2Configuration +from h2.connection import H2Connection +from h2.errors import ErrorCodes +from h2.events import ( + ConnectionTerminated, DataReceived, RemoteSettingsChanged, + ResponseReceived, StreamEnded, StreamReset, WindowUpdated, +) +from h2.exceptions import H2Error, StreamClosedError +from h2.settings import SettingCodes, Settings + +STREAM_RECEIVE_WINDOW = 256 * 1024 +CONNECTION_RECEIVE_WINDOW = 4 * 1024 * 1024 + + +class _ResponseStream(httpx.AsyncByteStream): + def __init__(self, connection, stream_id, read_timeout): + self.connection = connection + self.stream_id = stream_id + self.read_timeout = read_timeout + self.changed = asyncio.Event() + self.chunks = deque() + self.headers = None + self.response_started = None + self.error = None + self.ended = False + self.closed = False + self.extensions = {'http_version': b'HTTP/2', 'stream_id': stream_id} + + async def wait(self): + try: + await asyncio.wait_for(self.changed.wait(), self.read_timeout) + except asyncio.TimeoutError as exc: + raise httpx.ReadTimeout('Timed out waiting for H2 stream %d' % self.stream_id) from exc + + def check_error(self): + if self.error is not None: + raise self.error + + async def __aiter__(self): + while not self.closed: + self.changed.clear() + self.check_error() + if self.chunks: + data, credit = self.chunks.popleft() + yield data + self.connection.consumed(self, credit) + elif self.ended: + return + else: + await self.wait() + + async def aclose(self): + if not self.closed: + self.closed = True + self.chunks.clear() + self.connection.release(self) + self.changed.set() + + +class _Connection: + def __init__(self, reader, writer, max_streams, write_timeout): + self.reader, self.writer = reader, writer + self.max_streams = max_streams + self.write_timeout = write_timeout + self.h2 = H2Connection(config=H2Configuration(client_side=True)) + self.h2.local_settings = Settings( + client=True, initial_values={SettingCodes.ENABLE_PUSH: 0, + SettingCodes.INITIAL_WINDOW_SIZE: STREAM_RECEIVE_WINDOW}) + self.streams = {} + self.changed = asyncio.Event() + self.write_lock = asyncio.Lock() + self.settings_received = False + self.error = None + self.h2.initiate_connection() + self.h2.increment_flow_control_window(CONNECTION_RECEIVE_WINDOW - 65535) + self.flush() + self.reader_task = asyncio.create_task(self.read_loop()) + + def get_extra_info(self, name): + return self.writer.get_extra_info(name) + + def flush(self): + data = self.h2.data_to_send() + if data: + self.writer.write(data) + + async def drain(self): + try: + async with self.write_lock: + await asyncio.wait_for(self.writer.drain(), self.write_timeout) + except asyncio.TimeoutError as exc: + error = httpx.WriteTimeout('H2 connection write timed out') + self.fail(error) + raise error from exc + except OSError as exc: + error = httpx.WriteError(str(exc)) + self.fail(error) + raise error from exc + + def fail(self, error): + if self.error is None: + self.error = error + for stream in self.streams.values(): + if not stream.ended: + stream.error = error + stream.changed.set() + self.changed.set() + self.writer.close() + + async def read_loop(self): + try: + while self.error is None: + data = await self.reader.read(64 * 1024) + if not data: + raise httpx.ReadError('H2 peer closed the connection') + events = self.h2.receive_data(data) + terminated = any(isinstance(e, ConnectionTerminated) for e in events) + for event in events: + if isinstance(event, RemoteSettingsChanged): + self.settings_received = True + self.changed.set() + elif isinstance(event, WindowUpdated): + self.changed.set() + elif isinstance(event, ConnectionTerminated): + raise httpx.RemoteProtocolError('H2 GOAWAY error=%s last_stream=%s' % ( + event.error_code, event.last_stream_id)) + else: + stream = self.streams.get(getattr(event, 'stream_id', None)) + if isinstance(event, DataReceived): + if event.flow_controlled_length and not terminated: + self.h2.increment_flow_control_window(event.flow_controlled_length) + if stream is not None: + if event.data: + stream.chunks.append((event.data, event.flow_controlled_length)) + elif not terminated: + self.consumed(stream, event.flow_controlled_length) + elif stream is not None: + if isinstance(event, ResponseReceived): + stream.headers = event.headers + if stream.response_started is not None: + stream.response_started() + elif isinstance(event, StreamEnded): + stream.ended = True + elif isinstance(event, StreamReset): + stream.error = httpx.RemoteProtocolError( + 'H2 stream %d reset: %s' % (event.stream_id, event.error_code)) + self.changed.set() + if stream is not None: + stream.changed.set() + self.flush() + await self.drain() + except asyncio.CancelledError: + self.fail(httpx.ReadError('H2 reader stopped')) + raise + except (OSError, H2Error, httpx.HTTPError) as exc: + error = exc if isinstance(exc, httpx.HTTPError) else httpx.RemoteProtocolError(str(exc)) + self.fail(error) + + def check_error(self): + if self.error is not None: + raise self.error + + async def reserve(self, timeouts): + timeout = timeouts.get('pool') + deadline = None if timeout is None else time.monotonic() + timeout + while True: + self.changed.clear() + self.check_error() + if self.settings_received and len(self.streams) < min( + self.max_streams, self.h2.remote_settings.max_concurrent_streams): + stream_id = self.h2.get_next_available_stream_id() + stream = _ResponseStream(self, stream_id, timeouts.get('read')) + self.streams[stream_id] = stream + return stream + remaining = None if deadline is None else max(0, deadline - time.monotonic()) + try: + await asyncio.wait_for(self.changed.wait(), remaining) + except asyncio.TimeoutError as exc: + raise httpx.PoolTimeout('H2 stream capacity exhausted') from exc + + async def send_body(self, request, stream): + async for data in request.stream: + offset = 0 + while offset < len(data): + self.changed.clear() + self.check_error() + stream.check_error() + count = min(len(data) - offset, self.h2.max_outbound_frame_size, + self.h2.local_flow_control_window(stream.stream_id)) + if count <= 0: + try: + await asyncio.wait_for(self.changed.wait(), self.write_timeout) + except asyncio.TimeoutError as exc: + raise httpx.WriteTimeout('H2 upload flow control timed out') from exc + continue + self.h2.send_data(stream.stream_id, data[offset:offset + count]) + offset += count + self.flush() + await self.drain() + self.h2.end_stream(stream.stream_id) + self.flush() + await self.drain() + + async def request(self, request): + timeouts = request.extensions.get('timeout', {}) + trace = request.extensions.get('trace') + if trace: + await trace('http2.wait_for_stream.started', {}) + stream = await self.reserve(timeouts) + stream.response_started = request.extensions.get('h2_response_started') + stream_id = stream.stream_id + try: + headers = [(b':method', request.method.encode('ascii')), + (b':scheme', request.url.raw_scheme), + (b':authority', request.headers['host'].encode('ascii')), + (b':path', request.url.raw_path)] + headers.extend((k.lower(), v) for k, v in request.headers.raw if k.lower() not in ( + b'host', b'connection', b'keep-alive', b'proxy-connection', + b'upgrade', b'transfer-encoding')) + self.h2.send_headers(stream_id, headers) + self.flush() + if trace: + await trace('http2.send_request_headers.started', {'stream_id': stream_id}) + await self.send_body(request, stream) + if trace: + await trace('http2.send_request_body.complete', {'stream_id': stream_id}) + while stream.headers is None: + stream.changed.clear() + stream.check_error() + await stream.wait() + stream.check_error() + status = next(int(v) for k, v in stream.headers if k == b':status') + return httpx.Response(status, headers=[(k, v) for k, v in stream.headers + if not k.startswith(b':')], + stream=stream, extensions=stream.extensions) + except H2Error as exc: + await stream.aclose() + raise httpx.LocalProtocolError(str(exc)) from exc + except BaseException: + await stream.aclose() + raise + + def consumed(self, stream, credit): + if credit and not stream.ended and not stream.closed and self.error is None: + try: + self.h2.increment_flow_control_window(credit, stream.stream_id) + self.flush() + except StreamClosedError: + pass + + def release(self, stream): + self.streams.pop(stream.stream_id, None) + if self.error is None: + try: + self.h2.reset_stream(stream.stream_id, ErrorCodes.CANCEL) + self.flush() + except StreamClosedError: + pass + except OSError as exc: + self.fail(httpx.WriteError(str(exc))) + self.changed.set() + + async def aclose(self): + self.fail(httpx.ReadError('H2 connection closed')) + self.reader_task.cancel() + await asyncio.gather(self.reader_task, return_exceptions=True) + try: + await asyncio.wait_for(self.writer.wait_closed(), 2) + except (OSError, asyncio.TimeoutError): + pass + + +class H2Transport(httpx.AsyncBaseTransport): + def __init__(self, ssl_context, max_streams=64): + self.ssl_context = ssl_context + self.ssl_context.set_alpn_protocols(['h2']) + self.max_streams = max_streams + self.connection = None + self.origin = None + self.connect_lock = asyncio.Lock() + self.closed = False + + async def handle_async_request(self, request): + origin = (request.url.host, request.url.port or 443) + if request.url.scheme != 'https' or self.origin not in (None, origin): + raise httpx.UnsupportedProtocol('CF H2 transport requires a single HTTPS origin') + async with self.connect_lock: + if self.closed: + raise httpx.ConnectError('H2 transport closed') + self.origin = origin + if self.connection is None or self.connection.error is not None: + if self.connection is not None: + await self.connection.aclose() + timeouts = request.extensions.get('timeout', {}) + trace = request.extensions.get('trace') + if trace: + await trace('connection.connect_tcp.started', {}) + writer = None + try: + reader, writer = await asyncio.wait_for(asyncio.open_connection( + *origin, ssl=self.ssl_context, server_hostname=origin[0]), + timeouts.get('connect')) + if writer.get_extra_info('ssl_object').selected_alpn_protocol() != 'h2': + raise httpx.RemoteProtocolError('CF origin did not negotiate h2') + self.connection = _Connection(reader, writer, self.max_streams, timeouts.get('write')) + if trace: + await trace('connection.start_tls.complete', {'return_value': self.connection}) + except BaseException as exc: + if self.connection is not None: + self.connection.fail(httpx.ConnectError('H2 connection setup interrupted')) + if writer is not None: + writer.close() + if isinstance(exc, asyncio.TimeoutError): + raise httpx.ConnectTimeout('H2 TLS connect timed out') from exc + if isinstance(exc, OSError): + raise httpx.ConnectError(str(exc)) from exc + raise + connection = self.connection + return await connection.request(request) + + async def aclose(self): + self.closed = True + async with self.connect_lock: + if self.connection is not None: + await self.connection.aclose() diff --git a/proxy/network_debug.py b/proxy/network_debug.py new file mode 100644 index 0000000..908dd63 --- /dev/null +++ b/proxy/network_debug.py @@ -0,0 +1,88 @@ +import logging +import time +from collections import defaultdict + +log = logging.getLogger('tg-mtproto-proxy') + + +class WsActivity: + active = set() + + def __init__(self, label, dc): + self.label, self.dc = label, dc + self.started = time.monotonic() + self.native_up = self.ws_up = self.ws_down = self.native_down = 0 + self.first_input = self.first_reply = self.last_rx = self.last_tx = None + self.awaiting_rx = self.sending = self.writing_native = None + self.max_rx_gap = self.max_send = self.max_native_write = 0.0 + self.sample = (self.started, 0, 0, 0, 0) + self.active.add(self) + log.debug('[%s] WS OPEN %s', self.label, self.dc) + + def input(self, length): + self.native_up += length + if self.first_input is None: + self.first_input = time.monotonic() + + def send_started(self): + self.sending = time.monotonic() + + def sent(self, length): + now = time.monotonic() + self.ws_up += length + if self.sending is not None: + self.max_send = max(self.max_send, now - self.sending) + if self.awaiting_rx is None and (self.last_rx is None or self.last_rx < self.sending): + self.awaiting_rx = self.sending + self.last_tx, self.sending = now, None + + def received(self, length): + now = time.monotonic() + self.ws_down += length + self.last_rx = now + if self.awaiting_rx is not None: + self.max_rx_gap = max(self.max_rx_gap, now - self.awaiting_rx) + if self.first_reply is None and self.first_input is not None: + self.first_reply = now - self.first_input + self.awaiting_rx = None + self.writing_native = now + + def delivered(self, length): + self.native_down += length + if self.writing_native is not None: + self.max_native_write = max(self.max_native_write, time.monotonic() - self.writing_native) + self.writing_native = None + + def summary(self, now): + current = (now, self.native_up, self.ws_up, self.ws_down, self.native_down) + previous, self.sample = self.sample, current + return ('client=%s window_ms=%.0f native_up=%d ws_up=%d ws_down=%d native_down=%d ' + 'rx_gap_ms=%.0f last_rx_ms=%.0f ws_write_ms=%.0f native_write_ms=%.0f' % ( + self.label, (now - previous[0]) * 1000, + *(current[index] - previous[index] for index in range(1, 5)), + 0 if self.awaiting_rx is None else (now - self.awaiting_rx) * 1000, + -1 if self.last_rx is None else (now - self.last_rx) * 1000, + 0 if self.sending is None else (now - self.sending) * 1000, + 0 if self.writing_native is None else (now - self.writing_native) * 1000)) + + def close(self): + now = time.monotonic() + self.active.discard(self) + log.debug('WS END %s %s first_reply_ms=%.0f max_rx_gap_ms=%.0f ' + 'max_ws_write_ms=%.0f max_native_write_ms=%.0f', + self.dc, self.summary(now), + -1 if self.first_reply is None else self.first_reply * 1000, + self.max_rx_gap * 1000, self.max_send * 1000, self.max_native_write * 1000) + + +def log_ws_flow(now): + if not log.isEnabledFor(logging.DEBUG): + return + groups = defaultdict(list) + for activity in WsActivity.active: + groups[activity.dc].append(activity) + for dc, activities in sorted(groups.items()): + activities.sort(key=lambda item: item.label) + for offset in range(0, len(activities), 8): + log.debug('WS FLOW %s [%s]', dc, + '; '.join(item.summary(now) for item in activities[offset:offset + 8])) diff --git a/proxy/stats.py b/proxy/stats.py index b61390c..32cbbd7 100644 --- a/proxy/stats.py +++ b/proxy/stats.py @@ -5,6 +5,10 @@ class _Stats: self.connections_total = 0 self.connections_active = 0 self.connections_ws = 0 + self.connections_h2 = 0 + self.h2_tcp_connections = 0 + self.h2_requests = 0 + self.h2_errors = 0 self.connections_tcp_fallback = 0 self.connections_cfproxy = 0 self.connections_fronting = 0 @@ -28,6 +32,10 @@ class _Stats: return (f"total={self.connections_total} " f"active={self.connections_active} " f"ws={self.connections_ws} " + f"h2={self.connections_h2} " + f"h2_tcp={self.h2_tcp_connections} " + f"h2_req={self.h2_requests} " + f"h2_err={self.h2_errors} " f"tcp_fb={self.connections_tcp_fallback} " f"cf={self.connections_cfproxy} " f"front={self.connections_fronting} " @@ -40,4 +48,4 @@ class _Stats: f"down={human_bytes(self.bytes_down)}") -stats = _Stats() \ No newline at end of file +stats = _Stats() diff --git a/proxy/tg_ws_proxy.py b/proxy/tg_ws_proxy.py index 5c861b8..d526d11 100644 --- a/proxy/tg_ws_proxy.py +++ b/proxy/tg_ws_proxy.py @@ -8,7 +8,6 @@ import asyncio import hashlib import argparse import logging -import logging.handlers import socket as _socket from typing import Dict, Optional, Set, Tuple @@ -28,6 +27,8 @@ from .raw_websocket import RawWebSocket, WsHandshakeError, set_sock_opts from .fake_tls import proxy_to_masking_domain, verify_client_hello, build_server_hello, FakeTlsStream, TLS_RECORD_HANDSHAKE from .balancer import balancer from .pool import ws_pool, cf_worker_pool +from .cf_h2 import CfH2Pool +from .network_debug import log_ws_flow from ._aes import Cipher, algorithms, modes @@ -41,6 +42,7 @@ LISTENER_RESTART_DELAY = 1.0 ws_blacklist: Set[str] = set() dc_fail_until: Dict[str, float] = {} ip_fail_until: Dict[str, float] = {} +cf_h2_pool: Optional[CfH2Pool] = None def _try_handshake(handshake: bytes, secret: bytes) -> Optional[Tuple[int, bool, bytes, bytes]]: @@ -338,7 +340,7 @@ async def _handle_client(reader, writer, secret: bytes): ok = await do_fallback( clt_reader, clt_writer, relay_init, label, dc, is_test_dc, is_media, media_tag, - ctx, splitter=splitter) + ctx, splitter=splitter, h2_pool=cf_h2_pool, proto_tag=proto_tag) if not ok: log.warning("[%s] DC%d%s no fallback available", label, dc, media_tag) @@ -417,7 +419,7 @@ async def _handle_client(reader, writer, secret: bytes): ok = await do_fallback( clt_reader, clt_writer, relay_init, label, dc, is_test_dc, is_media, media_tag, - ctx, splitter=splitter_fb) + ctx, splitter=splitter_fb, h2_pool=cf_h2_pool, proto_tag=proto_tag) if ok: log.info("[%s] DC%d%s fallback closed", label, dc, media_tag) @@ -471,7 +473,7 @@ _client_tasks: Set[asyncio.Task] = set() async def _run(stop_event: Optional[asyncio.Event] = None): - global _server_instance, _server_stop_event + global _server_instance, _server_stop_event, cf_h2_pool _server_stop_event = stop_event ws_pool.reset() @@ -480,6 +482,7 @@ async def _run(stop_event: Optional[asyncio.Event] = None): dc_fail_until.clear() ip_fail_until.clear() _client_tasks.clear() + cf_h2_pool = CfH2Pool() if proxy_config.h2_enabled else None user_cf_domains = proxy_config.cfproxy_user_domains if user_cf_domains: @@ -488,8 +491,12 @@ async def _run(stop_event: Optional[asyncio.Event] = None): start_cfproxy_domain_refresh() secret_bytes = bytes.fromhex(proxy_config.secret) + stopping = False def client_cb(r, w): + if stopping: + w.close() + return task = asyncio.create_task(_handle_client(r, w, secret_bytes)) _client_tasks.add(task) task.add_done_callback(_client_tasks.discard) @@ -531,6 +538,8 @@ async def _run(stop_event: Optional[asyncio.Event] = None): if proxy_config.cfproxy_worker_domains: log.info(" CF worker: enabled (%s)", ", ".join(proxy_config.cfproxy_worker_domains)) + if cf_h2_pool is not None: + log.info(' CF H2 media: enabled') if proxy_config.disable_secure: log.info(" No secure: enabled (port 80 for CF proxy/worker)") log.info("=" * 60) @@ -542,11 +551,26 @@ async def _run(stop_event: Optional[asyncio.Event] = None): log.info("=" * 60) async def log_stats(): + next_stats = time.monotonic() + 60 + next_tick = time.monotonic() + 5 try: while True: - await asyncio.sleep(60) + await asyncio.sleep(5) + now = time.monotonic() + if log.isEnabledFor(logging.DEBUG): + if now - next_tick >= .25: + log.debug('NET LOOP late_ms=%.0f', (now - next_tick) * 1000) + log_ws_flow(now) + if cf_h2_pool is not None: + cf_h2_pool.log_flow(now) + next_tick = time.monotonic() + 5 + if now < next_stats: + continue + next_stats = now + 60 bl = ', '.join(f'DC{k}' for k in sorted(ws_blacklist)) or 'none' log.info("stats: %s | ws_bl: %s", stats.summary(), bl) + if cf_h2_pool is not None: + cf_h2_pool.log_stats() except asyncio.CancelledError: raise @@ -563,6 +587,7 @@ async def _run(stop_event: Optional[asyncio.Event] = None): except (asyncio.CancelledError, Exception): pass + serve_task = stop_task = watchdog_task = None try: while True: serve_task = asyncio.create_task(server.serve_forever()) @@ -585,19 +610,12 @@ async def _run(stop_event: Optional[asyncio.Event] = None): waiters, return_when=asyncio.FIRST_COMPLETED) if stop_task is not None and stop_task in done: - for task in list(_client_tasks): - task.cancel() - if _client_tasks: - await asyncio.gather( - *_client_tasks, return_exceptions=True) - await _quiet_cancel(watchdog_task) - await _quiet_cancel(serve_task) - server.close() - await server.wait_closed() break await _quiet_cancel(watchdog_task) await _quiet_cancel(serve_task) + if stop_task is not None: + await _quiet_cancel(stop_task) log.warning( "Listening socket died, restarting server") server.close() @@ -622,6 +640,20 @@ async def _run(stop_event: Optional[asyncio.Event] = None): log.warning("Server restored, listening on %s:%d", proxy_config.host, proxy_config.port) finally: + stopping = True + server.close() + clients = list(_client_tasks) + for task in clients: + task.cancel() + await asyncio.gather(*clients, return_exceptions=True) + for task in (serve_task, watchdog_task, stop_task): + if task is not None: + await _quiet_cancel(task) + await server.wait_closed() + if cf_h2_pool is not None: + await cf_h2_pool.close() + cf_h2_pool = None + _server_instance = _server_stop_event = None log_stats_task.cancel() try: await log_stats_task @@ -666,13 +698,17 @@ def main(): help='WS connection pool size per DC (default 4, min 0)') ap.add_argument('--cfproxy-domain', action='append', default=None, metavar='DOMAIN', - help='User defined Cloudflare-proxied domain for WS fallback ' + help='User defined Cloudflare-proxied domain for H2 media and WS fallback ' '(repeatable for multiple domains)') ap.add_argument('--cfproxy-worker-domain', action='append', default=None, metavar='DOMAIN', help='Cloudflare Worker domain for WS fallback ' '(tried before other fallback methods, ' 'repeatable for multiple domains)') + h2_flags = ap.add_mutually_exclusive_group() + h2_flags.add_argument('--no-h2', action='store_true', + help='Disable HTTP/2 multiplexing for CF media (enabled by default)') + h2_flags.add_argument('--cfproxy-h2-media', action='store_true', help=argparse.SUPPRESS) ap.add_argument('--no-cfproxy', action='store_true', help='Disable Cloudflare proxy fallback') ap.add_argument('--no-secure', action='store_true', @@ -728,6 +764,7 @@ def main(): proxy_config.fallback_cfproxy = not args.no_cfproxy proxy_config.cfproxy_user_domains = coerce_domain_list(args.cfproxy_domain) proxy_config.cfproxy_worker_domains = coerce_domain_list(args.cfproxy_worker_domain) + proxy_config.cfproxy_h2_media = not args.no_h2 proxy_config.disable_secure = args.no_secure proxy_config.fake_tls_domain = args.fake_tls_domain.strip() proxy_config.proxy_protocol = args.proxy_protocol @@ -756,6 +793,8 @@ def main(): root.addHandler(fh) logging.getLogger('asyncio').setLevel(logging.WARNING) + logging.getLogger('httpx').setLevel(logging.WARNING) + logging.getLogger('httpcore').setLevel(logging.WARNING) try: asyncio.run(_run()) diff --git a/pyproject.toml b/pyproject.toml index efb4344..0735678 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -37,6 +37,7 @@ classifiers = [ dependencies = [ "pyperclip==1.9.0", "certifi", + "httpx[http2]==0.28.1", "psutil==5.9.8; platform_system == 'Windows' and python_version < '3.9'", "cryptography==41.0.7; platform_system == 'Windows' and python_version < '3.9'", diff --git a/tests/test_bridge.py b/tests/test_bridge.py index 82e9d39..3386450 100644 --- a/tests/test_bridge.py +++ b/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() diff --git a/tests/test_cf_h2.py b/tests/test_cf_h2.py new file mode 100644 index 0000000..dda4a92 --- /dev/null +++ b/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('= 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(' 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) diff --git a/tests/test_h2_routing.py b/tests/test_h2_routing.py new file mode 100644 index 0000000..3cc34c9 --- /dev/null +++ b/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) diff --git a/tests/test_h2_settings.py b/tests/test_h2_settings.py new file mode 100644 index 0000000..00089fa --- /dev/null +++ b/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() diff --git a/tests/test_h2_transport.py b/tests/test_h2_transport.py new file mode 100644 index 0000000..9564100 --- /dev/null +++ b/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(' bool: pc.buffer_size = max(4, cfg.get("buf_kb", DEFAULT_CONFIG["buf_kb"])) * 1024 pc.pool_size = max(0, cfg.get("pool_size", DEFAULT_CONFIG["pool_size"])) pc.fallback_cfproxy = cfg.get("cfproxy", DEFAULT_CONFIG["cfproxy"]) + pc.cfproxy_h2_media = cfg.get("h2", DEFAULT_CONFIG["h2"]) cfproxy_user_domains = coerce_domain_list( cfg.get("cfproxy_user_domain", DEFAULT_CONFIG["cfproxy_user_domain"]) )