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:
+16
-1
@@ -328,6 +328,7 @@ class ConnectionHandler:
|
|||||||
|
|
||||||
async def _tunnel(self, dst_reader, dst_writer):
|
async def _tunnel(self, dst_reader, dst_writer):
|
||||||
"""双向隧道转发。"""
|
"""双向隧道转发。"""
|
||||||
|
f1 = f2 = None
|
||||||
try:
|
try:
|
||||||
# 获取限速参数
|
# 获取限速参数
|
||||||
user = None
|
user = None
|
||||||
@@ -352,12 +353,22 @@ class ConnectionHandler:
|
|||||||
except Exception as e:
|
except Exception as e:
|
||||||
log.exception("[%s] 隧道异常: %s", self.conn_id, e)
|
log.exception("[%s] 隧道异常: %s", self.conn_id, e)
|
||||||
finally:
|
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)
|
self._close_streams(dst_reader, dst_writer)
|
||||||
# 记录最终统计
|
# 记录最终统计
|
||||||
await self._record_stats()
|
await self._record_stats()
|
||||||
|
|
||||||
async def _forward(self, src, dst, direction, speed_limit_mbps, user):
|
async def _forward(self, src, dst, direction, speed_limit_mbps, user):
|
||||||
"""单向转发,带限速。"""
|
"""单向转发,带限速。"""
|
||||||
|
if src is None or dst is None:
|
||||||
|
return
|
||||||
total_transferred = 0
|
total_transferred = 0
|
||||||
try:
|
try:
|
||||||
while True:
|
while True:
|
||||||
@@ -372,7 +383,7 @@ class ConnectionHandler:
|
|||||||
if sleep_time > 0.001: # 小于 1ms 不节流
|
if sleep_time > 0.001: # 小于 1ms 不节流
|
||||||
await asyncio.sleep(min(sleep_time, 1.0))
|
await asyncio.sleep(min(sleep_time, 1.0))
|
||||||
|
|
||||||
await dst.write(data)
|
dst.write(data) # write() 不是协程,不需要 await
|
||||||
await dst.drain()
|
await dst.drain()
|
||||||
total_transferred += len(data)
|
total_transferred += len(data)
|
||||||
|
|
||||||
@@ -385,10 +396,14 @@ class ConnectionHandler:
|
|||||||
pass
|
pass
|
||||||
except asyncio.CancelledError:
|
except asyncio.CancelledError:
|
||||||
pass
|
pass
|
||||||
|
except Exception as e:
|
||||||
|
log.debug("[%s] 转发异常 %s: %s", self.conn_id, direction, e)
|
||||||
|
|
||||||
def _close_streams(self, dst_reader, dst_writer):
|
def _close_streams(self, dst_reader, dst_writer):
|
||||||
try:
|
try:
|
||||||
|
if dst_writer is not None:
|
||||||
dst_writer.close()
|
dst_writer.close()
|
||||||
|
if dst_reader is not None:
|
||||||
try:
|
try:
|
||||||
dst_reader.feed_eof()
|
dst_reader.feed_eof()
|
||||||
except Exception:
|
except Exception:
|
||||||
|
|||||||
Reference in New Issue
Block a user