Поддержка тестового окружения Telegram (тестовые DC) (#1087)

proxy/tg_ws_proxy.py

@@ -280,6 +280,11 @@ async def _handle_client(reader, writer, secret: bytes):
 
         dc, is_media, proto_tag, client_dec_prekey_iv = result
 
+        is_test_dc = proxy_config.force_test_dc or dc >= 10000
+        if dc >= 10000:
+            log.info("[%s] test DC%d -> DC%d", label, dc, dc - 10000)
+            dc -= 10000
+
         if proto_tag == PROTO_TAG_ABRIDGED:
             proto_int = PROTO_ABRIDGED_INT
         elif proto_tag == PROTO_TAG_INTERMEDIATE:
@@ -295,14 +300,15 @@ async def _handle_client(reader, writer, secret: bytes):
         relay_init = _generate_relay_init(proto_tag, dc_idx)
         ctx = _build_crypto_ctx(client_dec_prekey_iv, secret, relay_init)
 
-        dc_key = f'{dc}{"m" if is_media else ""}'
+        dc_key = f'{dc}{"t" if is_test_dc else ""}{"m" if is_media else ""}'
         media_tag = " media" if is_media else ""
         now = time.monotonic()
+        ws_path = WS_PATH_TEST if is_test_dc else WS_PATH
         target = proxy_config.dc_redirects.get(dc)
         is_any_cf_fallback = proxy_config.fallback_cfproxy or proxy_config.cfproxy_worker_domains
 
         # Fallback if DC not in config, if WS blacklisted for this DC/is_media or if connect to ip is timed out
-        if (dc not in proxy_config.dc_redirects 
+        if (dc not in proxy_config.dc_redirects
             or dc_key in ws_blacklist
             or now < ip_fail_until.get(target, 0) and is_any_cf_fallback):
 
@@ -322,7 +328,7 @@ async def _handle_client(reader, writer, secret: bytes):
                 pass
             ok = await do_fallback(
                 clt_reader, clt_writer, relay_init, label,
-                dc, is_media, media_tag,
+                dc, is_test_dc, is_media, media_tag,
                 ctx, splitter=splitter)
             if not ok:
                 log.warning("[%s] DC%d%s no fallback available",
@@ -338,7 +344,7 @@ async def _handle_client(reader, writer, secret: bytes):
         ws_timed_out = False
         all_redirects = True
 
-        ws = await ws_pool.get(dc, is_media, target, domains)
+        ws = await ws_pool.get(dc, is_media, target, domains, path=ws_path) if not is_test_dc else None
         if ws:
             log.info("[%s] DC%d%s -> pool hit via %s",
                      label, dc, media_tag, target)
@@ -348,7 +354,7 @@ async def _handle_client(reader, writer, secret: bytes):
                      label, dc, media_tag, domains[0])
             try:
                 ws = await RawWebSocket.connect(target, domains[0],
-                                                timeout=5.0,
+                                                timeout=5.0, path=ws_path,
                                                 sni="sprinthost.ru")
             except Exception as exc:
                 stats.ws_errors += 1
@@ -363,12 +369,13 @@ async def _handle_client(reader, writer, secret: bytes):
                 ws_pool.fronting_until = 0.0
         else:
             for domain in domains:
-                url = f'wss://{domain}/apiws'
+                url = f'wss://{domain}{ws_path}'
                 log.info("[%s] DC%d%s -> %s via %s",
                          label, dc, media_tag, url, target)
                 try:
                     ws = await RawWebSocket.connect(target, domain,
-                                                    timeout=ws_timeout)
+                                                    timeout=ws_timeout,
+                                                    path=ws_path)
                     all_redirects = False
                     break
                 except WsHandshakeError as exc:
@@ -404,7 +411,7 @@ async def _handle_client(reader, writer, secret: bytes):
                      label, dc, media_tag, domains[0])
             try:
                 ws = await RawWebSocket.connect(target, domains[0],
-                                                timeout=5.0,
+                                                timeout=5.0, path=ws_path,
                                                 sni="sprinthost.ru")
             except Exception as exc:
                 stats.ws_errors += 1
@@ -440,7 +447,7 @@ async def _handle_client(reader, writer, secret: bytes):
                 pass
             ok = await do_fallback(
                 clt_reader, clt_writer, relay_init, label,
-                dc, is_media, media_tag,
+                dc, is_test_dc, is_media, media_tag,
                 ctx, splitter=splitter_fb)
             if ok:
                 log.info("[%s] DC%d%s fallback closed",
@@ -703,6 +710,11 @@ def main():
                     metavar='DOMAIN',
                     help='Enable Fake TLS (ee-secret) masking with the given '
                          'SNI domain, e.g. example.com')
+    ap.add_argument('--force-test-dc', action='store_true',
+                    help='Force ALL traffic to Telegram TEST datacenters. '
+                         'Not needed for Telegram Desktop (test DCs 10001+ '
+                         'are detected automatically); use for clients that '
+                         'signal test DCs as plain 1-3')
     ap.add_argument('--proxy-protocol', action='store_true',
                     help='Accept PROXY protocol v1 header '
                          '(for use behind nginx/haproxy with proxy_protocol on)')
@@ -742,6 +754,7 @@ def main():
     proxy_config.cfproxy_worker_domains = coerce_domain_list(args.cfproxy_worker_domain)
     proxy_config.fake_tls_domain = args.fake_tls_domain.strip()
     proxy_config.proxy_protocol = args.proxy_protocol
+    proxy_config.force_test_dc = args.force_test_dc
 
     log_level = logging.DEBUG if args.verbose else logging.INFO
     log_fmt = logging.Formatter('%(asctime)s  %(levelname)-5s  %(message)s',