AI-Agent/socks-manager
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 2 Wiki Activity

fix: SOCKS5 隧道修复 - 子线程事件循环兼容性

Browse Source 
1. 去掉 asyncio.wait_for 包装 open_connection，改用直接调用
2. ConnectionHandler 增加 flask_app 参数，数据库操作使用 app context
3. InstanceManager 传递 flask_app 到 Socks5Server
4. 修复 _reply 中 writer.write 不需要 await 的问题

验证:
- 直连测试: HTTP/1.1 200 OK (ifconfig.me)
- 端口监听: 10800 正常工作
This commit is contained in:
Your Name
2026-07-16 19:00:53 +08:00
parent 4fbb63e0b9
commit 5ab5366534
3 changed files with 41 additions and 36 deletions
Download Patch File Download Diff File Expand all files Collapse all files
app.py
+1 -1
View File
@@ -44,7 +44,7 @@ def create_app():
app.user_service = user_service
# 实例管理器
instance_manager = InstanceManager(user_service)
instance_manager = InstanceManager(user_service, flask_app=app)
app.instance_manager = instance_manager
# 备份目录
engine/instances.py
+7 -4
View File
@@ -31,9 +31,10 @@ class InstanceConfig:
class Socks5Server:
"""单个 SOCKS5 服务器的 asyncio 事件循环。"""
def __init__(self, config: InstanceConfig, user_service):
def __init__(self, config: InstanceConfig, user_service, flask_app=None):
self.config = config
self.user_service = user_service
self.flask_app = flask_app
self.server = None
self.loop = None
self.thread = None
@@ -94,6 +95,7 @@ class Socks5Server:
self._active_connections += 1
handler = ConnectionHandler(
reader, writer, self.config, self.user_service,
flask_app=self.flask_app,
on_close_cb=self._on_connection_close
)
await handler.handle()
@@ -113,8 +115,9 @@ import time # noqa: E402
class InstanceManager:
"""管理所有 SOCKS5 实例。"""
def __init__(self, user_service):
def __init__(self, user_service, flask_app=None):
self.user_service = user_service
self.flask_app = flask_app
self._instances = {} # name -> Socks5Server
self._config_cache = {} # name -> InstanceConfig
self._lock = threading.Lock()
@@ -151,7 +154,7 @@ class InstanceManager:
# 启动缺失的实例
for name in current_names - running_names:
cfg = self._config_cache[name]
srv = Socks5Server(cfg, self.user_service)
srv = Socks5Server(cfg, self.user_service, self.flask_app)
srv.start()
with self._lock:
self._instances[name] = srv
@@ -192,7 +195,7 @@ class InstanceManager:
max_concurrent=inst.max_concurrent,
)
self._config_cache[inst.name] = cfg
srv = Socks5Server(cfg, self.user_service)
srv = Socks5Server(cfg, self.user_service, self.flask_app)
srv.start()
with self._lock:
self._instances[inst.name] = srv
engine/server.py
+33 -31
View File
@@ -4,7 +4,6 @@ import logging
import struct
import time
import uuid
import socket as _socket
import ipaddress
from datetime import datetime, timezone
@@ -43,11 +42,13 @@ TUNNEL_CHUNK = 65536
class ConnectionHandler:
"""单个 SOCKS5 连接的处理器。"""
def __init__(self, reader, writer, instance_config, user_service, on_close_cb):
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
@@ -274,12 +275,9 @@ class ConnectionHandler:
# 连接目标
try:
# 注意：不用 asyncio.wait_for 包装 open_connection，
# 在子线程事件循环中 wait_for 会导致连接成功后读取被取消
sock = _socket.socket(_socket.AF_INET, _socket.SOCK_STREAM)
sock.settimeout(self.config.timeout)
sock.connect((addr, port))
reader, writer = await asyncio.open_connection(sock=sock)
# 直接用 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
@@ -372,19 +370,21 @@ class ConnectionHandler:
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()
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)
@@ -392,17 +392,19 @@ class ConnectionHandler:
"""记录最终连接统计。"""
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()
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)
# 更新用户流量
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)