import asyncio import logging import struct import time from typing import Dict, List, Optional, Set, Tuple from urllib.parse import urlencode from .utils import * from .stats import stats from .balancer import balancer from .config import proxy_config from .raw_websocket import RawWebSocket from .pool import cf_worker_pool from .network_debug import WsActivity from .cf_h2 import bridge_h2 log = logging.getLogger('tg-mtproto-proxy') _st_I_le = struct.Struct(' None: _tcp_failures.clear() _tcp_retry_after.clear() _tcp_connecting.clear() class CryptoCtx: __slots__ = ('clt_dec', 'clt_enc', 'tg_enc', 'tg_dec') def __init__(self, clt_dec, clt_enc, tg_enc, tg_dec): self.clt_dec = clt_dec # decrypt from client self.clt_enc = clt_enc # encrypt to client self.tg_enc = tg_enc # encrypt to telegram self.tg_dec = tg_dec # decrypt from telegram class MsgSplitter: """ Splits TCP stream data into individual MTProto transport packets so each can be sent as a separate WS frame. """ __slots__ = ('_proto', '_cipher_buf', '_plain_buf', '_disabled') def __init__(self, proto_int: int): self._proto = proto_int self._cipher_buf = bytearray() self._plain_buf = bytearray() self._disabled = False def split(self, chunk: bytes, plain: bytes) -> List[bytes]: if not chunk: return [] if self._disabled: return [chunk] self._cipher_buf.extend(chunk) self._plain_buf.extend(plain) parts = [] offset = 0 buf_len = len(self._cipher_buf) while offset < buf_len: packet_len = self._next_packet_len(offset, buf_len - offset) if packet_len is None: break if packet_len <= 0: parts.append(bytes(self._cipher_buf[offset:])) offset = buf_len self._disabled = True break parts.append(bytes(self._cipher_buf[offset:offset + packet_len])) offset += packet_len if offset: del self._cipher_buf[:offset] del self._plain_buf[:offset] return parts def flush(self) -> List[bytes]: if not self._cipher_buf: return [] tail = bytes(self._cipher_buf) self._cipher_buf.clear() self._plain_buf.clear() return [tail] def _next_packet_len(self, offset: int, avail: int) -> Optional[int]: if avail <= 0: return None if self._proto == PROTO_ABRIDGED_INT: return self._next_abridged_len(offset, avail) if self._proto in (PROTO_INTERMEDIATE_INT, PROTO_PADDED_INTERMEDIATE_INT): return self._next_intermediate_len(offset, avail) return 0 def _next_abridged_len(self, offset: int, avail: int) -> Optional[int]: first = self._plain_buf[offset] if first in (0x7F, 0xFF): if avail < 4: return None payload_len = int.from_bytes( self._plain_buf[offset + 1:offset + 4], 'little') * 4 header_len = 4 else: payload_len = (first & 0x7F) * 4 header_len = 1 if payload_len <= 0: return 0 packet_len = header_len + payload_len if avail < packet_len: return None return packet_len def _next_intermediate_len(self, offset: int, avail: int) -> Optional[int]: if avail < 4: return None payload_len = _st_I_le.unpack_from(self._plain_buf, offset)[0] & 0x7FFFFFFF if payload_len <= 0: return 0 packet_len = 4 + payload_len if avail < packet_len: return None return packet_len 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, *, 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 worker_domains = proxy_config.cfproxy_worker_domains methods: List[str] = [] if worker_domains and fallback_dst: methods.append('cf_worker') if use_cf: methods.append('cf') if fallback_dst: methods.append('tcp') for method in methods: if method == 'cf_worker' and fallback_dst: ok = await _cfproxy_worker_fallback( reader, writer, relay_init, label, ctx, dc=dc, is_test_dc=is_test_dc, is_media=is_media, fallback_dst=fallback_dst, splitter=splitter) 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, splitter=splitter) if ok: return True elif method == 'tcp' and fallback_dst: ok = await _tcp_fallback( reader, writer, fallback_dst, 443, relay_init, label, ctx) if ok: return True return False async def _cfproxy_worker_fallback(reader, writer, relay_init, label, ctx: CryptoCtx, dc: int, is_test_dc: bool, is_media: bool, fallback_dst: str, splitter=None): media_tag = ' media' if is_media else '' worker_domains = proxy_config.cfproxy_worker_domains if not worker_domains: return False pooled = None if is_test_dc else await cf_worker_pool.get( dc, fallback_dst, worker_domains) if pooled: ws, worker_domain = pooled log.info("[%s] DC%d%s -> CF worker pool hit via %s for %s", label, dc, media_tag, worker_domain, fallback_dst) else: query = urlencode({ 'dst': fallback_dst, 'dc': str(dc), }) path = f'/apiws?{query}' ws = None for worker_domain in cf_worker_pool.available_domains(worker_domains): log.info("[%s] DC%d%s -> trying CF worker %s for %s", label, dc, media_tag, worker_domain, fallback_dst) try: ws = await RawWebSocket.connect(worker_domain, worker_domain, timeout=10.0, path=path, secure=not proxy_config.disable_secure) break except Exception as exc: cf_worker_pool.report_failure(worker_domain, exc) log.warning("[%s] DC%d%s CF worker %s failed: %s", label, dc, media_tag, worker_domain, repr(exc)) continue if ws is None: return False stats.connections_cfproxy += 1 await ws.send(relay_init) await bridge_ws_reencrypt(reader, writer, ws, label, ctx, dc=dc, is_media=is_media, splitter=None) return True async def _cfproxy_fallback(reader, writer, relay_init, label, ctx: CryptoCtx, dc: int, is_media: bool, splitter=None): media_tag = ' media' if is_media else '' ws = None chosen_domain = None 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}' try: ws = await RawWebSocket.connect(domain, domain, timeout=10.0, secure=not proxy_config.disable_secure) chosen_domain = base_domain break except Exception as exc: log.warning("[%s] DC%d%s CF proxy failed: %s", label, dc, media_tag, repr(exc)) 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) stats.connections_cfproxy += 1 await ws.send(relay_init) await bridge_ws_reencrypt(reader, writer, ws, label, ctx, dc=dc, is_media=is_media, splitter=splitter) return True async def _tcp_fallback(reader, writer, dst, port, relay_init, label, ctx: CryptoCtx): key = (dst, port) if key in _tcp_connecting or time.monotonic() < _tcp_retry_after.get(key, 0): log.debug("[%s] TCP fallback to %s:%d skipped (backoff or connecting)", label, dst, port) return False _tcp_connecting.add(key) rw = None try: try: log.info("[%s] TCP fallback to %s:%d", label, dst, port) rr, rw = await asyncio.wait_for( asyncio.open_connection(dst, port), timeout=10) rw.write(relay_init) await rw.drain() except Exception as exc: failures = _tcp_failures.get(key, 0) + 1 _tcp_failures[key] = failures delay = min(TCP_BACKOFF_INITIAL * 2 ** min(failures - 1, 7), TCP_BACKOFF_MAX) _tcp_retry_after[key] = time.monotonic() + delay log.warning("[%s] TCP fallback to %s:%d failed: %r; retry in %.0fs", label, dst, port, exc, delay) return False finally: _tcp_connecting.discard(key) _tcp_failures.pop(key, None) _tcp_retry_after.pop(key, None) stats.connections_tcp_fallback += 1 await _bridge_tcp_reencrypt(reader, writer, rr, rw, label, ctx) return True finally: if rw is not None: rw.close() try: await rw.wait_closed() except (OSError, ConnectionError): pass async def bridge_ws_reencrypt(reader, writer, ws: RawWebSocket, label, ctx: CryptoCtx, dc=None, is_media=False, splitter: Optional[MsgSplitter] = None): """ Bidirectional TCP(client) <-> WS(telegram) with re-encryption. client ciphertext → decrypt(clt_key) → encrypt(tg_key) → WS WS data → decrypt(tg_key) → encrypt(clt_key) → client TCP """ dc_tag = f"DC{dc}{'m' if is_media else ''}" if dc else "DC?" up_bytes = 0 down_bytes = 0 up_packets = 0 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 try: while True: chunk = await reader.read(65536) if not chunk: 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 plain = ctx.clt_dec.update(chunk) chunk = ctx.tg_enc.update(plain) if splitter: parts = splitter.split(chunk, plain) 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: close_reason = f"client: {type(e).__name__}" except Exception as e: close_reason = f"client: {type(e).__name__}: {e}" log.debug("[%s] tcp->ws ended: %s", label, e) async def ws_to_tcp(): nonlocal down_bytes, down_packets, close_reason try: while True: data = await ws.recv() if data is None: if close_reason == 'normal': 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 plain = ctx.tg_dec.update(data) 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: close_reason = f"upstream: {type(e).__name__}" except asyncio.IncompleteReadError: close_reason = 'upstream: tcp_reset' except Exception as e: close_reason = f"upstream: {type(e).__name__}: {e}" log.debug("[%s] ws->tcp ended: %s", label, e) tasks = [asyncio.create_task(tcp_to_ws()), asyncio.create_task(ws_to_tcp())] try: await asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED) finally: for t in tasks: t.cancel() for t in tasks: try: await t 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, human_bytes(up_bytes), up_packets, human_bytes(down_bytes), down_packets, elapsed) try: await ws.close() except BaseException: pass try: writer.close() await writer.wait_closed() except BaseException: pass async def _bridge_tcp_reencrypt(reader, writer, remote_reader, remote_writer, label, ctx: CryptoCtx): """Bidirectional TCP <-> TCP with re-encryption.""" async def forward(src, dst_w, is_up): try: while True: data = await src.read(65536) if not data: break n = len(data) if is_up: stats.bytes_up += n plain = ctx.clt_dec.update(data) data = ctx.tg_enc.update(plain) else: stats.bytes_down += n plain = ctx.tg_dec.update(data) data = ctx.clt_enc.update(plain) dst_w.write(data) await dst_w.drain() except asyncio.CancelledError: pass except Exception as e: log.debug("[%s] forward ended: %s", label, e) tasks = [ asyncio.create_task(forward(reader, remote_writer, True)), asyncio.create_task(forward(remote_reader, writer, False)), ] try: await asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED) finally: for t in tasks: t.cancel() for t in tasks: try: await t except BaseException: pass for w in (writer, remote_writer): try: w.close() await w.wait_closed() except BaseException: pass