update upd

edge_udp.py

@@ -11,6 +11,7 @@ from dataclasses import dataclass, field
 from typing import Dict
 
 from .config_udp import UdpConfig, UdpRelayNode
+from .logging_utils import log_print as print
 from .scheduler_udp import UdpScheduler
 from .protocol import AUTH, STATUS_ERR, STATUS_OK, UDP_RECV, UDP_SEND, Frame, read_frame, write_frame, encode_json
 
@@ -61,7 +62,10 @@ class RelayLink:
     async def _pump(self) -> None:
         try:
             while True:
-                frame = await read_frame(self.reader)
+                try:
+                    frame = await read_frame(self.reader)
+                except (asyncio.IncompleteReadError, ConnectionResetError, BrokenPipeError, OSError):
+                    break
                 if frame.kind == UDP_RECV and self.udp_server:
                     await self.udp_server.handle_from_relay(frame, self)
         finally:

logging_utils.py

@@ -0,0 +1,15 @@
+from __future__ import annotations
+
+from datetime import datetime
+import builtins
+from typing import Any
+
+
+def log_print(*args: Any, **kwargs: Any) -> None:
+    sep = kwargs.get("sep", " ")
+    end = kwargs.get("end", "\n")
+    file = kwargs.get("file", None)
+    flush = kwargs.get("flush", False)
+    message = sep.join(str(arg) for arg in args)
+    timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
+    builtins.print(f"{timestamp} {message}", end=end, file=file, flush=flush)

relay_client_tcp.py

@@ -10,6 +10,7 @@ import time
 from typing import Awaitable, Callable, Dict
 
 from .config_tcp import TcpConfig, TcpRelayNode
+from .logging_utils import log_print as print
 from .protocol import AUTH, PING, PONG, STATUS_OK, TCP_CLOSE, Frame, encode_json, read_frame, write_frame
 from .scheduler_tcp import TcpScheduler
 

relay_client_udp.py

@@ -10,6 +10,7 @@ import time
 from typing import Awaitable, Callable, Dict
 
 from .config_udp import UdpConfig, UdpRelayNode
+from .logging_utils import log_print as print
 from .protocol import AUTH, PING, PONG, STATUS_OK, TCP_CLOSE, Frame, encode_json, read_frame, write_frame
 from .scheduler_udp import UdpScheduler
 

relay_server_tcp.py

@@ -5,6 +5,7 @@ import contextlib
 from dataclasses import dataclass, field
 from typing import Dict
 
+from .logging_utils import log_print as print
 from .protocol import AUTH, PING, PONG, STATUS_ERR, STATUS_OK, TCP_CLOSE, TCP_DATA, TCP_OPEN, TCP_STATUS, Frame, decode_json, encode_json, read_frame, write_frame
 
 

relay_server_udp.py

@@ -4,6 +4,7 @@ import asyncio
 import contextlib
 from dataclasses import dataclass, field
 
+from .logging_utils import log_print as print
 from .protocol import AUTH, PING, PONG, STATUS_ERR, STATUS_OK, UDP_RECV, UDP_SEND, Frame, decode_json, encode_json, read_frame, write_frame
 
 
@@ -145,10 +146,19 @@ class UdpRelayServer:
     async def start(self, host: str, port: int) -> None:
         async def accept(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None:
             try:
+                peer = writer.get_extra_info("peername")
+                if peer:
+                    print(f"[relay] udp client connected peer={peer[0]}:{peer[1]}")
                 await UdpRelayChannel(reader, writer, self.token).run()
             except (asyncio.IncompleteReadError, ConnectionResetError, BrokenPipeError, OSError):
                 pass
+            finally:
+                with contextlib.suppress(Exception):
+                    peer = writer.get_extra_info("peername")
+                    if peer:
+                        print(f"[relay] udp client closed peer={peer[0]}:{peer[1]}")
 
+        print(f"[relay] udp server listening on {host}:{port}")
         server = await asyncio.start_server(accept, host, port)
         async with server:
             await server.serve_forever()