diff --git a/engine/server.py b/engine/server.py index 9f6df53..dd9cd2a 100644 --- a/engine/server.py +++ b/engine/server.py @@ -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