proxynt 2.0.61rc1__tar.gz → 2.1.0__tar.gz

{proxynt-2.0.61rc1 → proxynt-2.1.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: proxynt
-Version: 2.0.61rc1
+Version: 2.1.0
 Summary: UNKNOWN
 Home-page: https://github.com/sazima/proxynt
 License: UNKNOWN

{proxynt-2.0.61rc1 → proxynt-2.1.0}/client/tcp_forward_client.py

@@ -9,7 +9,7 @@ from typing import Dict, List
 
 from common.logger_factory import LoggerFactory
 from common.nat_serialization import NatSerialization
-from common.pool import SelectPool
+from common.pool import SelectPool, PAUSE_REASON_FLOW_CONTROL
 from common.register_append_data import ResisterAppendData
 from common.speed_limit import SpeedLimiter
 from constant.message_type_constnat import MessageTypeConstant
@@ -38,6 +38,7 @@ class TcpForwardClient:
         self.ws = ws_sender  # WsSender — thread-safe, send via IOLoop.add_callback
         self.lock = Lock()
         self.socket_event_loop = SelectPool()
+        self.flow_control_enabled = False  # 服务端是否支持按连接流控, 由 PUSH_CONFIG 协商
 
         # C2C client-to-client forward state
         self.c2c_rules: List[dict] = []
@@ -52,6 +53,10 @@ class TcpForwardClient:
         self.pending_data_lock = Lock()
         self.pending_data_timeout = 10
 
+        # 已关闭的 uid -> 关闭时间；服务端的关闭通知晚于本地关闭时，避免为旧 uid 重建后端连接
+        self.recently_closed: Dict[bytes, float] = {}
+        self.recently_closed_ttl = 60
+
     def set_running(self, running: bool):
         self.socket_event_loop.is_running = running
 
@@ -223,6 +228,9 @@ class TcpForwardClient:
         with self.lock:
             if uid in self.uid_to_socket_connection:
                 return True
+            if uid in self.recently_closed:
+                LoggerFactory.get_logger().info('ignore data for closed uid: %s' % uid.hex())
+                return False
             try:
                 if LoggerFactory.get_logger().isEnabledFor(logging.DEBUG):
                     LoggerFactory.get_logger().debug('create socket %s, %s' % (name, uid))
@@ -301,6 +309,7 @@ class TcpForwardClient:
             connection: PrivateSocketConnection = self.socket_to_socket_connection.pop(socket_client)
             LoggerFactory.get_logger().info('close uid: %s, sent=%d recv=%d, alive=%.3fs' % (
                 connection.uid.hex(), connection.bytes_sent, connection.bytes_recv, time.time() - connection.create_time))
+            self._mark_closed(connection.uid)
             self.socket_event_loop.unregister(socket_client)
             try:
                 socket_client.shutdown(socket.SHUT_RDWR)
@@ -312,6 +321,13 @@ class TcpForwardClient:
                 self.uid_to_socket_connection.pop(connection.uid)
             self.close_remote_socket(connection)
 
+    def _mark_closed(self, uid: bytes):
+        now = time.time()
+        self.recently_closed[uid] = now
+        if len(self.recently_closed) > 1000:
+            for k in [k for k, t in self.recently_closed.items() if now - t > self.recently_closed_ttl]:
+                self.recently_closed.pop(k, None)
+
     def close(self):
         LoggerFactory.get_logger().info('Starting to close %s' % self.c2c_listeners)
         with self.lock:
@@ -367,6 +383,18 @@ class TcpForwardClient:
         )
         LoggerFactory.get_logger().debug('Send to websocket cost time %s' % (time.time() - start_time))
 
+    def set_local_reading_paused(self, uid: bytes, pause: bool):
+        """服务端通知该 uid 积压过多, 暂停/恢复读取本地连接"""
+        connection = self.uid_to_socket_connection.get(uid)
+        if not connection:
+            return
+        if LoggerFactory.get_logger().isEnabledFor(logging.DEBUG):
+            LoggerFactory.get_logger().debug('flow control %s uid: %s' % ('pause' if pause else 'resume', uid.hex()))
+        if pause:
+            self.socket_event_loop.pause_reading(connection.socket, PAUSE_REASON_FLOW_CONTROL)
+        else:
+            self.socket_event_loop.resume_reading(connection.socket, PAUSE_REASON_FLOW_CONTROL)
+
     def send_by_uid(self, uid: bytes, msg: bytes):
         connection = self.uid_to_socket_connection.get(uid)
         if not connection:

{proxynt-2.0.61rc1 → proxynt-2.1.0}/common/nat_serialization_v1.py

@@ -74,6 +74,8 @@ class NatSerializationV1:
                                    magic, mode_flag, len(target_client), len(target_service), len(source_rule_name), len(protocol),
                                    uid, target_client, target_service, source_rule_name, protocol)
 
+        elif type_ in (MessageTypeConstant.FLOW_PAUSE, MessageTypeConstant.FLOW_RESUME):
+            body = struct.pack(f'{UID_LEN}s', data['data']['uid'])
         elif type_ == MessageTypeConstant.PUSH_CONFIG:
             body = json.dumps(data).encode()
         elif type_ == MessageTypeConstant.PING:
@@ -190,6 +192,13 @@ class NatSerializationV1:
             }
             return return_data
 
+        elif type_.decode() in (MessageTypeConstant.FLOW_PAUSE, MessageTypeConstant.FLOW_RESUME):
+            return_data: MessageEntity = {
+                'type_': type_.decode(),
+                'data': {'uid': body[:UID_LEN]}
+            }
+            return return_data
+
         elif type_ == MessageTypeConstant.PUSH_CONFIG.encode():
             return_data: MessageEntity = json.loads(body.decode())
             return return_data

{proxynt-2.0.61rc1 → proxynt-2.1.0}/common/pool.py

@@ -4,6 +4,7 @@ import threading
 import time
 import traceback
 from selectors import DefaultSelector, EVENT_READ
+from typing import Dict, Set
 
 from common.logger_factory import LoggerFactory
 from common.register_append_data import ResisterAppendData
@@ -14,6 +15,9 @@ from constant.system_constant import SystemConstant
 参考 shadowsocks eventloop 模式
 """
 
+PAUSE_REASON_SPEED_LIMIT = 'speed_limit'
+PAUSE_REASON_FLOW_CONTROL = 'flow_control'
+
 
 class SelectPool:
 
@@ -21,6 +25,27 @@ class SelectPool:
         self.is_running = True
         self.fileno_to_client = dict()
         self.selector = DefaultSelector()
+        # 其他线程 register 时唤醒阻塞中的 select，否则 Windows 上新 socket 要等本轮 select 超时才被监听
+        self._wakeup_r, self._wakeup_w = socket.socketpair()
+        self._wakeup_r.setblocking(False)
+        self._wakeup_w.setblocking(False)
+        self.selector.register(self._wakeup_r, EVENT_READ, None)
+        # 每个 socket 的暂停原因 (如 speed_limit / flow_control), 全部解除才恢复读取
+        self._pause_reasons: Dict[socket.socket, Set[str]] = dict()
+        self._pause_lock = threading.Lock()
+
+    def _wakeup(self):
+        try:
+            self._wakeup_w.send(b'\x00')
+        except OSError:
+            pass
+
+    def _drain_wakeup(self):
+        try:
+            while self._wakeup_r.recv(4096):
+                pass
+        except OSError:
+            pass
 
     def stop(self):
         self.is_running = False
@@ -31,6 +56,7 @@ class SelectPool:
     def register(self, s: socket.socket, data: ResisterAppendData):
         self.fileno_to_client[s.fileno()] = s
         self.selector.register(s, EVENT_READ, data)
+        self._wakeup()
 
     def unregister(self, s: socket.socket):
         fileno = -1
@@ -40,6 +66,8 @@ class SelectPool:
             pass
         if fileno in self.fileno_to_client:
             self.fileno_to_client.pop(fileno)
+        with self._pause_lock:
+            self._pause_reasons.pop(s, None)
         try:
             self.selector.unregister(s)
         except (KeyError, ValueError):
@@ -47,26 +75,43 @@ class SelectPool:
         except OSError:
             LoggerFactory.get_logger().error(traceback.format_exc())
 
-    def pause_reading(self, s: socket.socket):
+    def pause_reading(self, s: socket.socket, reason: str = PAUSE_REASON_SPEED_LIMIT):
         """暂停监听可读事件 (modify 掩码为 0)"""
-        try:
-            key = self.selector.get_key(s)
-            self.selector.modify(s, 0, key.data)
-        except (KeyError, ValueError):
-            pass
+        with self._pause_lock:
+            try:
+                key = self.selector.get_key(s)
+            except (KeyError, ValueError):
+                return
+            reasons = self._pause_reasons.setdefault(s, set())
+            first_pause = not reasons
+            reasons.add(reason)
+            if first_pause:
+                try:
+                    self.selector.modify(s, 0, key.data)
+                except (KeyError, ValueError):
+                    pass
 
-    def resume_reading(self, s: socket.socket):
-        """恢复监听可读事件"""
-        try:
-            key = self.selector.get_key(s)
-            self.selector.modify(s, EVENT_READ, key.data)
-        except (KeyError, ValueError):
-            pass
+    def resume_reading(self, s: socket.socket, reason: str = PAUSE_REASON_SPEED_LIMIT):
+        """解除一个暂停原因, 没有其他原因时恢复监听可读事件"""
+        with self._pause_lock:
+            reasons = self._pause_reasons.get(s)
+            if not reasons or reason not in reasons:
+                return
+            reasons.discard(reason)
+            if reasons:
+                return
+            self._pause_reasons.pop(s, None)
+            try:
+                key = self.selector.get_key(s)
+                self.selector.modify(s, EVENT_READ, key.data)
+            except (KeyError, ValueError):
+                return
+        self._wakeup()
 
     def pause_and_resume_later(self, s: socket.socket, delay_time: float):
         """暂停可读监听, delay_time 秒后恢复"""
-        self.pause_reading(s)
-        threading.Timer(delay_time, self.resume_reading, args=(s,)).start()
+        self.pause_reading(s, PAUSE_REASON_SPEED_LIMIT)
+        threading.Timer(delay_time, self.resume_reading, args=(s, PAUSE_REASON_SPEED_LIMIT)).start()
 
     def run(self):
         while True:
@@ -80,6 +125,9 @@ class SelectPool:
                     time.sleep(0.5)
                     continue
                 for key, mask in ready:
+                    if key.data is None:
+                        self._drain_wakeup()
+                        continue
                     fileno = key.fd
                     client = self.fileno_to_client.get(fileno)
                     if client is None:

{proxynt-2.0.61rc1 → proxynt-2.1.0}/constant/message_type_constnat.py

@@ -28,4 +28,9 @@ class MessageTypeConstant:
     P2P_PRE_CONNECT = 'g'
     P2P_EXCHANGE = 'h'        # UDP exchange to establish NAT mapping
     P2P_PEER_INFO = 'i'       # Server sends peer's actual UDP port
-    P2P_PUNCH_REQUEST = 'j'   # Request to initiate P2P hole punching
+    P2P_PUNCH_REQUEST = 'j'   # Request to initiate P2P hole punching
+
+    # 按连接流控: 接收方某个 uid 积压过多时通知发送方暂停/恢复读取该 uid 的数据源
+    # 仅在双方通过 PUSH_CONFIG 协商支持后发送 (旧版本反序列化未知类型会报错)
+    FLOW_PAUSE = 'k'
+    FLOW_RESUME = 'l'

{proxynt-2.0.61rc1 → proxynt-2.1.0}/constant/system_constant.py

@@ -15,7 +15,7 @@ class SystemConstant:
 
     COOKIE_EXPIRE_SECONDS = 3600 * 24
 
-    VERSION = '2.0.61rc1'
+    VERSION = '2.1.0'
 
     GITHUB = 'https://github.com/sazima/proxynt'
 

proxynt-2.1.0/p2ptest/tls_handshake_timeout_test.py

@@ -0,0 +1,66 @@
+"""
+测试后端 TLS 握手超时：发送 ClientHello，收完服务端握手数据后故意不回 Finished，
+测量后端多久主动断开。
+
+用法: python tls_handshake_timeout_test.py 10.253.68.143 31015 [最长等待秒数]
+"""
+import socket
+import ssl
+import sys
+import time
+
+
+def main():
+    host = sys.argv[1]
+    port = int(sys.argv[2])
+    max_wait = float(sys.argv[3]) if len(sys.argv) > 3 else 30
+
+    ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT)
+    ctx.check_hostname = False
+    ctx.verify_mode = ssl.CERT_NONE
+
+    incoming = ssl.MemoryBIO()
+    outgoing = ssl.MemoryBIO()
+    tls = ctx.wrap_bio(incoming, outgoing)
+
+    try:
+        tls.do_handshake()
+    except ssl.SSLWantReadError:
+        pass
+    client_hello = outgoing.read()
+
+    s = socket.create_connection((host, port), timeout=5)
+    s.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
+    t0 = time.time()
+    s.sendall(client_hello)
+    print(f'[{0:.3f}s] connected, sent ClientHello {len(client_hello)} bytes')
+
+    s.settimeout(0.2)
+    total = 0
+    last_data = None
+    while True:
+        elapsed = time.time() - t0
+        if elapsed > max_wait:
+            print(f'[{elapsed:.3f}s] backend still open after {max_wait}s, giving up (no short handshake timeout)')
+            break
+        try:
+            data = s.recv(65536)
+        except socket.timeout:
+            continue
+        except OSError as e:
+            print(f'[{elapsed:.3f}s] recv error: {type(e).__name__} {e}  (total recv {total} bytes, '
+                  f'last data at {last_data:.3f}s)')
+            break
+        if not data:
+            print(f'[{elapsed:.3f}s] backend closed (EOF)  (total recv {total} bytes, '
+                  f'last data at {last_data:.3f}s)')
+            break
+        total += len(data)
+        last_data = elapsed
+        print(f'[{elapsed:.3f}s] recv {len(data)} bytes (total {total}), not replying...')
+
+    s.close()
+
+
+if __name__ == '__main__':
+    main()

{proxynt-2.0.61rc1 → proxynt-2.1.0}/proxynt.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: proxynt
-Version: 2.0.61rc1
+Version: 2.1.0
 Summary: UNKNOWN
 Home-page: https://github.com/sazima/proxynt
 License: UNKNOWN

{proxynt-2.0.61rc1 → proxynt-2.1.0}/proxynt.egg-info/SOURCES.txt

@@ -119,6 +119,7 @@ setup.py
 ./p2ptest/__init__.py
 ./p2ptest/client.py
 ./p2ptest/n4.py
+./p2ptest/tls_handshake_timeout_test.py
 ./server/__init__.py
 ./server/admin_http_handler.py
 ./server/n4.py

{proxynt-2.0.61rc1 → proxynt-2.1.0}/run_client.py

@@ -79,6 +79,7 @@ class WsSender:
     """Thread-safe WebSocket sender wrapping Tornado WebSocketClientConnection."""
 
     MAX_PENDING_WRITES = 32
+    WRITE_BACKPRESSURE_TIMEOUT = 10
 
     def __init__(self):
         self._ws = None
@@ -104,24 +105,30 @@ class WsSender:
             io.add_callback(ws.write_message, data, binary=True)
             return
 
-        acquired = self._write_sem.acquire(timeout=10)
-        if not acquired or self._ws is None:
+        write_sem = self._write_sem
+        acquired = write_sem.acquire(timeout=self.WRITE_BACKPRESSURE_TIMEOUT)
+        if self._ws is None:
+            if acquired:
+                write_sem.release()
             return
-        io.add_callback(self._do_write, data)
+        if not acquired:
+            # 不能丢弃: TCP 流中间少一块数据会导致对端协议错乱(如 sshd Bad packet length)
+            LoggerFactory.get_logger().warning(f'websocket write backlog over {self.WRITE_BACKPRESSURE_TIMEOUT}s, send without backpressure')
+        io.add_callback(self._do_write, data, write_sem if acquired else None)
 
-    def _do_write(self, data):
+    def _do_write(self, data, write_sem):
         ws = self._ws
         if ws is None:
-            self._write_sem.release()
+            if write_sem:
+                write_sem.release()
             return
         try:
             future = ws.write_message(data, binary=True)
-            ioloop.IOLoop.current().add_future(future, lambda f: self._on_write_done())
+            if write_sem:
+                ioloop.IOLoop.current().add_future(future, lambda f: write_sem.release())
         except Exception:
-            self._write_sem.release()
-
-    def _on_write_done(self):
-        self._write_sem.release()
+            if write_sem:
+                write_sem.release()
 
     def close(self):
         ws = self._ws
@@ -295,6 +302,7 @@ class WebsocketClient:
                 push_config: PushConfigEntity = message_data['data']
                 self.public_ip = push_config.get('public_ip')
                 self.public_port = push_config.get('public_port')
+                self.forward_client.flow_control_enabled = bool(push_config.get('server_flow_control', False))
 
                 for d in push_config['config_list']:
                     if d.get('speed_limit'):
@@ -306,6 +314,11 @@ class WebsocketClient:
                     self.forward_client.setup_c2c_tcp_listeners(c2c_rules)
                     self.udp_forward_client.setup_c2c_udp_listeners(c2c_rules)
 
+            elif msg_type == MessageTypeConstant.FLOW_PAUSE:
+                self.forward_client.set_local_reading_paused(message_data['data']['uid'], True)
+            elif msg_type == MessageTypeConstant.FLOW_RESUME:
+                self.forward_client.set_local_reading_paused(message_data['data']['uid'], False)
+
             # P2P punch disabled
             elif msg_type == MessageTypeConstant.P2P_PUNCH_REQUEST:
                 LoggerFactory.get_logger().info('P2P punch disabled, ignoring request')
@@ -337,7 +350,8 @@ class WebsocketClient:
                     'config_list': push_client_data,
                     "client_name": client_name,
                     'version': SystemConstant.VERSION,
-                    'p2p_supported': False
+                    'p2p_supported': False,
+                    'flow_control': True
                 }
                 message: MessageEntity = {
                     'type_': MessageTypeConstant.PUSH_CONFIG,