multiplexing media through cfproxy

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,

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('<I', await plaintext(4))
+        quick, length = bool(value & 0x80000000), value & 0x7fffffff
+    if not 24 <= length <= MAX_PACKET + (15 if tag == PROTO_TAG_SECURE else 0):
+        raise ValueError('native packet length outside supported range: %d' % length)
+    body = await plaintext(length)
+    if tag == PROTO_TAG_SECURE:
+        if body[:8] == b'\x00' * 8:
+            packet_length = 20 + int.from_bytes(body[16:20], 'little')
+        else:
+            packet_length = 24 + ((length - 24) // 16) * 16
+        if not 24 <= packet_length <= length or length - packet_length > 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('<I', len(body) + len(padding)) + body + padding
+
+
+class _HttpChannel:
+    def __init__(self, lane, channel_id: int, label: str):
+        self.lane, self.channel_id, self.label = lane, channel_id, label
+        self.closed = False
+        self.pending: Dict[asyncio.Task, int] = {}
+        self.pending_since: Dict[asyncio.Task, float] = {}
+        self.sent_since: Dict[asyncio.Task, float] = {}
+        self.pending_bytes = 0
+        self.reply_bytes = 0
+        self.queue = asyncio.Queue(maxsize=16)
+        self.capacity = asyncio.Condition()
+        self.recovery_wakeup = asyncio.Event()
+        self.requests = self.up = self.down = self.quick_requests = 0
+        self.started = time.monotonic()
+        self.receiving_tasks: Set[asyncio.Task] = set()
+        self.delivering_since: Optional[float] = None
+        self.replay_pending: Set[asyncio.Task] = set()
+        self.replay_history: List[_ReplayPacket] = []
+        self.replay_key: Optional[bytes] = None
+        self.replay_requests = 0
+        self.poll_requests = 0
+        self.first_poll_gap = 0.0
+        self.replay_rotations = 0
+        self.capacity_rotations = 0
+        self.last_progress = self.started
+        self.last_native_send = self.started
+        self.close_lock = asyncio.Lock()
+        self.send_lock = asyncio.Lock()
+        self.close_done = False
+        self.request_traces: Dict[asyncio.Task, _RequestTrace] = {}
+        self.recent_traces = deque(maxlen=16)
+        self.upload_wait_since: Optional[float] = None
+        self.upload_wait_phase = '-'
+        self.upload_wait_bytes = 0
+        self.http_bytes = self.small_responses = self.large_responses = 0
+        self.flow_sample = (self.started, 0, 0, 0, 0, 0, 0, 0)
+        lane.channels.add(self)
+
+    @property
+    def receiving_http(self) -> 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('<i', content)[0])
+                self.responses += 1
+                if channel is not None:
+                    channel.http_bytes += received
+                    channel.small_responses += int(received <= 4096)
+                    channel.large_responses += int(received > 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()

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()
 

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()

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]))

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()
+stats = _Stats()

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())