Your Name
ae8334140d
问题: 协程被取消/销毁时，finally 块中的异步操作触发 GeneratorExit 原因: GeneratorExit 是 BaseException 子类，不会被 Exception 捕获 修复内容: - server.py handle(): 捕获 CancelledError + GeneratorExit - server.py _tunnel(): 被取消时不执行 _record_stats() 异步操作 - instances.py _accept_connection(): 添加异常处理和 writer 关闭 - 所有取消场景都静默处理（不抛异常/不记录错误日志）
476 lines
17 KiB
Python
476 lines
17 KiB
Python
"""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
|