more secure defaul settings

3.8 removes the asyncio.streams re-export of `IncompleteReadError`, so just access it directly from `asyncio`, like everywhere else in the code already does.

Dockerfile

@@ -7,6 +7,8 @@ RUN apk add --no-cache python3 py3-cryptography ca-certificates libcap
 RUN chown -R tgproxy:tgproxy /home/tgproxy
 RUN setcap cap_net_bind_service=+ep /usr/bin/python3.7
 
+COPY mtprotoproxy.py config.py /home/tgproxy/
+
 USER tgproxy
 
 WORKDIR /home/tgproxy/

config.py

@@ -1,9 +1,9 @@
-PORT = 3256
+PORT = 443
 
 # name -> secret (32 hex chars)
 USERS = {
-    "tg":  "00000000000000000000000000000000",
-    "tg2": "0123456789abcdef0123456789abcdef"
+    "tg":  "00000000000000000000000000000001",
+#    "tg2": "0123456789abcdef0123456789abcdef",
 }
 
 # Makes the proxy harder to detect

mtprotoproxy.py

@@ -66,6 +66,8 @@ TLS_HANDSHAKE_LEN = 1 + 2 + 2 + 512
 PROTO_TAG_POS = 56
 DC_IDX_POS = 60
 
+MIN_CERT_LEN = 1024
+
 PROTO_TAG_ABRIDGED = b"\xef\xef\xef\xef"
 PROTO_TAG_INTERMEDIATE = b"\xee\xee\xee\xee"
 PROTO_TAG_SECURE = b"\xdd\xdd\xdd\xdd"
@@ -236,6 +238,9 @@ def init_config():
     # export proxy link to prometheus
     conf_dict.setdefault("METRICS_EXPORT_LINKS", False)
 
+    # default prefix for metrics
+    conf_dict.setdefault("METRICS_PREFIX", "mtprotoproxy_")
+
     # allow access to config by attributes
     config = type("config", (dict,), conf_dict)(conf_dict)
 
@@ -431,6 +436,60 @@ class MyRandom(random.Random):
 myrandom = MyRandom()
 
 
+class TgConnectionPool:
+    MAX_CONNS_IN_POOL = 64
+
+    def __init__(self):
+        self.pools = {}
+
+    async def open_tg_connection(self, host, port, init_func=None):
+        task = asyncio.open_connection(host, port, limit=get_to_clt_bufsize())
+        reader_tgt, writer_tgt = await asyncio.wait_for(task, timeout=config.TG_CONNECT_TIMEOUT)
+
+        set_keepalive(writer_tgt.get_extra_info("socket"))
+        set_bufsizes(writer_tgt.get_extra_info("socket"), get_to_clt_bufsize(), get_to_tg_bufsize())
+
+        if init_func:
+            return await asyncio.wait_for(init_func(host, port, reader_tgt, writer_tgt),
+                                          timeout=config.TG_CONNECT_TIMEOUT)
+        return reader_tgt, writer_tgt
+
+    def register_host_port(self, host, port, init_func):
+        if (host, port, init_func) not in self.pools:
+            self.pools[(host, port, init_func)] = []
+
+        while len(self.pools[(host, port, init_func)]) < TgConnectionPool.MAX_CONNS_IN_POOL:
+            connect_task = asyncio.ensure_future(self.open_tg_connection(host, port, init_func))
+            self.pools[(host, port, init_func)].append(connect_task)
+
+    async def get_connection(self, host, port, init_func=None):
+        self.register_host_port(host, port, init_func)
+
+        ret = None
+        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():
+                    self.pools[(host, port, init_func)].remove(task)
+                    continue
+
+                if not ret:
+                    self.pools[(host, port, init_func)].remove(task)
+                    ret = (reader, writer, *other)
+
+        self.register_host_port(host, port, init_func)
+        if ret:
+            return ret
+        return await self.open_tg_connection(host, port, init_func)
+
+
+tg_connection_pool = TgConnectionPool()
+
+
 class LayeredStreamReaderBase:
     __slots__ = ("upstream", )
 
@@ -1191,21 +1250,6 @@ async def handle_handshake(reader, writer):
     return False
 
 
-async def open_connection_tryer(addr, port, limit, timeout, max_attempts=3):
-    for attempt in range(max_attempts-1):
-        try:
-            task = asyncio.open_connection(addr, port, limit=limit)
-            reader_tgt, writer_tgt = await asyncio.wait_for(task, timeout=timeout)
-            return reader_tgt, writer_tgt
-        except (OSError, asyncio.TimeoutError):
-            continue
-
-    # the last attempt
-    task = asyncio.open_connection(addr, port, limit=limit)
-    reader_tgt, writer_tgt = await asyncio.wait_for(task, timeout=timeout)
-    return reader_tgt, writer_tgt
-
-
 async def do_direct_handshake(proto_tag, dc_idx, dec_key_and_iv=None):
     RESERVED_NONCE_FIRST_CHARS = [b"\xef"]
     RESERVED_NONCE_BEGININGS = [b"\x48\x45\x41\x44", b"\x50\x4F\x53\x54",
@@ -1214,6 +1258,7 @@ async def do_direct_handshake(proto_tag, dc_idx, dec_key_and_iv=None):
     RESERVED_NONCE_CONTINUES = [b"\x00\x00\x00\x00"]
 
     global my_ip_info
+    global tg_connection_pool
 
     dc_idx = abs(dc_idx) - 1
 
@@ -1227,18 +1272,17 @@ async def do_direct_handshake(proto_tag, dc_idx, dec_key_and_iv=None):
         dc = TG_DATACENTERS_V4[dc_idx]
 
     try:
-        reader_tgt, writer_tgt = await open_connection_tryer(
-            dc, TG_DATACENTER_PORT, limit=get_to_clt_bufsize(), timeout=config.TG_CONNECT_TIMEOUT)
+        reader_tgt, writer_tgt = await tg_connection_pool.get_connection(dc, TG_DATACENTER_PORT)
     except ConnectionRefusedError as E:
         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))
+        return False
     except (OSError, asyncio.TimeoutError) as E:
         print_err("Unable to connect to", dc, TG_DATACENTER_PORT)
         return False
 
-    set_keepalive(writer_tgt.get_extra_info("socket"))
-    set_bufsizes(writer_tgt.get_extra_info("socket"), get_to_clt_bufsize(), get_to_tg_bufsize())
-
     while True:
         rnd = bytearray(myrandom.getrandbytes(HANDSHAKE_LEN))
         if rnd[:1] in RESERVED_NONCE_FIRST_CHARS:
@@ -1301,49 +1345,21 @@ def get_middleproxy_aes_key_and_iv(nonce_srv, nonce_clt, clt_ts, srv_ip, clt_por
     return key, iv
 
 
-async def do_middleproxy_handshake(proto_tag, dc_idx, cl_ip, cl_port):
+async def middleproxy_handshake(host, port, reader_tgt, writer_tgt):
+    """ The most logic of middleproxy handshake, launched in pool """
     START_SEQ_NO = -2
     NONCE_LEN = 16
 
-    RPC_NONCE = b"\xaa\x87\xcb\x7a"
     RPC_HANDSHAKE = b"\xf5\xee\x82\x76"
+    RPC_NONCE = b"\xaa\x87\xcb\x7a"
+    # pass as consts to simplify code
+    RPC_FLAGS = b"\x00\x00\x00\x00"
     CRYPTO_AES = b"\x01\x00\x00\x00"
 
     RPC_NONCE_ANS_LEN = 32
     RPC_HANDSHAKE_ANS_LEN = 32
 
-    # pass as consts to simplify code
-    RPC_FLAGS = b"\x00\x00\x00\x00"
-
-    global my_ip_info
-
-    use_ipv6_tg = (my_ip_info["ipv6"] and (config.PREFER_IPV6 or not my_ip_info["ipv4"]))
-    use_ipv6_clt = (":" in cl_ip)
-
-    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])
-    else:
-        if dc_idx not in TG_MIDDLE_PROXIES_V4:
-            return False
-        addr, port = myrandom.choice(TG_MIDDLE_PROXIES_V4[dc_idx])
-
-    try:
-        reader_tgt, writer_tgt = await open_connection_tryer(addr, port, limit=get_to_clt_bufsize(),
-                                                             timeout=config.TG_CONNECT_TIMEOUT)
-    except ConnectionRefusedError as E:
-        print_err("The Telegram server %d (%s %s) is refusing connections" % (dc_idx, addr, port))
-        return False
-    except (OSError, asyncio.TimeoutError) as E:
-        print_err("Unable to connect to the Telegram server %d (%s %s)" % (dc_idx, addr, port))
-        return False
-
-    set_keepalive(writer_tgt.get_extra_info("socket"))
-    set_bufsizes(writer_tgt.get_extra_info("socket"), get_to_clt_bufsize(), get_to_tg_bufsize())
-
     writer_tgt = MTProtoFrameStreamWriter(writer_tgt, START_SEQ_NO)
-
     key_selector = PROXY_SECRET[:4]
     crypto_ts = int.to_bytes(int(time.time()) % (256**4), 4, "little")
 
@@ -1354,24 +1370,25 @@ async def do_middleproxy_handshake(proto_tag, dc_idx, cl_ip, cl_port):
     writer_tgt.write(msg)
     await writer_tgt.drain()
 
-    old_reader = reader_tgt
     reader_tgt = MTProtoFrameStreamReader(reader_tgt, START_SEQ_NO)
     ans = await reader_tgt.read(get_to_clt_bufsize())
 
     if len(ans) != RPC_NONCE_ANS_LEN:
-        return False
+        raise ConnectionAbortedError("bad rpc answer length")
 
     rpc_type, rpc_key_selector, rpc_schema, rpc_crypto_ts, rpc_nonce = (
         ans[:4], ans[4:8], ans[8:12], ans[12:16], ans[16:32]
     )
 
     if rpc_type != RPC_NONCE or rpc_key_selector != key_selector or rpc_schema != CRYPTO_AES:
-        return False
+        raise ConnectionAbortedError("bad rpc answer")
 
     # get keys
     tg_ip, tg_port = writer_tgt.upstream.get_extra_info('peername')[:2]
     my_ip, my_port = writer_tgt.upstream.get_extra_info('sockname')[:2]
 
+    use_ipv6_tg = (":" in tg_ip)
+
     if not use_ipv6_tg:
         if my_ip_info["ipv4"]:
             # prefer global ip settings to work behind NAT
@@ -1422,11 +1439,42 @@ async def do_middleproxy_handshake(proto_tag, dc_idx, cl_ip, cl_port):
 
     handshake_ans = await reader_tgt.read(1)
     if len(handshake_ans) != RPC_HANDSHAKE_ANS_LEN:
-        return False
+        raise ConnectionAbortedError("bad rpc handshake answer length")
 
     handshake_type, handshake_flags, handshake_sender_pid, handshake_peer_pid = (
         handshake_ans[:4], handshake_ans[4:8], handshake_ans[8:20], handshake_ans[20:32])
     if handshake_type != RPC_HANDSHAKE or handshake_peer_pid != SENDER_PID:
+        raise ConnectionAbortedError("bad rpc handshake answer")
+
+    return reader_tgt, writer_tgt, my_ip, my_port
+
+
+async def do_middleproxy_handshake(proto_tag, dc_idx, cl_ip, cl_port):
+    global my_ip_info
+    global tg_connection_pool
+
+    use_ipv6_tg = (my_ip_info["ipv6"] and (config.PREFER_IPV6 or not my_ip_info["ipv4"]))
+
+    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])
+    else:
+        if dc_idx not in TG_MIDDLE_PROXIES_V4:
+            return False
+        addr, port = myrandom.choice(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))
         return False
 
     writer_tgt = ProxyReqStreamWriter(writer_tgt, cl_ip, cl_port, my_ip, my_port, proto_tag)
@@ -1528,7 +1576,7 @@ async def handle_client(reader_clt, writer_clt):
                     update_user_stats(user, octets=len(data), msgs=1)
                     wr.write(data, extra)
                     await wr.drain()
-        except (OSError, asyncio.streams.IncompleteReadError) as e:
+        except (OSError, asyncio.IncompleteReadError) as e:
             # print_err(e)
             pass
 
@@ -1586,6 +1634,7 @@ def make_metrics_pkt(metrics):
     used_names = set()
 
     for name, m_type, desc, val in metrics:
+        name = config.METRICS_PREFIX + name
         if name not in used_names:
             pkt_body_list.append("# HELP %s %s" % (name, desc))
             pkt_body_list.append("# TYPE %s %s" % (name, m_type))
@@ -1803,7 +1852,11 @@ async def get_mask_host_cert_len():
             task = get_encrypted_cert(config.MASK_HOST, config.MASK_PORT, config.TLS_DOMAIN)
             cert = await asyncio.wait_for(task, timeout=GET_CERT_TIMEOUT)
             if cert:
-                if len(cert) != fake_cert_len:
+                if len(cert) < MIN_CERT_LEN:
+                    msg = ("The MASK_HOST %s returned several TLS records, this is not supported" %
+                           config.MASK_HOST)
+                    print_err(msg)
+                elif len(cert) != fake_cert_len:
                     fake_cert_len = len(cert)
                     print_err("Got cert from the MASK_HOST %s, its length is %d" %
                               (config.MASK_HOST, fake_cert_len))
@@ -2001,7 +2054,8 @@ def print_tg_info():
             proxy_links.append({"user": user, "link": tls_link})
             print("{}: {}".format(user, tls_link), flush=True)
 
-        if secret in ["00000000000000000000000000000000", "0123456789abcdef0123456789abcdef"]:
+        if secret in ["00000000000000000000000000000000", "0123456789abcdef0123456789abcdef",
+                      "00000000000000000000000000000001"]:
             msg = "The default secret {} is used, this is not recommended".format(secret)
             print(msg, flush=True)
             random_secret = "".join(myrandom.choice("0123456789abcdef") for i in range(32))
@@ -2133,29 +2187,29 @@ def main():
 
     if config.LISTEN_ADDR_IPV4:
         task = asyncio.start_server(handle_client_wrapper, config.LISTEN_ADDR_IPV4, config.PORT,
-                                    limit=get_to_tg_bufsize(), reuse_port=reuse_port, loop=loop)
+                                    limit=get_to_tg_bufsize(), reuse_port=reuse_port)
         servers.append(loop.run_until_complete(task))
 
     if config.LISTEN_ADDR_IPV6 and socket.has_ipv6:
         task = asyncio.start_server(handle_client_wrapper, config.LISTEN_ADDR_IPV6, config.PORT,
-                                    limit=get_to_tg_bufsize(), reuse_port=reuse_port, loop=loop)
+                                    limit=get_to_tg_bufsize(), reuse_port=reuse_port)
         servers.append(loop.run_until_complete(task))
 
     if config.LISTEN_UNIX_SOCK and has_unix:
         remove_unix_socket(config.LISTEN_UNIX_SOCK)
         task = asyncio.start_unix_server(handle_client_wrapper, config.LISTEN_UNIX_SOCK,
-                                         limit=get_to_tg_bufsize(), loop=loop)
+                                         limit=get_to_tg_bufsize())
         servers.append(loop.run_until_complete(task))
         os.chmod(config.LISTEN_UNIX_SOCK, 0o666)
 
     if config.METRICS_PORT is not None:
         if config.METRICS_LISTEN_ADDR_IPV4:
             task = asyncio.start_server(handle_metrics, config.METRICS_LISTEN_ADDR_IPV4,
-                                        config.METRICS_PORT, loop=loop)
+                                        config.METRICS_PORT)
             servers.append(loop.run_until_complete(task))
         if config.METRICS_LISTEN_ADDR_IPV6 and socket.has_ipv6:
             task = asyncio.start_server(handle_metrics, config.METRICS_LISTEN_ADDR_IPV6,
-                                        config.METRICS_PORT, loop=loop)
+                                        config.METRICS_PORT)
             servers.append(loop.run_until_complete(task))
 
     try: