fix: 修复 asyncio 任务清理和 StreamWriter.write() 协程错误
问题1: StreamWriter.write() 不是协程,不需要 await,错误使用会导致 TypeError 问题2: 隧道关闭时未正确取消转发任务,导致 'Task was destroyed but it is pending!' 警告 问题3: 未对 None 的 reader/writer 做前置检查,导致空指针异常 修复内容: - 移除 dst.write(data) 前的 await - _tunnel() finally 块中主动取消转发任务并 await - _forward() 开头添加 None 检查 - _close_streams() 增加 None 检查 - 添加异常捕获日志
This commit is contained in:
+21
-6
@@ -328,6 +328,7 @@ class ConnectionHandler:
|
||||
|
||||
async def _tunnel(self, dst_reader, dst_writer):
|
||||
"""双向隧道转发。"""
|
||||
f1 = f2 = None
|
||||
try:
|
||||
# 获取限速参数
|
||||
user = None
|
||||
@@ -352,12 +353,22 @@ class ConnectionHandler:
|
||||
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:
|
||||
pass
|
||||
self._close_streams(dst_reader, dst_writer)
|
||||
# 记录最终统计
|
||||
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:
|
||||
@@ -372,7 +383,7 @@ class ConnectionHandler:
|
||||
if sleep_time > 0.001: # 小于 1ms 不节流
|
||||
await asyncio.sleep(min(sleep_time, 1.0))
|
||||
|
||||
await dst.write(data)
|
||||
dst.write(data) # write() 不是协程,不需要 await
|
||||
await dst.drain()
|
||||
total_transferred += len(data)
|
||||
|
||||
@@ -385,14 +396,18 @@ class ConnectionHandler:
|
||||
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:
|
||||
dst_writer.close()
|
||||
try:
|
||||
dst_reader.feed_eof()
|
||||
except Exception:
|
||||
pass
|
||||
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
|
||||
|
||||
|
||||
Reference in New Issue
Block a user