from __future__ import annotations import asyncio import contextlib from dataclasses import dataclass, field from typing import Dict from .protocol import AUTH, PING, PONG, STATUS_ERR, STATUS_OK, TCP_CLOSE, TCP_DATA, TCP_OPEN, TCP_STATUS, UDP_RECV, UDP_SEND, Frame, decode_json, read_frame, write_frame @dataclass class TcpSession: session_id: int stream_id: int writer: asyncio.StreamWriter task: asyncio.Task @dataclass class UdpSession: session_id: int stream_id: int transport: asyncio.DatagramTransport | None = None protocol: "RelayUdpProtocol | None" = None class RelayUdpProtocol(asyncio.DatagramProtocol): def __init__(self, channel: "RelayChannel", session_id: int, stream_id: int) -> None: self.channel = channel self.session_id = session_id self.stream_id = stream_id def datagram_received(self, data: bytes, _addr) -> None: if self.channel.closed: return asyncio.create_task(self.channel.safe_send(Frame(UDP_RECV, self.session_id, self.stream_id, 0, 0, data))) @dataclass class RelayChannel: reader: asyncio.StreamReader writer: asyncio.StreamWriter token: str tcp_sessions: Dict[tuple[int, int], TcpSession] = field(default_factory=dict) udp_sessions: Dict[tuple[int, int], UdpSession] = field(default_factory=dict) closed: bool = False send_lock: asyncio.Lock = field(default_factory=asyncio.Lock) async def run(self) -> None: peer = self.writer.get_extra_info("peername") authed = False try: auth = await read_frame(self.reader) if auth.kind != AUTH or decode_json(auth.payload).get("token") != self.token: raise PermissionError("invalid token") authed = True print(f"[relay] auth ok peer={peer}") await self.safe_send(Frame(AUTH, 0, 0, 0, STATUS_OK, b"ok")) while True: frame = await read_frame(self.reader) await self.handle(frame) except asyncio.IncompleteReadError: if authed: print(f"[relay] disconnected peer={peer}") except asyncio.CancelledError: pass except Exception as exc: if authed: print(f"[relay] channel error peer={peer} error={exc!r}") finally: await self.close() async def safe_send(self, frame: Frame) -> bool: if self.closed: return False try: async with self.send_lock: if self.closed: return False await write_frame(self.writer, frame) return True except (BrokenPipeError, ConnectionResetError, RuntimeError, OSError, asyncio.CancelledError): return False async def handle(self, frame: Frame) -> None: key = (frame.session_id, frame.stream_id) if frame.kind == PING: await self.safe_send(Frame(PONG, 0, 0, frame.seq, 0, b"pong")) return if frame.kind == TCP_OPEN: meta = decode_json(frame.payload) family = int(meta.get("family", 0)) or 0 try: reader, writer = await asyncio.open_connection(meta["host"], int(meta["port"]), family=family or 0) task = asyncio.create_task(self._tcp_pump(frame.session_id, frame.stream_id, reader)) self.tcp_sessions[key] = TcpSession(frame.session_id, frame.stream_id, writer, task) await self.safe_send(Frame(TCP_STATUS, frame.session_id, frame.stream_id, 0, STATUS_OK, b"ok")) except Exception as exc: await self.safe_send(Frame(TCP_STATUS, frame.session_id, frame.stream_id, 0, STATUS_ERR, str(exc).encode())) return if frame.kind == TCP_DATA: session = self.tcp_sessions.get(key) if session: try: session.writer.write(frame.payload) await session.writer.drain() except Exception: await self._close_tcp(key) return if frame.kind == TCP_CLOSE: await self._close_tcp(key) return if frame.kind == UDP_SEND: meta = decode_json(frame.payload[: frame.packet_id]) payload = frame.payload[frame.packet_id :] session = self.udp_sessions.get(key) if session is None: family = int(meta.get("family", 0)) or 0 transport, protocol = await asyncio.get_running_loop().create_datagram_endpoint( lambda: RelayUdpProtocol(self, frame.session_id, frame.stream_id), remote_addr=(meta["host"], int(meta["port"])), family=family or 0, ) session = UdpSession(frame.session_id, frame.stream_id, transport, protocol) self.udp_sessions[key] = session with contextlib.suppress(Exception): session.transport.sendto(payload) return async def _tcp_pump(self, session_id: int, stream_id: int, reader: asyncio.StreamReader) -> None: try: while True: chunk = await reader.read(65536) if not chunk: break sent = await self.safe_send(Frame(TCP_DATA, session_id, stream_id, 0, 0, chunk)) if not sent: break except asyncio.CancelledError: pass except Exception: pass finally: if not self.closed: await self.safe_send(Frame(TCP_CLOSE, session_id, stream_id, 0, 0, b"")) await self._close_tcp((session_id, stream_id), from_task=True) async def _close_tcp(self, key: tuple[int, int], from_task: bool = False) -> None: session = self.tcp_sessions.pop(key, None) if session is None: return if not from_task and session.task is not asyncio.current_task(): session.task.cancel() with contextlib.suppress(Exception): await session.task session.writer.close() with contextlib.suppress(Exception): await session.writer.wait_closed() async def close(self) -> None: if self.closed: return self.closed = True for key in list(self.tcp_sessions): await self._close_tcp(key) for session in self.udp_sessions.values(): if session.transport: session.transport.close() self.udp_sessions.clear() self.writer.close() with contextlib.suppress(Exception): await self.writer.wait_closed() class RelayServer: def __init__(self, token: str) -> None: self.token = token async def start(self, host: str, port: int) -> None: server = await asyncio.start_server(self._accept, host, port) sockets = ", ".join(str(sock.getsockname()) for sock in server.sockets or []) print(f"[relay] listening on {sockets}") async with server: await server.serve_forever() async def _accept(self, reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None: await RelayChannel(reader, writer, self.token).run()