Gogs il y a 2 semaines
Parent
commit
ca8e5c319f
2 fichiers modifiés avec 13 ajouts et 1 suppressions
  1. 8 0
      relay_server.py
  2. 5 1
      socks_edge.py

+ 8 - 0
relay_server.py

@@ -2,6 +2,7 @@ from __future__ import annotations
 
 import asyncio
 import contextlib
+import json
 import time
 from dataclasses import dataclass, field
 from typing import Dict
@@ -58,6 +59,9 @@ class RelayChannel:
             return False
         return lived >= 15 or self.frame_count > 20
 
+    def _should_log_auth(self) -> bool:
+        return self.authed_kind not in {"probe", "normal"}
+
     async def run(self) -> None:
         peer = self.writer.get_extra_info("peername")
         authed = False
@@ -67,6 +71,8 @@ class RelayChannel:
                 return
             try:
                 payload = decode_json(auth.payload) if auth.payload else {}
+            except json.JSONDecodeError:
+                return
             except Exception:
                 return
             if payload.get("token") != self.token:
@@ -74,6 +80,8 @@ class RelayChannel:
             authed = True
             self.authed_at = time.monotonic()
             self.authed_kind = payload.get("purpose", "normal")
+            if self._should_log_auth():
+                print(f"[relay] auth ok peer={peer} kind={self.authed_kind}")
             ack_payload = {"status": "ok", "kind": self.authed_kind}
             await self.safe_send(Frame(AUTH, 0, 0, 0, STATUS_OK, encode_json(ack_payload)))
             while True:

+ 5 - 1
socks_edge.py

@@ -27,6 +27,7 @@ class RelayLink:
     pump: asyncio.Task | None = None
     closed_event: asyncio.Event = field(default_factory=asyncio.Event)
     maintain_task: asyncio.Task | None = None
+    send_lock: asyncio.Lock = field(default_factory=asyncio.Lock)
     tcp_sessions: Dict[tuple[int, int], "TcpRaceSession"] = field(default_factory=dict)
     udp_server: "UdpAssociateServer | None" = None
     closed: bool = False
@@ -62,7 +63,10 @@ class RelayLink:
         if self.closed:
             raise ConnectionError(f"relay closed: {self.node.name}")
         try:
-            await write_frame(self.writer, frame)
+            async with self.send_lock:
+                if self.closed:
+                    raise ConnectionError(f"relay closed: {self.node.name}")
+                await write_frame(self.writer, frame)
         except (BrokenPipeError, ConnectionResetError, RuntimeError, OSError, asyncio.CancelledError) as exc:
             await self.close()
             raise ConnectionError(f"relay closed: {self.node.name}") from exc