fix: fix middle proxy mode without AD_TAG and retry another endpoint when one fails

Dockerfile

@@ -1,7 +1,7 @@
-FROM ubuntu:22.04
+FROM ubuntu:24.04
 
 RUN apt-get update && apt-get install --no-install-recommends -y python3 python3-uvloop python3-cryptography python3-socks libcap2-bin ca-certificates && rm -rf /var/lib/apt/lists/*
-RUN setcap cap_net_bind_service=+ep /usr/bin/python3.10
+RUN setcap cap_net_bind_service=+ep /usr/bin/python3.12
 
 RUN useradd tgproxy -u 10000
 

mtprotoproxy.py

@@ -39,7 +39,7 @@ TG_MIDDLE_PROXIES_V4 = {
     2: [("149.154.161.144", 8888)], -2: [("149.154.161.144", 8888)],
     3: [("149.154.175.100", 8888)], -3: [("149.154.175.100", 8888)],
     4: [("91.108.4.136", 8888)], -4: [("149.154.165.109", 8888)],
-    5: [("91.108.56.181", 8888)], -5: [("91.108.56.181", 8888)]
+    5: [("91.108.56.183", 8888)], -5: [("91.108.56.183", 8888)]
 }
 
 TG_MIDDLE_PROXIES_V6 = {
@@ -47,7 +47,7 @@ TG_MIDDLE_PROXIES_V6 = {
     2: [("2001:67c:04e8:f002::d", 80)], -2: [("2001:67c:04e8:f002::d", 80)],
     3: [("2001:b28:f23d:f003::d", 8888)], -3: [("2001:b28:f23d:f003::d", 8888)],
     4: [("2001:67c:04e8:f004::d", 8888)], -4: [("2001:67c:04e8:f004::d", 8888)],
-    5: [("2001:b28:f23f:f005::d", 8888)], -5: [("2001:67c:04e8:f004::d", 8888)]
+    5: [("2001:b28:f23f:f005::d", 8888)], -5: [("2001:b28:f23f:f005::d", 8888)]
 }
 
 PROXY_SECRET = bytes.fromhex(
@@ -62,7 +62,6 @@ PREKEY_LEN = 32
 KEY_LEN = 32
 IV_LEN = 16
 HANDSHAKE_LEN = 64
-TLS_HANDSHAKE_LEN = 1 + 2 + 2 + 512
 PROTO_TAG_POS = 56
 DC_IDX_POS = 60
 
@@ -140,10 +139,10 @@ def init_config():
     # use middle proxy, necessary to show ad
     conf_dict.setdefault("USE_MIDDLE_PROXY", len(conf_dict["AD_TAG"]) == 16)
 
-    # if IPv6 avaliable, use it by default
-    conf_dict.setdefault("PREFER_IPV6", socket.has_ipv6)
+    # if IPv6 available, use it by default, IPv6 with middle proxies is unstable now
+    conf_dict.setdefault("PREFER_IPV6", socket.has_ipv6 and not conf_dict["USE_MIDDLE_PROXY"])
 
-    # disables tg->client trafic reencryption, faster but less secure
+    # disables tg->client traffic reencryption, faster but less secure
     conf_dict.setdefault("FAST_MODE", True)
 
     # enables some working modes
@@ -265,6 +264,9 @@ def init_config():
     # telegram servers connect timeout in seconds
     conf_dict.setdefault("TG_CONNECT_TIMEOUT", 10)
 
+    # drop connection if no data from telegram server for this many seconds
+    conf_dict.setdefault("TG_READ_TIMEOUT", 60)
+
     # listen address for IPv4
     conf_dict.setdefault("LISTEN_ADDR_IPV4", "0.0.0.0")
 
@@ -484,7 +486,7 @@ myrandom = MyRandom()
 
 
 class TgConnectionPool:
-    MAX_CONNS_IN_POOL = 64
+    MAX_CONNS_IN_POOL = 16
 
     def __init__(self):
         self.pools = {}
@@ -501,6 +503,16 @@ class TgConnectionPool:
                                           timeout=config.TG_CONNECT_TIMEOUT)
         return reader_tgt, writer_tgt
 
+    def is_conn_dead(self, reader, writer):
+        if writer.transport.is_closing():
+            return True
+        raw_reader = reader
+        while hasattr(raw_reader, 'upstream'):
+            raw_reader = raw_reader.upstream
+        if raw_reader.at_eof():
+            return True
+        return False
+
     def register_host_port(self, host, port, init_func):
         if (host, port, init_func) not in self.pools:
             self.pools[(host, port, init_func)] = []
@@ -513,15 +525,16 @@ class TgConnectionPool:
         self.register_host_port(host, port, init_func)
 
         ret = None
-        for task in self.pools[(host, port, init_func)][::]:
+        for task in self.pools[(host, port, init_func)][:]:
             if task.done():
                 if task.exception():
                     self.pools[(host, port, init_func)].remove(task)
                     continue
 
                 reader, writer, *other = task.result()
-                if writer.transport.is_closing():
+                if self.is_conn_dead(reader, writer):
                     self.pools[(host, port, init_func)].remove(task)
+                    writer.transport.abort()
                     continue
 
                 if not ret:
@@ -658,7 +671,7 @@ class CryptoWrappedStreamReader(LayeredStreamReaderBase):
 
             needed_till_full_block = -len(data) % self.block_size
             if needed_till_full_block > 0:
-                data += self.upstream.readexactly(needed_till_full_block)
+                data += await self.upstream.readexactly(needed_till_full_block)
             return self.decryptor.decrypt(data)
 
     async def readexactly(self, n):
@@ -779,7 +792,7 @@ class MTProtoCompactFrameStreamWriter(LayeredStreamWriterBase):
 
     def write(self, data, extra={}):
         SMALL_PKT_BORDER = 0x7f
-        LARGE_PKT_BORGER = 256 ** 3
+        LARGE_PKT_BORDER = 256 ** 3
 
         if len(data) % 4 != 0:
             print_err("BUG: MTProtoFrameStreamWriter attempted to send msg with len", len(data))
@@ -792,7 +805,7 @@ class MTProtoCompactFrameStreamWriter(LayeredStreamWriterBase):
 
         if len_div_four < SMALL_PKT_BORDER:
             return self.upstream.write(bytes([len_div_four]) + data)
-        elif len_div_four < LARGE_PKT_BORGER:
+        elif len_div_four < LARGE_PKT_BORDER:
             return self.upstream.write(b'\x7f' + int.to_bytes(len_div_four, 3, 'little') + data)
         else:
             print_err("Attempted to send too large pkt len =", len(data))
@@ -936,7 +949,7 @@ class ProxyReqStreamWriter(LayeredStreamWriterBase):
             print_err("BUG: attempted to send msg with len %d" % len(msg))
             return 0
 
-        flags = FLAG_HAS_AD_TAG | FLAG_MAGIC | FLAG_EXTMODE2
+        flags = FLAG_MAGIC | FLAG_EXTMODE2
 
         if self.proto_tag == PROTO_TAG_ABRIDGED:
             flags |= FLAG_ABRIDGED
@@ -951,10 +964,15 @@ class ProxyReqStreamWriter(LayeredStreamWriterBase):
         if msg.startswith(b"\x00" * 8):
             flags |= FLAG_NOT_ENCRYPTED
 
+        if len(config.AD_TAG) == 16:
+            flags |= FLAG_HAS_AD_TAG
+
         full_msg = bytearray()
         full_msg += RPC_PROXY_REQ + int.to_bytes(flags, 4, "little") + self.out_conn_id
-        full_msg += self.remote_ip_port + self.our_ip_port + EXTRA_SIZE + PROXY_TAG
-        full_msg += bytes([len(config.AD_TAG)]) + config.AD_TAG + FOUR_BYTES_ALIGNER
+        full_msg += self.remote_ip_port + self.our_ip_port
+        if len(config.AD_TAG) == 16:
+            full_msg += EXTRA_SIZE + PROXY_TAG
+            full_msg += bytes([len(config.AD_TAG)]) + config.AD_TAG + FOUR_BYTES_ALIGNER
         full_msg += msg
 
         return self.upstream.write(full_msg)
@@ -1236,7 +1254,7 @@ async def handle_handshake(reader, writer):
     global last_client_ips
     global last_clients_with_same_handshake
 
-    TLS_START_BYTES = b"\x16\x03\x01\x02\x00\x01\x00\x01\xfc\x03\x03"
+    TLS_START_BYTES = b"\x16\x03\x01"
 
     if writer.transport.is_closing() or writer.get_extra_info("peername") is None:
         return False
@@ -1262,7 +1280,13 @@ async def handle_handshake(reader, writer):
             break
 
     if is_tls_handshake:
-        handshake += await reader.readexactly(TLS_HANDSHAKE_LEN - len(handshake))
+        handshake += await reader.readexactly(2)
+        tls_handshake_len = int.from_bytes(handshake[-2:], "big")
+        if tls_handshake_len < 512:
+            is_tls_handshake = False
+
+    if is_tls_handshake:
+        handshake += await reader.readexactly(tls_handshake_len)
         tls_handshake_result = await handle_fake_tls_handshake(handshake, reader, writer, peer)
 
         if not tls_handshake_result:
@@ -1359,7 +1383,7 @@ async def do_direct_handshake(proto_tag, dc_idx, dec_key_and_iv=None):
         print_err("Got connection refused while trying to connect to", dc, TG_DATACENTER_PORT)
         return False
     except ConnectionAbortedError as E:
-        print_err("The Telegram server connection is bad: %d (%s %s) %s" % (dc_idx, addr, port, E))
+        print_err("The Telegram server connection is bad: %d (%s %s) %s" % (dc_idx, dc, TG_DATACENTER_PORT, E))
         return False
     except (OSError, asyncio.TimeoutError) as E:
         print_err("Unable to connect to", dc, TG_DATACENTER_PORT)
@@ -1540,23 +1564,26 @@ async def do_middleproxy_handshake(proto_tag, dc_idx, cl_ip, cl_port):
     if use_ipv6_tg:
         if dc_idx not in TG_MIDDLE_PROXIES_V6:
             return False
-        addr, port = myrandom.choice(TG_MIDDLE_PROXIES_V6[dc_idx])
+        proxies = list(TG_MIDDLE_PROXIES_V6[dc_idx])
     else:
         if dc_idx not in TG_MIDDLE_PROXIES_V4:
             return False
-        addr, port = myrandom.choice(TG_MIDDLE_PROXIES_V4[dc_idx])
+        proxies = list(TG_MIDDLE_PROXIES_V4[dc_idx])
 
-    try:
-        ret = await tg_connection_pool.get_connection(addr, port, middleproxy_handshake)
-        reader_tgt, writer_tgt, my_ip, my_port = ret
-    except ConnectionRefusedError as E:
-        print_err("The Telegram server %d (%s %s) is refusing connections" % (dc_idx, addr, port))
-        return False
-    except ConnectionAbortedError as E:
-        print_err("The Telegram server connection is bad: %d (%s %s) %s" % (dc_idx, addr, port, E))
-        return False
-    except (OSError, asyncio.TimeoutError) as E:
-        print_err("Unable to connect to the Telegram server %d (%s %s)" % (dc_idx, addr, port))
+    myrandom.shuffle(proxies)
+
+    for addr, port in proxies:
+        try:
+            reader_tgt, writer_tgt, my_ip, my_port = await tg_connection_pool.get_connection(
+                addr, port, middleproxy_handshake)
+            break
+        except ConnectionRefusedError as E:
+            print_err("The Telegram server %d (%s %s) is refusing connections, trying next" % (dc_idx, addr, port))
+        except ConnectionAbortedError as E:
+            print_err("The Telegram server connection is bad: %d (%s %s) %s, trying next" % (dc_idx, addr, port, E))
+        except (OSError, asyncio.TimeoutError) as E:
+            print_err("Unable to connect to the Telegram server %d (%s %s), trying next" % (dc_idx, addr, port))
+    else:
         return False
 
     writer_tgt = ProxyReqStreamWriter(writer_tgt, cl_ip, cl_port, my_ip, my_port, proto_tag)
@@ -1568,7 +1595,11 @@ async def do_middleproxy_handshake(proto_tag, dc_idx, cl_ip, cl_port):
 async def tg_connect_reader_to_writer(rd, wr, user, rd_buf_size, is_upstream):
     try:
         while True:
-            data = await rd.read(rd_buf_size)
+            if not is_upstream:
+                data = await asyncio.wait_for(rd.read(rd_buf_size),
+                                              timeout=config.TG_READ_TIMEOUT)
+            else:
+                data = await rd.read(rd_buf_size)
             if isinstance(data, tuple):
                 data, extra = data
             else:
@@ -1589,7 +1620,7 @@ async def tg_connect_reader_to_writer(rd, wr, user, rd_buf_size, is_upstream):
 
                 wr.write(data, extra)
                 await wr.drain()
-    except (OSError, asyncio.IncompleteReadError) as e:
+    except (OSError, asyncio.IncompleteReadError, asyncio.TimeoutError) as e:
         # print_err(e)
         pass
 
@@ -1995,7 +2026,8 @@ async def get_srv_time():
                     continue
                 line = line[len("Date: "):].decode()
                 srv_time = datetime.datetime.strptime(line, "%a, %d %b %Y %H:%M:%S %Z")
-                now_time = datetime.datetime.utcnow()
+                srv_time = srv_time.replace(tzinfo=datetime.timezone.utc)
+                now_time = datetime.datetime.now(datetime.timezone.utc)
                 is_time_skewed = (now_time-srv_time).total_seconds() > MAX_TIME_SKEW
                 if is_time_skewed and config.USE_MIDDLE_PROXY and not disable_middle_proxy:
                     print_err("Time skew detected, please set the clock")
@@ -2141,15 +2173,15 @@ def print_tg_info():
         for ip in ip_addrs:
             if config.MODES["classic"]:
                 params = {"server": ip, "port": config.PORT, "secret": secret}
-                params_encodeded = urllib.parse.urlencode(params, safe=':')
-                classic_link = "tg://proxy?{}".format(params_encodeded)
+                params_encoded = urllib.parse.urlencode(params, safe=':')
+                classic_link = "tg://proxy?{}".format(params_encoded)
                 proxy_links.append({"user": user, "link": classic_link})
                 print("{}: {}".format(user, classic_link), flush=True)
 
             if config.MODES["secure"]:
                 params = {"server": ip, "port": config.PORT, "secret": "dd" + secret}
-                params_encodeded = urllib.parse.urlencode(params, safe=':')
-                dd_link = "tg://proxy?{}".format(params_encodeded)
+                params_encoded = urllib.parse.urlencode(params, safe=':')
+                dd_link = "tg://proxy?{}".format(params_encoded)
                 proxy_links.append({"user": user, "link": dd_link})
                 print("{}: {}".format(user, dd_link), flush=True)
 
@@ -2159,8 +2191,8 @@ def print_tg_info():
                 # tls_secret = bytes.fromhex("ee" + secret) + config.TLS_DOMAIN.encode()
                 # tls_secret_base64 = base64.urlsafe_b64encode(tls_secret)
                 params = {"server": ip, "port": config.PORT, "secret": tls_secret}
-                params_encodeded = urllib.parse.urlencode(params, safe=':')
-                tls_link = "tg://proxy?{}".format(params_encodeded)
+                params_encoded = urllib.parse.urlencode(params, safe=':')
+                tls_link = "tg://proxy?{}".format(params_encoded)
                 proxy_links.append({"user": user, "link": tls_link})
                 print("{}: {}".format(user, tls_link), flush=True)