feat: 重构为 SOCKS5 代理管理平台 v2.0
架构: - engine/server.py: asyncio SOCKS5 服务器引擎(含隧道/认证/限速) - engine/instances.py: 多实例管理器(独立事件循环/热启动/热停止) - services/user_service.py: 用户认证/流量统计/IP白名单/防爆破 - services/stats_service.py: 实时监控/流量趋势/连接追踪 - services/backup_service.py: 数据库备份与恢复 - api/v1.py: 完整 RESTful API - web/routes.py: Web 管理面板(暗色 Bootstrap 5 主题) 功能覆盖: ✅ 多实例 SOCKS5 管理(创建/启动/停止/重启/热加载) ✅ 用户认证(无认证/账号密码/流量上限/限速/并发限制) ✅ IP 白名单/黑名单/防爆破 ✅ 实时仪表盘(CPU/内存/网络/活跃连接) ✅ 流量统计(30天趋势图/用户排名/实例统计) ✅ 活跃连接追踪 ✅ 审计日志(事件/用户/IP 筛选) ✅ 数据库备份与恢复 ✅ 管理员密码保护 ✅ 16 项功能测试全部通过
This commit is contained in:
Your Name
@@ -0,0 +1,246 @@
|
||||
"""多实例管理器：每个实例是一个独立的 asyncio 服务器。"""
|
||||
import asyncio
|
||||
import logging
|
||||
import threading
|
||||
from dataclasses import dataclass, field
|
||||
|
||||
from engine.server import ConnectionHandler
|
||||
|
||||
log = logging.getLogger("socks.instances")
|
||||
|
||||
|
||||
@dataclass
|
||||
class InstanceConfig:
|
||||
"""单个实例的配置快照（运行时使用）。"""
|
||||
instance_id: int
|
||||
name: str
|
||||
listen_host: str
|
||||
listen_port: int
|
||||
timeout: int
|
||||
auth_method: str
|
||||
bandwidth_down: float
|
||||
bandwidth_up: float
|
||||
max_concurrent: int
|
||||
|
||||
def match(self, new_name, new_host, new_port):
|
||||
return (self.name == new_name and
|
||||
self.listen_host == new_host and
|
||||
self.listen_port == new_port)
|
||||
|
||||
|
||||
class Socks5Server:
|
||||
"""单个 SOCKS5 服务器的 asyncio 事件循环。"""
|
||||
|
||||
def __init__(self, config: InstanceConfig, user_service):
|
||||
self.config = config
|
||||
self.user_service = user_service
|
||||
self.server = None
|
||||
self.loop = None
|
||||
self.thread = None
|
||||
self.running = False
|
||||
self._active_connections = 0
|
||||
self._lock = threading.Lock()
|
||||
|
||||
def start(self):
|
||||
"""启动服务器（在新线程中运行独立事件循环）。"""
|
||||
if self.running:
|
||||
return
|
||||
self.loop = asyncio.new_event_loop()
|
||||
asyncio.set_event_loop(self.loop)
|
||||
|
||||
self._active_connections = 0
|
||||
self.running = True
|
||||
|
||||
def run_loop():
|
||||
asyncio.set_event_loop(self.loop)
|
||||
self.server = self.loop.run_until_complete(
|
||||
self._create_server()
|
||||
)
|
||||
self.loop.run_forever()
|
||||
|
||||
self.thread = threading.Thread(target=run_loop, daemon=True,
|
||||
name=f"socks5-{self.config.name}")
|
||||
self.thread.start()
|
||||
# 等 server 创建完成
|
||||
while self.server is None and self.running:
|
||||
time.sleep(0.05)
|
||||
if self.server:
|
||||
log.info("[%s] SOCKS5 服务器已启动: %s:%d",
|
||||
self.config.name, self.config.listen_host, self.config.listen_port)
|
||||
|
||||
def stop(self):
|
||||
"""停止服务器。"""
|
||||
if not self.running:
|
||||
return
|
||||
self.running = False
|
||||
if self.server:
|
||||
self.server.close()
|
||||
if self.loop:
|
||||
try:
|
||||
self.loop.call_soon_threadsafe(self.loop.stop)
|
||||
except Exception:
|
||||
pass
|
||||
if self.thread:
|
||||
self.thread.join(timeout=5)
|
||||
log.info("[%s] SOCKS5 服务器已停止", self.config.name)
|
||||
|
||||
async def _create_server(self):
|
||||
return await asyncio.start_server(
|
||||
self._accept_connection,
|
||||
self.config.listen_host, self.config.listen_port
|
||||
)
|
||||
|
||||
async def _accept_connection(self, reader, writer):
|
||||
self._active_connections += 1
|
||||
handler = ConnectionHandler(
|
||||
reader, writer, self.config, self.user_service,
|
||||
on_close_cb=self._on_connection_close
|
||||
)
|
||||
await handler.handle()
|
||||
|
||||
def _on_connection_close(self, handler):
|
||||
self._active_connections = max(0, self._active_connections - 1)
|
||||
|
||||
@property
|
||||
def active_connections(self):
|
||||
with self._lock:
|
||||
return self._active_connections
|
||||
|
||||
|
||||
import time # noqa: E402
|
||||
|
||||
|
||||
class InstanceManager:
|
||||
"""管理所有 SOCKS5 实例。"""
|
||||
|
||||
def __init__(self, user_service):
|
||||
self.user_service = user_service
|
||||
self._instances = {} # name -> Socks5Server
|
||||
self._config_cache = {} # name -> InstanceConfig
|
||||
self._lock = threading.Lock()
|
||||
|
||||
def reload_config(self):
|
||||
"""从数据库重新加载实例配置。"""
|
||||
from models import Instance
|
||||
from database import db
|
||||
new_configs = {}
|
||||
with db.app.app_context():
|
||||
for inst in Instance.query.filter_by(enabled=True).all():
|
||||
new_configs[inst.name] = InstanceConfig(
|
||||
instance_id=inst.id,
|
||||
name=inst.name,
|
||||
listen_host=inst.listen_host,
|
||||
listen_port=inst.listen_port,
|
||||
timeout=inst.timeout,
|
||||
auth_method=inst.auth_method,
|
||||
bandwidth_down=inst.bandwidth_down,
|
||||
bandwidth_up=inst.bandwidth_up,
|
||||
max_concurrent=inst.max_concurrent,
|
||||
)
|
||||
self._config_cache = new_configs
|
||||
|
||||
def sync_instances(self):
|
||||
"""根据数据库配置同步实例运行状态。"""
|
||||
self.reload_config()
|
||||
|
||||
current_names = set(self._config_cache.keys())
|
||||
running_names = set(self._instances.keys())
|
||||
|
||||
# 启动缺失的实例
|
||||
for name in current_names - running_names:
|
||||
cfg = self._config_cache[name]
|
||||
srv = Socks5Server(cfg, self.user_service)
|
||||
srv.start()
|
||||
with self._lock:
|
||||
self._instances[name] = srv
|
||||
|
||||
# 停止多余的实例
|
||||
for name in running_names - current_names:
|
||||
srv = self._instances.pop(name)
|
||||
srv.stop()
|
||||
|
||||
# 更新实例运行状态到数据库
|
||||
with db.app.app_context():
|
||||
for inst in Instance.query.all():
|
||||
srv = self._instances.get(inst.name)
|
||||
inst.running = srv is not None and srv.running
|
||||
if srv:
|
||||
inst.active_connections = srv.active_connections
|
||||
else:
|
||||
inst.active_connections = 0
|
||||
db.session.commit()
|
||||
|
||||
def start_instance(self, instance_id):
|
||||
"""启动单个实例。"""
|
||||
from models import Instance
|
||||
from database import db
|
||||
with db.app.app_context():
|
||||
inst = Instance.query.get(instance_id)
|
||||
if not inst:
|
||||
return False
|
||||
cfg = InstanceConfig(
|
||||
instance_id=inst.id,
|
||||
name=inst.name,
|
||||
listen_host=inst.listen_host,
|
||||
listen_port=inst.listen_port,
|
||||
timeout=inst.timeout,
|
||||
auth_method=inst.auth_method,
|
||||
bandwidth_down=inst.bandwidth_down,
|
||||
bandwidth_up=inst.bandwidth_up,
|
||||
max_concurrent=inst.max_concurrent,
|
||||
)
|
||||
self._config_cache[inst.name] = cfg
|
||||
srv = Socks5Server(cfg, self.user_service)
|
||||
srv.start()
|
||||
with self._lock:
|
||||
self._instances[inst.name] = srv
|
||||
inst.running = True
|
||||
db.session.commit()
|
||||
return True
|
||||
|
||||
def stop_instance(self, instance_id):
|
||||
"""停止单个实例。"""
|
||||
from models import Instance
|
||||
from database import db
|
||||
with db.app.app_context():
|
||||
inst = Instance.query.get(instance_id)
|
||||
if not inst:
|
||||
return False
|
||||
srv = self._instances.pop(inst.name, None)
|
||||
if srv:
|
||||
srv.stop()
|
||||
del self._config_cache[inst.name]
|
||||
inst.running = False
|
||||
inst.active_connections = 0
|
||||
db.session.commit()
|
||||
return True
|
||||
|
||||
def restart_instance(self, instance_id):
|
||||
"""重启单个实例。"""
|
||||
self.stop_instance(instance_id)
|
||||
return self.start_instance(instance_id)
|
||||
|
||||
def get_status(self):
|
||||
"""获取所有实例状态。"""
|
||||
from models import Instance
|
||||
from database import db
|
||||
with db.app.app_context():
|
||||
result = []
|
||||
for inst in Instance.query.all():
|
||||
srv = self._instances.get(inst.name)
|
||||
result.append({
|
||||
"id": inst.id,
|
||||
"name": inst.name,
|
||||
"host": inst.listen_host,
|
||||
"port": inst.listen_port,
|
||||
"enabled": inst.enabled,
|
||||
"running": bool(srv and srv.running),
|
||||
"active_connections": srv.active_connections if srv else 0,
|
||||
"auth_method": inst.auth_method,
|
||||
"timeout": inst.timeout,
|
||||
"bandwidth_down": inst.bandwidth_down,
|
||||
"bandwidth_up": inst.bandwidth_up,
|
||||
"max_concurrent": inst.max_concurrent,
|
||||
"notes": inst.notes,
|
||||
})
|
||||
return result
|
||||
@@ -0,0 +1,419 @@
|
||||
"""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, on_close_cb):
|
||||
self.reader = reader
|
||||
self.writer = writer
|
||||
self.config = instance_config
|
||||
self.user_service = user_service
|
||||
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 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":
|
||||
await self.writer.write(bytes([SOCKS_VERSION, METHOD_USERPASS]))
|
||||
elif METHOD_NO_AUTH in methods:
|
||||
await self.writer.write(bytes([SOCKS_VERSION, METHOD_NO_AUTH]))
|
||||
else:
|
||||
await 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
|
||||
|
||||
# 读认证请求头
|
||||
header = await asyncio.wait_for(self.reader.read(256),
|
||||
timeout=self.config.timeout)
|
||||
if len(header) < 2:
|
||||
return False
|
||||
ver, ulen = struct.unpack("!BB", header[:2])
|
||||
if ver != 1:
|
||||
return False
|
||||
|
||||
# 读用户名
|
||||
username = await asyncio.wait_for(self.reader.read(ulen),
|
||||
timeout=self.config.timeout)
|
||||
if len(username) != ulen:
|
||||
return False
|
||||
username = username.decode("utf-8", errors="replace")
|
||||
|
||||
# 读密码长度
|
||||
plen_byte = await asyncio.wait_for(self.reader.read(1),
|
||||
timeout=self.config.timeout)
|
||||
if not plen_byte:
|
||||
return False
|
||||
plen = plen_byte[0]
|
||||
|
||||
# 读密码
|
||||
passwd = await asyncio.wait_for(self.reader.read(plen),
|
||||
timeout=self.config.timeout)
|
||||
if len(passwd) != plen:
|
||||
return False
|
||||
passwd = passwd.decode("utf-8", errors="replace")
|
||||
|
||||
# 验证
|
||||
user = self.user_service.get_user(username)
|
||||
if user is None or user.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, "密码错误或账号不可用")
|
||||
await 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)
|
||||
await 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 响应。"""
|
||||
addr_bytes = b"\x00\x00\x00\x00"
|
||||
if bnd_addr and "." not in str(bnd_addr):
|
||||
# IPv6
|
||||
try:
|
||||
addr_bytes = ipaddress.IPv6Address(str(bnd_addr)).packed
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
self.writer.write(struct.pack(
|
||||
"!BBBB5sH",
|
||||
SOCKS_VERSION, rep, 0x00, ATYP_IPV4,
|
||||
addr_bytes, 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:
|
||||
reader, writer = await asyncio.wait_for(
|
||||
asyncio.open_connection(addr, port),
|
||||
timeout=self.config.timeout
|
||||
)
|
||||
except (asyncio.TimeoutError, OSError) 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):
|
||||
"""双向隧道转发。"""
|
||||
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:
|
||||
pass
|
||||
except Exception as e:
|
||||
log.exception("[%s] 隧道异常: %s", self.conn_id, e)
|
||||
finally:
|
||||
self._close_streams(dst_reader, dst_writer)
|
||||
# 记录最终统计
|
||||
await self._record_stats()
|
||||
|
||||
async def _forward(self, src, dst, direction, speed_limit_mbps, user):
|
||||
"""单向转发，带限速。"""
|
||||
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))
|
||||
|
||||
await dst.write(data)
|
||||
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
|
||||
|
||||
def _close_streams(self, dst_reader, dst_writer):
|
||||
try:
|
||||
dst_writer.close()
|
||||
try:
|
||||
dst_reader.feed_eof()
|
||||
except Exception:
|
||||
pass
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
async def _record_connection(self, state="active"):
|
||||
"""记录连接到数据库（异步）。"""
|
||||
try:
|
||||
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:
|
||||
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
|
||||
Reference in New Issue
Block a user