"""SOCKS5 异步服务器引擎 (asyncio)。""" import asyncio import logging import struct import time import uuid import ipaddress from datetime import datetime, timezone import models from services.user_service import UserService log = logging.getLogger("socks.engine") # ── SOCKS5 协议常量 ───────────────────────────────────────────── SOCKS_VERSION = 5 METHOD_NO_AUTH = 0x00 METHOD_USERPASS = 0x02 METHOD_NOT_ACCEPTED = 0xFF CMD_CONNECT = 0x01 CMD_BIND = 0x02 CMD_UDP_ASSOCIATE = 0x03 ATYP_IPV4 = 0x01 ATYP_DOMAIN = 0x03 ATYP_IPV6 = 0x04 REP_SUCCEEDED = 0x00 REP_GENERAL_FAILURE = 0x01 REP_CONN_FORBIDDEN = 0x02 REP_NETWORK_UNREACHABLE = 0x03 REP_HOST_UNREACHABLE = 0x04 REP_CONN_REFUSED = 0x05 REP_TTL_EXPIRED = 0x06 REP_CMD_NOT_SUPPORTED = 0x07 REP_ADDR_NOT_SUPPORTED = 0x08 TUNNEL_CHUNK = 65536 class ConnectionHandler: """单个 SOCKS5 连接的处理器。""" def __init__(self, reader, writer, instance_config, user_service, flask_app=None, on_close_cb=None): self.reader = reader self.writer = writer self.config = instance_config self.user_service = user_service self.flask_app = flask_app self.on_close = on_close_cb self.username = None self.src_ip = None self.src_port = None self.dst_addr = None self.dst_port = None self.conn_id = str(uuid.uuid4())[:16] self.bytes_in = 0 self.bytes_out = 0 self.start_time = time.time() self._closed = False # 限速节流器 self._throttle_writer_in = None # client->dst 方向 self._throttle_writer_out = None # dst->client 方向 self._speed_semaphore_in = None self._speed_semaphore_out = None # 上游代理 self.upstream = None async def handle(self): """主入口：处理完整的 SOCKS5 会话。""" try: self.src_ip, self.src_port = self._peer_info() # 1. 方法协商 if not await self._negotiate_method(): return # 2. 认证（如果需要） if not await self._authenticate(): return # 3. 解析命令请求 cmd, atyp, addr, port = await self._read_request() if cmd is None: return # 4. 记录连接 await self._record_connection("active") # 5. 处理命令（目前只支持 CONNECT） if cmd == CMD_CONNECT: await self._handle_connect(addr, port) else: await self._reply(REP_CMD_NOT_SUPPORTED) return except asyncio.TimeoutError: log.debug("[%s] 超时", self.conn_id) except ConnectionError: log.debug("[%s] 连接错误", self.conn_id) except (asyncio.CancelledError, GeneratorExit): # 正常关闭，不记录错误 pass except Exception as e: log.exception("[%s] 异常: %s", self.conn_id, e) finally: if not self._closed: self._close() def _peer_info(self): extra = self.writer.get_extra_info("peername") return extra async def _negotiate_method(self): """方法协商。""" header = await asyncio.wait_for(self.reader.read(256), timeout=self.config.timeout) if len(header) < 2: return False ver, n_methods = struct.unpack("!BB", header[:2]) if ver != SOCKS_VERSION: log.debug("协议版本错误: %d", ver) return False methods = header[2:2 + n_methods] # 检查是否支持 if METHOD_USERPASS in methods and self.config.auth_method == "userpass": self.writer.write(bytes([SOCKS_VERSION, METHOD_USERPASS])) elif METHOD_NO_AUTH in methods: self.writer.write(bytes([SOCKS_VERSION, METHOD_NO_AUTH])) else: self.writer.write(bytes([SOCKS_VERSION, METHOD_NOT_ACCEPTED])) return False await self.writer.drain() return True async def _authenticate(self): """USERPASS 认证。""" if self.config.auth_method != "userpass": return True # RFC 1929 认证请求格式： # +----+------+----------+------+----------+ # |VER | ULEN | UNAME | PLEN | PASSWD | # +----+------+----------+------+----------+ # | 1 | 1 | 1 to 255 | 1 | 1 to 255 | # +----+------+----------+------+----------+ # 注意：整个请求是一个连续 TCP 流，可能一次性全部到达 data = b"" # 1. 先读版本和用户名长度 (2 字节) while len(data) < 2: chunk = await asyncio.wait_for(self.reader.read(256), timeout=self.config.timeout) if not chunk: return False data += chunk ver, ulen = struct.unpack("!BB", data[:2]) if ver != 1: return False offset = 2 # 2. 读取剩余需要的数据 needed = ulen + 1 # 用户名 + 密码长度字节 while len(data) < offset + needed: chunk = await asyncio.wait_for(self.reader.read(256), timeout=self.config.timeout) if not chunk: return False data += chunk username = data[offset:offset + ulen].decode("utf-8", errors="replace") offset += ulen plen = data[offset] offset += 1 # 3. 读取密码 while len(data) < offset + plen: chunk = await asyncio.wait_for(self.reader.read(256), timeout=self.config.timeout) if not chunk: return False data += chunk passwd = data[offset:offset + plen].decode("utf-8", errors="replace") # 验证（使用 bcrypt 兼容验证） user = self.user_service.get_user(username) if user is None or not user.check_password(passwd) or not user.is_active(): log.warning("认证失败: %s (user=%s)", self.src_ip, username) await self.user_service.log_event("auth_fail", username, self.src_ip, "密码错误或账号不可用") self.writer.write(bytes([1, 0x01])) # auth fail await self.writer.drain() return False self.username = username log.info("[%s] 认证成功: %s (from %s)", self.conn_id, username, self.src_ip) await self.user_service.log_event("auth_success", username, self.src_ip) self.writer.write(bytes([1, 0x00])) # auth ok await self.writer.drain() return True async def _read_request(self): """解析 SOCKS5 命令请求。""" req = await asyncio.wait_for(self.reader.read(512), timeout=self.config.timeout) if len(req) < 8: return None, None, None, None ver, cmd, rsv, atyp = struct.unpack("!BBBB", req[:4]) if ver != SOCKS_VERSION: return None, None, None, None offset = 4 if atyp == ATYP_IPV4: addr_bytes = req[offset:offset + 4] addr = ".".join(str(b) for b in addr_bytes) port = struct.unpack("!H", req[offset + 4:offset + 6])[0] return cmd, atyp, addr, port elif atyp == ATYP_IPV6: addr_bytes = req[offset:offset + 16] addr = str(ipaddress.IPv6Address(addr_bytes)) port = struct.unpack("!H", req[offset + 16:offset + 18])[0] return cmd, atyp, addr, port elif atyp == ATYP_DOMAIN: dlen = req[offset] addr = req[offset + 1:offset + 1 + dlen].decode("utf-8", errors="replace") port = struct.unpack("!H", req[offset + 1 + dlen:offset + 3 + dlen])[0] return cmd, atyp, addr, port return None, None, None, None async def _reply(self, rep, bnd_addr="0.0.0.0", bnd_port=0): """发送 SOCKS5 响应 (RFC 1928 §6). 格式: VER(1) REP(1) RSV(1) ATYP(1) BND.ADDR(4/16字节) BND.PORT(2字节) IPv4 是 4 字节 packed, IPv6 是 16 字节, 域名不允许出现在 reply. """ atyp = ATYP_IPV4 if bnd_addr: try: ip = ipaddress.ip_address(str(bnd_addr)) if ip.version == 6: atyp = ATYP_IPV6 addr_bytes = ip.packed else: atyp = ATYP_IPV4 addr_bytes = ip.packed except ValueError: atyp = ATYP_IPV4 addr_bytes = b"\x00\x00\x00\x00" else: atyp = ATYP_IPV4 addr_bytes = b"\x00\x00\x00\x00" self.writer.write(struct.pack( "!BBBB", SOCKS_VERSION, rep, 0x00, atyp, ) + addr_bytes + struct.pack("!H", bnd_port)) await self.writer.drain() async def _handle_connect(self, addr, port): """处理 CONNECT 命令。""" self.dst_addr = addr self.dst_port = port # IP 白名单/黑名单检查 if self.username: user = self.user_service.get_user(self.username) if user and user.ip_blacklist: black = {x.strip() for x in user.ip_blacklist.split(",") if x.strip()} if self.src_ip in black: log.info("[%s] 源IP在黑名单: %s", self.conn_id, self.src_ip) await self._reply(REP_CONN_FORBIDDEN) await self.user_service.log_event("connect_reject", self.username, self.src_ip, f"源IP黑名单 {self.src_ip}") return if user and user.ip_whitelist: white = {x.strip() for x in user.ip_whitelist.split(",") if x.strip()} if self.src_ip not in white: log.info("[%s] 源IP不在白名单: %s", self.conn_id, self.src_ip) await self._reply(REP_CONN_FORBIDDEN) await self.user_service.log_event("connect_reject", self.username, self.src_ip, f"源IP不在白名单 {self.src_ip}") return # 并发连接检查 if self.username: cur_conns = self.user_service.get_active_connections(self.username) if self.username: user = self.user_service.get_user(self.username) if user and cur_conns >= user.max_concurrent: log.info("[%s] 并发超限: %s (当前 %d >= %d)", self.conn_id, self.username, cur_conns, user.max_concurrent) await self._reply(REP_CONN_FORBIDDEN) return # 连接目标 try: # 直接用 open_connection，不预连接 socket # 在子线程事件循环中 asyncio.wait_for 会导致读取被取消 reader, writer = await asyncio.open_connection(addr, port) except (OSError, Exception) as e: log.warning("[%s] 无法连接 %s:%d: %s", self.conn_id, addr, port, e) rep = REP_CONN_REFUSED if "refused" in str(e).lower() else REP_HOST_UNREACHABLE await self._reply(rep) await self.user_service.log_event("connect_reject", self.username, self.src_ip, f"连接失败 {addr}:{port}: {e}") return # 成功响应 await self._reply(REP_SUCCEEDED) # 开始隧道 log.info("[%s] 隧道建立: %s -> %s:%d", self.conn_id, self.src_ip, addr, port) await self.user_service.log_event("connect_success", self.username, self.src_ip, f"{addr}:{port}") await self._tunnel(reader, writer) async def _tunnel(self, dst_reader, dst_writer): """双向隧道转发。""" f1 = f2 = None cancelled = False try: # 获取限速参数 user = None bw_down = self.config.bandwidth_down bw_up = self.config.bandwidth_up if self.username: user = self.user_service.get_user(self.username) if user and user.bandwidth_down: bw_down = user.bandwidth_down if user and user.bandwidth_up: bw_up = user.bandwidth_up # 启动双向转发 f1 = asyncio.create_task(self._forward(dst_reader, self.writer, "dst->client", bw_down, user)) f2 = asyncio.create_task(self._forward(self.reader, dst_writer, "client->dst", bw_up, user)) await asyncio.gather(f1, f2, return_exceptions=True) except asyncio.CancelledError: cancelled = True except GeneratorExit: cancelled = True except Exception as e: log.exception("[%s] 隧道异常: %s", self.conn_id, e) finally: # 确保转发任务被取消 for task in [f1, f2]: if task and not task.done(): task.cancel() try: await task except (asyncio.CancelledError, GeneratorExit): pass self._close_streams(dst_reader, dst_writer) # 记录最终统计 - 协程被取消时不做异步操作（避免 GeneratorExit） if not cancelled: await self._record_stats() async def _forward(self, src, dst, direction, speed_limit_mbps, user): """单向转发，带限速。""" if src is None or dst is None: return total_transferred = 0 try: while True: data = await src.read(TUNNEL_CHUNK) if not data: break # 限速：如果设置了，则节流 if speed_limit_mbps and speed_limit_mbps > 0: bytes_per_sec = speed_limit_mbps * 1000000 / 8 sleep_time = len(data) / bytes_per_sec if sleep_time > 0.001: # 小于 1ms 不节流 await asyncio.sleep(min(sleep_time, 1.0)) dst.write(data) # write() 不是协程，不需要 await await dst.drain() total_transferred += len(data) if direction == "dst->client": self.bytes_in += len(data) else: self.bytes_out += len(data) except (ConnectionError, BrokenPipeError): pass except asyncio.CancelledError: pass except Exception as e: log.debug("[%s] 转发异常 %s: %s", self.conn_id, direction, e) def _close_streams(self, dst_reader, dst_writer): try: if dst_writer is not None: dst_writer.close() if dst_reader is not None: try: dst_reader.feed_eof() except Exception: pass except Exception: pass async def _record_connection(self, state="active"): """记录连接到数据库（异步）。""" try: ctx = self.flask_app.app_context() if self.flask_app else None with ctx: models.db.session.add(models.Connection( id=self.conn_id, instance_id=self.config.instance_id, username=self.username, src_ip=self.src_ip, src_port=self.src_port, dst_addr=self.dst_addr, dst_port=self.dst_port, protocol="SOCKS5", state=state, started_at=datetime.now(timezone.utc), )) models.db.session.commit() except Exception as e: log.error("记录连接失败: %s", e) async def _record_stats(self): """记录最终连接统计。""" duration = time.time() - self.start_time try: ctx = self.flask_app.app_context() if self.flask_app else None with ctx: conn = models.db.session.query(models.Connection).filter_by(id=self.conn_id).first() if conn: conn.state = "closed" conn.ended_at = datetime.now(timezone.utc) conn.bytes_in = self.bytes_in conn.bytes_out = self.bytes_out models.db.session.commit() # 更新用户流量 if self.username and (self.bytes_in or self.bytes_out): await self.user_service.add_traffic(self.username, self.bytes_in, self.bytes_out) except Exception as e: log.error("记录统计失败: %s", e) def _close(self): """关闭连接。""" if self._closed: return self._closed = True try: self.writer.close() except Exception: pass if self.on_close: try: self.on_close(self) except Exception: pass