diff --git a/engine/server.py b/engine/server.py index c7d4f2b..b99cbae 100644 --- a/engine/server.py +++ b/engine/server.py @@ -373,13 +373,27 @@ class ConnectionHandler: await self._record_stats() async def _forward(self, src, dst, direction, speed_limit_mbps, user): - """单向转发，带限速。""" + """单向转发，带限速与 idle timeout。 + + timeout 行为：握手/认证/请求阶段和这里都共用 self.config.timeout。 + 客户端在 tunnel 阶段不发数据超过 timeout 秒，会被服务端主动关闭， + 释放 fd + 关闭对端 writer + 触发 _tunnel 的 finally 清理。 + """ if src is None or dst is None: return total_transferred = 0 + # 读超时短于 config.timeout 时宁可提前踢,不放过慢客户端 + read_timeout = max(1, int(self.config.timeout)) try: while True: - data = await src.read(TUNNEL_CHUNK) + try: + data = await asyncio.wait_for( + src.read(TUNNEL_CHUNK), timeout=read_timeout + ) + except asyncio.TimeoutError: + log.info("[%s] %s idle timeout (%ds), closing tunnel", + self.conn_id, direction, read_timeout) + break if not data: break diff --git a/models.py b/models.py index 7349a54..7d02bda 100644 --- a/models.py +++ b/models.py @@ -13,6 +13,9 @@ class Instance(db.Model): engine = db.Column(db.String(16), default="builtin") # builtin / 3proxy listen_host = db.Column(db.String(64), default="0.0.0.0") listen_port = db.Column(db.Integer, nullable=False) + # 统一超时(秒):覆盖 SOCKS5 握手/认证/请求读取,以及隧道阶段每段 read。 + # 客户端在任意阶段(含 tunnel)空闲超过此值会被服务端主动断开,避免空连接 + # 占满 fd 直到 LimitNOFILE 上限。 timeout = db.Column(db.Integer, default=30) enabled = db.Column(db.Boolean, default=True) notes = db.Column(db.Text) diff --git a/tests/smoke_tunnel_timeout.py b/tests/smoke_tunnel_timeout.py new file mode 100644 index 0000000..14ccf00 --- /dev/null +++ b/tests/smoke_tunnel_timeout.py @@ -0,0 +1,138 @@ +"""最小烟雾测试: 验证 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 + async def add_traffic(self, *a, **kw): pass + 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)