from __future__ import annotations import asyncio import contextlib import itertools import socket import struct from dataclasses import dataclass, field from typing import Dict from .config import Config, RelayNode from .protocol import AUTH, STATUS_OK, TCP_CLOSE, TCP_DATA, TCP_OPEN, TCP_STATUS, UDP_RECV, UDP_SEND, Frame, decode_json, encode_json, read_frame, write_frame from .scheduler import Scheduler SOCKS_VERSION = 5 async def read_exact(reader: asyncio.StreamReader, size: int) -> bytes: return await reader.readexactly(size) @dataclass(eq=False) class RelayLink: node: RelayNode reader: asyncio.StreamReader writer: asyncio.StreamWriter pump: asyncio.Task | None = None tcp_sessions: Dict[tuple[int, int], "TcpRaceSession"] = field(default_factory=dict) udp_server: "UdpAssociateServer | None" = None closed: bool = False async def start(self) -> None: await write_frame(self.writer, Frame(AUTH, 0, 0, 0, 0, encode_json({"token": self.node.token}))) frame = await read_frame(self.reader) if frame.kind != AUTH or frame.packet_id != STATUS_OK: raise ConnectionError(f"relay auth failed: {self.node.name}") self.pump = asyncio.create_task(self._pump()) async def _pump(self) -> None: try: while True: frame = await read_frame(self.reader) key = (frame.session_id, frame.stream_id) if frame.kind in (TCP_STATUS, TCP_DATA, TCP_CLOSE): session = self.tcp_sessions.get(key) if session: await session.handle_frame(self, frame) elif frame.kind == UDP_RECV and self.udp_server: await self.udp_server.handle_from_relay(frame, self) except asyncio.IncompleteReadError: pass finally: await self.close() async def send(self, frame: Frame) -> None: if self.closed: raise ConnectionError(f"relay closed: {self.node.name}") await write_frame(self.writer, frame) async def close(self) -> None: if self.closed: return self.closed = True self.writer.close() with contextlib.suppress(Exception): await self.writer.wait_closed() @dataclass class TcpRaceSession: session_id: int stream_id: int target_host: str target_port: int local_reader: asyncio.StreamReader local_writer: asyncio.StreamWriter links: list[RelayLink] warmup_bytes: int winning_link: RelayLink | None = None winner_name: str | None = None opened: int = 0 open_errors: list[str] = field(default_factory=list) uplink_bytes: int = 0 closed: bool = False open_event: asyncio.Event = field(default_factory=asyncio.Event) winner_event: asyncio.Event = field(default_factory=asyncio.Event) pump_task: asyncio.Task | None = None win_counts: Dict[str, int] = field(default_factory=dict) async def start(self) -> None: meta = encode_json({"host": self.target_host, "port": self.target_port}) for link in self.links: link.tcp_sessions[(self.session_id, self.stream_id)] = self await link.send(Frame(TCP_OPEN, self.session_id, self.stream_id, 0, 0, meta)) await asyncio.wait_for(self.open_event.wait(), timeout=10) if self.opened == 0: raise ConnectionError(self.open_errors[0] if self.open_errors else "all relays failed") self.pump_task = asyncio.create_task(self._pump_local()) async def _pump_local(self) -> None: try: while True: chunk = await self.local_reader.read(65536) if not chunk: break self.uplink_bytes += len(chunk) if self.winning_link is None and self.uplink_bytes <= self.warmup_bytes: await asyncio.gather(*(link.send(Frame(TCP_DATA, self.session_id, self.stream_id, 0, 0, chunk)) for link in self.links if not link.closed), return_exceptions=True) else: if self.winning_link is None: await self.winner_event.wait() if self.winning_link: await self.winning_link.send(Frame(TCP_DATA, self.session_id, self.stream_id, 0, 0, chunk)) except Exception: pass finally: await self.close() async def handle_frame(self, link: RelayLink, frame: Frame) -> None: if self.closed: return if frame.kind == TCP_STATUS: if frame.packet_id == STATUS_OK: self.opened += 1 else: self.open_errors.append(frame.payload.decode("utf-8", errors="replace")) if self.opened > 0 or len(self.open_errors) == len(self.links): self.open_event.set() return if frame.kind == TCP_DATA: if self.winning_link is None: self.winning_link = link self.winner_name = link.node.name self.win_counts[link.node.name] = self.win_counts.get(link.node.name, 0) + 1 node_total = self.win_counts[link.node.name] relay_detail = ", ".join(f"{name}={count}" for name, count in sorted(self.win_counts.items())) or "none" print(f"[edge] tcp win session={self.session_id} target={self.target_host}:{self.target_port} winner={link.node.name} node_total={node_total} win_breakdown={relay_detail}") self.winner_event.set() await self._close_losers(except_link=link) if link is self.winning_link: self.local_writer.write(frame.payload) await self.local_writer.drain() return if frame.kind == TCP_CLOSE: if self.winning_link is None: self.winning_link = link self.winner_event.set() if link is self.winning_link: await self.close() async def _close_losers(self, except_link: RelayLink) -> None: await asyncio.gather(*(link.send(Frame(TCP_CLOSE, self.session_id, self.stream_id, 0, 0, b"")) for link in self.links if link is not except_link and not link.closed), return_exceptions=True) async def close(self) -> None: if self.closed: return self.closed = True if self.pump_task and self.pump_task is not asyncio.current_task(): self.pump_task.cancel() with contextlib.suppress(Exception): await self.pump_task await asyncio.gather(*(link.send(Frame(TCP_CLOSE, self.session_id, self.stream_id, 0, 0, b"")) for link in self.links if not link.closed), return_exceptions=True) for link in self.links: link.tcp_sessions.pop((self.session_id, self.stream_id), None) self.local_writer.close() with contextlib.suppress(Exception): await self.local_writer.wait_closed() class UdpAssociateServer(asyncio.DatagramProtocol): def __init__(self, edge: "SocksEdge") -> None: self.edge = edge self.transport: asyncio.DatagramTransport | None = None self.client_addr = None self.packet_counter = itertools.count(1) self.pending: set[int] = set() def connection_made(self, transport) -> None: self.transport = transport def datagram_received(self, data: bytes, addr) -> None: if len(data) < 10: return if self.client_addr is None: self.client_addr = addr if addr != self.client_addr: return host, port, payload = self._parse_socks_udp(data) packet_id = next(self.packet_counter) self.pending.add(packet_id) asyncio.create_task(self.edge.forward_udp(host, port, payload, packet_id, self)) async def handle_from_relay(self, frame: Frame, _link: RelayLink) -> None: if frame.packet_id not in self.pending or self.transport is None or self.client_addr is None: return self.pending.discard(frame.packet_id) host = self.edge.udp_targets.get(frame.packet_id, ("0.0.0.0", 0))[0] port = self.edge.udp_targets.get(frame.packet_id, ("0.0.0.0", 0))[1] packet = self._build_socks_udp(host, port, frame.payload) self.transport.sendto(packet, self.client_addr) def _parse_socks_udp(self, packet: bytes) -> tuple[str, int, bytes]: atyp = packet[3] offset = 4 if atyp == 1: host = socket.inet_ntoa(packet[offset:offset + 4]) offset += 4 elif atyp == 3: size = packet[offset] offset += 1 host = packet[offset:offset + size].decode() offset += size else: raise ValueError("unsupported udp atyp") port = struct.unpack("!H", packet[offset:offset + 2])[0] offset += 2 return host, port, packet[offset:] def _build_socks_udp(self, host: str, port: int, payload: bytes) -> bytes: try: addr = socket.inet_aton(host) header = b"\x00\x00\x00\x01" + addr + struct.pack("!H", port) except OSError: raw = host.encode() header = b"\x00\x00\x00\x03" + bytes([len(raw)]) + raw + struct.pack("!H", port) return header + payload class SocksEdge: def __init__(self, listen_host: str, listen_port: int, config: Config) -> None: self.listen_host = listen_host self.listen_port = listen_port self.config = config self.scheduler = Scheduler(config) self.links: list[RelayLink] = [] self.session_ids = itertools.count(1) self.udp_targets: dict[int, tuple[str, int]] = {} self.udp_server: UdpAssociateServer | None = None async def start(self) -> None: await self.scheduler.start() await self._connect_relays() server = await asyncio.start_server(self._accept, self.listen_host, self.listen_port) sockets = ", ".join(str(sock.getsockname()) for sock in server.sockets or []) print(f"[edge] socks5 listening on {sockets}") async with server: await server.serve_forever() async def _connect_relays(self) -> None: for node in self.config.relays: reader, writer = await asyncio.open_connection(node.host, node.port) link = RelayLink(node, reader, writer) await link.start() self.links.append(link) loop = asyncio.get_running_loop() transport, protocol = await loop.create_datagram_endpoint(lambda: UdpAssociateServer(self), local_addr=(self.listen_host, 0)) self.udp_server = protocol for link in self.links: link.udp_server = protocol self.udp_transport = transport async def _accept(self, reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None: try: host, port, udp_mode = await self._handshake(reader, writer) if udp_mode: return links = self._selected_links() session = TcpRaceSession( session_id=next(self.session_ids), stream_id=0, target_host=host, target_port=port, local_reader=reader, local_writer=writer, links=links, warmup_bytes=self.config.tcp_warmup_bytes, ) await session.start() except Exception: writer.close() with contextlib.suppress(Exception): await writer.wait_closed() def _selected_links(self) -> list[RelayLink]: chosen = {node.name for node in self.scheduler.choose()} links = [link for link in self.links if link.node.name in chosen and not link.closed] return links or [link for link in self.links if not link.closed][:1] async def forward_udp(self, host: str, port: int, payload: bytes, packet_id: int, udp_server: UdpAssociateServer) -> None: self.udp_targets[packet_id] = (host, port) meta = encode_json({"host": host, "port": port}) links = self._selected_links() for index, link in enumerate(links): body = meta + payload if index == 0 else payload await link.send(Frame(UDP_SEND, 1, index, 0, packet_id if index == 0 else 0, body)) async def _handshake(self, reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> tuple[str, int, bool]: version, methods_len = (await read_exact(reader, 2)) if version != SOCKS_VERSION: raise ValueError("unsupported socks version") await read_exact(reader, methods_len) writer.write(b"\x05\x00") await writer.drain() version, command, _, atyp = await read_exact(reader, 4) if version != SOCKS_VERSION: raise ValueError("unsupported socks version") if atyp == 1: host = socket.inet_ntoa(await read_exact(reader, 4)) elif atyp == 3: size = (await read_exact(reader, 1))[0] host = (await read_exact(reader, size)).decode() else: raise ValueError("unsupported atyp") port = struct.unpack("!H", await read_exact(reader, 2))[0] if command == 1: writer.write(b"\x05\x00\x00\x01\x00\x00\x00\x00\x00\x00") await writer.drain() return host, port, False if command == 3 and self.udp_server and self.udp_server.transport: bind_host, bind_port = self.udp_server.transport.get_extra_info("sockname")[:2] writer.write(b"\x05\x00\x00\x01" + socket.inet_aton(bind_host) + struct.pack("!H", bind_port)) await writer.drain() return host, port, True raise ValueError("unsupported socks command")