socks-manager/tests/smoke_tunnel_timeout.py

"""最小烟雾测试: 验证 SOCKS5 隧道转发阶段会被 idle timeout 踢掉。

不需要 Flask app_context, 直接构造 fake deps 喂给 Socks5Server。
"""
import asyncio
import logging
import os
import socket
import struct
import sys
import time

# 让 import 找到 engine/ — 走 cwd 而非 __file__, 因为目录名带中文/空格时不可靠
sys.path.insert(0, os.getcwd())

from engine.instances import Socks5Server, InstanceConfig


class FakeUser:
    """最少够 server.py 跑通"""
    def __init__(self):
        self.bandwidth_down = 0
        self.bandwidth_up = 0
        self.ip_whitelist = ""
        self.ip_blacklist = ""
        self.max_concurrent = 100

    def check_password(self, p): return False
    def is_active(self): return True


class FakeUserService:
    def get_user(self, name): return FakeUser()
    def get_active_connections(self, name): return 0
    def add_traffic(self, *a, **kw): pass   # 同步, 不要 await (见 services/user_service.py:162)
    async def log_event(self, *a, **kw): pass


class FakeFlaskApp:
    """handle() 会调 self.flask_app.app_context(), 给个最小 ctx manager"""
    class _Ctx:
        def __enter__(self): return self
        def __exit__(self, *a): return False
    def app_context(self): return self._Ctx()


def find_free_port():
    s = socket.socket()
    s.bind(("127.0.0.1", 0))
    p = s.getsockname()[1]
    s.close()
    return p


async def silent_listener(port):
    """起一个真在 listen 但永远不 accept 的目标, 逼客户端进 tunnel"""
    async def cb(r, w):
        # 不 accept 也不做任何事, 保持端口占用即可
        await asyncio.sleep(3600)
    return await asyncio.start_server(cb, "127.0.0.1", port)


async def handshake_no_auth(host, port, target_host="127.0.0.1", target_port=1):
    """完整 SOCKS5 握手 + 认证 (no-auth) + CONNECT"""
    r, w = await asyncio.open_connection(host, port)
    # 1. method negotiation
    w.write(struct.pack("!BB", 5, 1) + bytes([0x00]))
    await w.drain()
    sel = await r.readexactly(2)
    assert sel == bytes([5, 0]), f"method sel = {sel!r}"
    # 2. CONNECT
    req = struct.pack("!BBBB", 5, 1, 0, 1) + socket.inet_aton(target_host) + struct.pack("!H", target_port)
    w.write(req)
    await w.drain()
    # 3. 读 response
    resp_head = await asyncio.wait_for(r.readexactly(4), timeout=5)
    rep = resp_head[1]
    atyp = resp_head[3]
    extra = {1: 4, 3: 1, 4: 16}.get(atyp)
    assert extra is not None
    tail = await asyncio.wait_for(r.readexactly(extra + 2), timeout=5)
    return r, w, rep


async def run_test(timeout_sec: int):
    socks_port = find_free_port()
    target_port = find_free_port()
    # 起一个真 listen 但不 accept 的目标
    target_srv = await silent_listener(target_port)
    try:
        cfg = InstanceConfig(
            instance_id=1, name="t", listen_host="127.0.0.1", listen_port=socks_port,
            timeout=timeout_sec, auth_method="none",
            bandwidth_down=0, bandwidth_up=0, max_concurrent=100,
        )
        srv = Socks5Server(cfg, FakeUserService(), flask_app=FakeFlaskApp())
        srv.start()
        for _ in range(50):
            if srv.server: break
            await asyncio.sleep(0.05)
        assert srv.server, "SOCKS5 server not listening"

        try:
            # 客户端握手 + CONNECT 到一个 listen 但不 accept 的目标
            r, w, rep = await handshake_no_auth("127.0.0.1", socks_port,
                                                target_host="127.0.0.1",
                                                target_port=target_port)
            assert rep == 0, f"expected REP_SUCCEEDED, got {rep}"
            print(f"  connected to {target_port}, entering tunnel, idle waiting {timeout_sec}s...")
            t0 = time.time()
            try:
                data = await asyncio.wait_for(r.read(1), timeout=timeout_sec + 3)
                elapsed = time.time() - t0
                if data == b"":
                    print(f"PASS: 服务端在 {elapsed:.1f}s 关闭了连接 (EOF, timeout={timeout_sec}s)")
                    return elapsed >= timeout_sec * 0.8  # 容许 20% 误差
                else:
                    print(f"FAIL: 收到意外数据 {data!r}")
                    return False
            except asyncio.TimeoutError:
                elapsed = time.time() - t0
                print(f"FAIL: {elapsed:.1f}s 后仍未关闭, timeout 未生效")
                return False
            finally:
                try: w.close()
                except: pass
        finally:
            srv.stop()
    finally:
        target_srv.close()
        await target_srv.wait_closed()


if __name__ == "__main__":
    logging.basicConfig(level=logging.WARNING,
                        format="%(asctime)s [%(levelname)s] %(name)s: %(message)s")
    rc = asyncio.run(run_test(timeout_sec=2))
    sys.exit(0 if rc else 1)