socks_edge.py 27 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599
  1. from __future__ import annotations
  2. import asyncio
  3. import contextlib
  4. import itertools
  5. import socket
  6. import struct
  7. from dataclasses import dataclass, field
  8. from typing import Dict
  9. from .config import Config, RelayNode
  10. 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
  11. from .scheduler import Scheduler
  12. SOCKS_VERSION = 5
  13. async def read_exact(reader: asyncio.StreamReader, size: int) -> bytes:
  14. return await reader.readexactly(size)
  15. @dataclass(eq=False)
  16. class RelayLink:
  17. node: RelayNode
  18. reader: asyncio.StreamReader
  19. writer: asyncio.StreamWriter
  20. pump: asyncio.Task | None = None
  21. closed_event: asyncio.Event = field(default_factory=asyncio.Event)
  22. maintain_task: asyncio.Task | None = None
  23. tcp_sessions: Dict[tuple[int, int], "TcpRaceSession"] = field(default_factory=dict)
  24. udp_server: "UdpAssociateServer | None" = None
  25. closed: bool = False
  26. async def start(self) -> None:
  27. await write_frame(self.writer, Frame(AUTH, 0, 0, 0, 0, encode_json({"token": self.node.token})))
  28. frame = await read_frame(self.reader)
  29. if frame.kind != AUTH or frame.packet_id != STATUS_OK:
  30. raise ConnectionError(f"relay auth failed: {self.node.name}")
  31. self.closed = False
  32. self.closed_event.clear()
  33. self.pump = asyncio.create_task(self._pump())
  34. async def _pump(self) -> None:
  35. try:
  36. while True:
  37. frame = await read_frame(self.reader)
  38. key = (frame.session_id, frame.stream_id)
  39. if frame.kind in (TCP_STATUS, TCP_DATA, TCP_CLOSE):
  40. session = self.tcp_sessions.get(key)
  41. if session:
  42. await session.handle_frame(self, frame)
  43. elif frame.kind == UDP_RECV and self.udp_server:
  44. await self.udp_server.handle_from_relay(frame, self)
  45. except (asyncio.IncompleteReadError, ConnectionResetError, BrokenPipeError, OSError):
  46. pass
  47. except Exception:
  48. pass
  49. finally:
  50. await self.close()
  51. async def send(self, frame: Frame) -> None:
  52. if self.closed:
  53. raise ConnectionError(f"relay closed: {self.node.name}")
  54. await write_frame(self.writer, frame)
  55. async def close(self) -> None:
  56. if self.closed:
  57. return
  58. self.closed = True
  59. self.closed_event.set()
  60. if self.pump and self.pump is not asyncio.current_task():
  61. self.pump.cancel()
  62. with contextlib.suppress(Exception):
  63. await self.pump
  64. self.writer.close()
  65. with contextlib.suppress(Exception):
  66. await self.writer.wait_closed()
  67. @dataclass
  68. class UdpFlowState:
  69. flow_id: int
  70. client_addr: tuple[str, int]
  71. target_host: str
  72. target_port: int
  73. created_at: float
  74. last_activity: float
  75. packets_sent: int = 0
  76. packets_received: int = 0
  77. duplicate_responses: int = 0
  78. winner_name: str | None = None
  79. candidate_names: tuple[str, ...] = ()
  80. link_streams: dict[str, int] = field(default_factory=dict)
  81. initialized_links: set[str] = field(default_factory=set)
  82. direct_sockets: dict[str, socket.socket] = field(default_factory=dict)
  83. direct_tasks: dict[str, asyncio.Task] = field(default_factory=dict)
  84. direct_failures: set[str] = field(default_factory=set)
  85. relay_failures: dict[str, int] = field(default_factory=dict)
  86. relay_error_seen: set[str] = field(default_factory=set)
  87. def touch(self, now: float) -> None:
  88. self.last_activity = now
  89. @dataclass
  90. class TcpRaceSession:
  91. session_id: int
  92. stream_id: int
  93. target_host: str
  94. target_port: int
  95. local_reader: asyncio.StreamReader
  96. local_writer: asyncio.StreamWriter
  97. links: list[RelayLink]
  98. warmup_bytes: int
  99. winning_link: RelayLink | None = None
  100. winner_name: str | None = None
  101. opened: int = 0
  102. open_errors: list[str] = field(default_factory=list)
  103. uplink_bytes: int = 0
  104. closed: bool = False
  105. open_event: asyncio.Event = field(default_factory=asyncio.Event)
  106. winner_event: asyncio.Event = field(default_factory=asyncio.Event)
  107. pump_task: asyncio.Task | None = None
  108. win_counts: Dict[str, int] = field(default_factory=dict)
  109. async def start(self) -> None:
  110. meta = encode_json({"host": self.target_host, "port": self.target_port})
  111. for link in self.links:
  112. link.tcp_sessions[(self.session_id, self.stream_id)] = self
  113. await link.send(Frame(TCP_OPEN, self.session_id, self.stream_id, 0, 0, meta))
  114. await asyncio.wait_for(self.open_event.wait(), timeout=10)
  115. if self.opened == 0:
  116. raise ConnectionError(self.open_errors[0] if self.open_errors else "all relays failed")
  117. self.pump_task = asyncio.create_task(self._pump_local())
  118. async def _pump_local(self) -> None:
  119. try:
  120. while True:
  121. chunk = await self.local_reader.read(65536)
  122. if not chunk:
  123. break
  124. self.uplink_bytes += len(chunk)
  125. if self.winning_link is None and self.uplink_bytes <= self.warmup_bytes:
  126. 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)
  127. else:
  128. if self.winning_link is None:
  129. await self.winner_event.wait()
  130. if self.winning_link:
  131. await self.winning_link.send(Frame(TCP_DATA, self.session_id, self.stream_id, 0, 0, chunk))
  132. except Exception:
  133. pass
  134. finally:
  135. await self.close()
  136. async def handle_frame(self, link: RelayLink, frame: Frame) -> None:
  137. if self.closed:
  138. return
  139. if frame.kind == TCP_STATUS:
  140. if frame.packet_id == STATUS_OK:
  141. self.opened += 1
  142. else:
  143. self.open_errors.append(frame.payload.decode("utf-8", errors="replace"))
  144. if self.opened > 0 or len(self.open_errors) == len(self.links):
  145. self.open_event.set()
  146. return
  147. if frame.kind == TCP_DATA:
  148. if self.winning_link is None:
  149. self.winning_link = link
  150. self.winner_name = link.node.name
  151. self.win_counts[link.node.name] = self.win_counts.get(link.node.name, 0) + 1
  152. node_total = self.win_counts[link.node.name]
  153. relay_detail = ", ".join(f"{name}={count}" for name, count in sorted(self.win_counts.items())) or "none"
  154. 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}")
  155. self.winner_event.set()
  156. await self._close_losers(except_link=link)
  157. if link is self.winning_link:
  158. self.local_writer.write(frame.payload)
  159. await self.local_writer.drain()
  160. return
  161. if frame.kind == TCP_CLOSE:
  162. if self.winning_link is None:
  163. self.winning_link = link
  164. self.winner_event.set()
  165. if link is self.winning_link:
  166. await self.close()
  167. async def _close_losers(self, except_link: RelayLink) -> None:
  168. 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)
  169. async def close(self) -> None:
  170. if self.closed:
  171. return
  172. self.closed = True
  173. if self.pump_task and self.pump_task is not asyncio.current_task():
  174. self.pump_task.cancel()
  175. with contextlib.suppress(Exception):
  176. await self.pump_task
  177. 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)
  178. for link in self.links:
  179. link.tcp_sessions.pop((self.session_id, self.stream_id), None)
  180. self.local_writer.close()
  181. with contextlib.suppress(Exception):
  182. await self.local_writer.wait_closed()
  183. class UdpAssociateServer(asyncio.DatagramProtocol):
  184. def __init__(self, edge: "SocksEdge") -> None:
  185. self.edge = edge
  186. self.transport: asyncio.DatagramTransport | None = None
  187. self.client_addr = None
  188. self.associate_peer = None
  189. self.packet_counter = itertools.count(1)
  190. self.client_flows: dict[tuple[tuple[str, int], str, int], UdpFlowState] = {}
  191. self.flow_counter = itertools.count(1)
  192. self.last_summary_at = 0.0
  193. self.win_counts: Dict[str, int] = {}
  194. self.relay_error_counts: Dict[str, int] = {}
  195. def connection_made(self, transport) -> None:
  196. self.transport = transport
  197. def register_associate(self, peer) -> None:
  198. peer_text = f"{peer[0]}:{peer[1]}" if isinstance(peer, tuple) and len(peer) >= 2 else str(peer)
  199. if self.associate_peer != peer_text:
  200. print(f"[edge] udp associate peer={peer_text}")
  201. self.associate_peer = peer_text
  202. def datagram_received(self, data: bytes, addr) -> None:
  203. if len(data) < 10:
  204. return
  205. if self.client_addr is None:
  206. self.client_addr = addr
  207. print(f"[edge] udp client bound addr={addr[0]}:{addr[1]}")
  208. elif addr != self.client_addr:
  209. print(f"[edge] udp client rebound old={self.client_addr[0]}:{self.client_addr[1]} new={addr[0]}:{addr[1]}")
  210. self._reset_client_state(addr)
  211. host, port, payload = self._parse_socks_udp(data)
  212. loop = asyncio.get_running_loop()
  213. now = loop.time()
  214. flow_key = ((addr[0], addr[1]), host, port)
  215. flow = self.client_flows.get(flow_key)
  216. if flow is None:
  217. flow = UdpFlowState(
  218. flow_id=next(self.flow_counter),
  219. client_addr=(addr[0], addr[1]),
  220. target_host=host,
  221. target_port=port,
  222. created_at=now,
  223. last_activity=now,
  224. )
  225. self.client_flows[flow_key] = flow
  226. flow.touch(now)
  227. flow.packets_sent += 1
  228. packet_id = next(self.packet_counter)
  229. asyncio.create_task(self.edge.forward_udp(flow, payload, packet_id, self))
  230. self._log_udp_summary()
  231. def _reset_client_state(self, addr) -> None:
  232. for flow in list(self.client_flows.values()):
  233. for task in list(flow.direct_tasks.values()):
  234. task.cancel()
  235. for sock in list(flow.direct_sockets.values()):
  236. with contextlib.suppress(Exception):
  237. sock.close()
  238. for stream_id in list(flow.link_streams.values()):
  239. self.edge.udp_flow_sessions.pop((flow.flow_id, stream_id), None)
  240. self.client_flows.clear()
  241. self.client_addr = addr
  242. self.win_counts.clear()
  243. print(f"[edge] udp client bound addr={addr[0]}:{addr[1]}")
  244. async def handle_from_relay(self, frame: Frame, link: RelayLink) -> None:
  245. if self.transport is None or self.client_addr is None:
  246. return
  247. flow = self.edge.udp_flow_sessions.get((frame.session_id, frame.stream_id))
  248. if flow is None:
  249. return
  250. await self._deliver_flow_packet(flow, frame.packet_id, frame.payload, link.node.name)
  251. async def handle_from_direct(self, flow: UdpFlowState, path_name: str, payload: bytes) -> None:
  252. if self.transport is None or self.client_addr is None:
  253. return
  254. await self._deliver_flow_packet(flow, 0, payload, path_name)
  255. async def _deliver_flow_packet(self, flow: UdpFlowState, packet_id: int, payload: bytes, source_name: str) -> None:
  256. if self.transport is None or self.client_addr is None:
  257. return
  258. packet = self._build_socks_udp(flow.target_host, flow.target_port, payload)
  259. now = asyncio.get_running_loop().time()
  260. flow.touch(now)
  261. flow.packets_received += 1
  262. if flow.winner_name is None:
  263. flow.winner_name = source_name
  264. self.win_counts[source_name] = self.win_counts.get(source_name, 0) + 1
  265. self._log_udp_summary(force=True)
  266. elif flow.winner_name != source_name:
  267. flow.duplicate_responses += 1
  268. if flow.winner_name == source_name:
  269. self.transport.sendto(packet, self.client_addr)
  270. self._log_udp_summary()
  271. def set_flow_candidates(self, flow: UdpFlowState, candidate_names: tuple[str, ...]) -> None:
  272. if not flow.candidate_names:
  273. flow.candidate_names = candidate_names
  274. def note_unsent(self, flow: UdpFlowState, packet_id: int) -> None:
  275. flow.touch(asyncio.get_running_loop().time())
  276. flow.relay_failures["unsent"] = flow.relay_failures.get("unsent", 0) + 1
  277. self._log_udp_summary(force=True)
  278. def _log_udp_summary(self, force: bool = False) -> None:
  279. now = asyncio.get_running_loop().time()
  280. if not force and now - self.last_summary_at < 10:
  281. return
  282. self.last_summary_at = now
  283. active_flows = len(self.client_flows)
  284. winners = sum(1 for flow in self.client_flows.values() if flow.winner_name)
  285. packets_sent = sum(flow.packets_sent for flow in self.client_flows.values())
  286. packets_received = sum(flow.packets_received for flow in self.client_flows.values())
  287. duplicates = sum(flow.duplicate_responses for flow in self.client_flows.values())
  288. direct_paths = sum(len(flow.direct_sockets) for flow in self.client_flows.values())
  289. relay_candidates = sum(len(flow.link_streams) for flow in self.client_flows.values())
  290. winner_detail = ", ".join(f"{flow.flow_id}:{flow.winner_name}" for flow in self.client_flows.values() if flow.winner_name) or "none"
  291. relay_errors: list[str] = []
  292. for flow in self.client_flows.values():
  293. for name, count in flow.relay_failures.items():
  294. relay_errors.append(f"{name}={count}")
  295. relay_error_detail = ", ".join(sorted(relay_errors)) or "none"
  296. if self.client_addr:
  297. print(
  298. f"[edge] udp summary bind={self.client_addr[0]}:{self.client_addr[1]} active_flows={active_flows} "
  299. f"winner_flows={winners} winner_detail={winner_detail} packets_sent={packets_sent} packets_received={packets_received} dup={duplicates} "
  300. f"direct_paths={direct_paths} relay_paths={relay_candidates} relay_errors={relay_error_detail}"
  301. )
  302. else:
  303. print(
  304. f"[edge] udp summary bind=unbound active_flows={active_flows} winner_flows={winners} winner_detail={winner_detail} packets_sent={packets_sent} packets_received={packets_received} dup={duplicates} "
  305. f"direct_paths={direct_paths} relay_paths={relay_candidates} relay_errors={relay_error_detail}"
  306. )
  307. def _parse_socks_udp(self, packet: bytes) -> tuple[str, int, bytes]:
  308. atyp = packet[3]
  309. offset = 4
  310. if atyp == 1:
  311. host = socket.inet_ntoa(packet[offset:offset + 4])
  312. offset += 4
  313. elif atyp == 3:
  314. size = packet[offset]
  315. offset += 1
  316. host = packet[offset:offset + size].decode()
  317. offset += size
  318. else:
  319. raise ValueError("unsupported udp atyp")
  320. port = struct.unpack("!H", packet[offset:offset + 2])[0]
  321. offset += 2
  322. return host, port, packet[offset:]
  323. def _build_socks_udp(self, host: str, port: int, payload: bytes) -> bytes:
  324. try:
  325. addr = socket.inet_aton(host)
  326. header = b"\x00\x00\x00\x01" + addr + struct.pack("!H", port)
  327. except OSError:
  328. raw = host.encode()
  329. header = b"\x00\x00\x00\x03" + bytes([len(raw)]) + raw + struct.pack("!H", port)
  330. return header + payload
  331. class SocksEdge:
  332. def __init__(self, listen_host: str, listen_port: int, config: Config) -> None:
  333. self.listen_host = listen_host
  334. self.listen_port = listen_port
  335. self.config = config
  336. self.scheduler = Scheduler(config)
  337. self.links: list[RelayLink] = []
  338. self.session_ids = itertools.count(1)
  339. self.udp_stream_ids = itertools.count(1)
  340. self.udp_flow_sessions: dict[tuple[int, int], UdpFlowState] = {}
  341. self.udp_server: UdpAssociateServer | None = None
  342. async def start(self) -> None:
  343. await self.scheduler.start()
  344. await self._connect_relays()
  345. server = await asyncio.start_server(self._accept, self.listen_host, self.listen_port)
  346. sockets = ", ".join(str(sock.getsockname()) for sock in server.sockets or [])
  347. print(f"[edge] socks5 listening on {sockets}")
  348. async with server:
  349. await server.serve_forever()
  350. async def _connect_relays(self) -> None:
  351. loop = asyncio.get_running_loop()
  352. transport, protocol = await loop.create_datagram_endpoint(lambda: UdpAssociateServer(self), local_addr=(self.listen_host, 0))
  353. self.udp_server = protocol
  354. self.udp_transport = transport
  355. for node in self.config.relays:
  356. link = RelayLink(node=node, reader=None, writer=None) # type: ignore[arg-type]
  357. link.udp_server = protocol
  358. self.links.append(link)
  359. link.maintain_task = asyncio.create_task(self._maintain_link(link))
  360. async def _maintain_link(self, link: RelayLink) -> None:
  361. backoff = 1.0
  362. while True:
  363. try:
  364. reader, writer = await asyncio.open_connection(link.node.host, link.node.port)
  365. sock = writer.get_extra_info("socket")
  366. if sock is not None:
  367. with contextlib.suppress(OSError):
  368. sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
  369. link.reader = reader
  370. link.writer = writer
  371. await link.start()
  372. backoff = 1.0
  373. await link.closed_event.wait()
  374. except asyncio.CancelledError:
  375. raise
  376. except Exception:
  377. await asyncio.sleep(backoff)
  378. backoff = min(10.0, backoff * 2)
  379. async def _accept(self, reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None:
  380. try:
  381. peer = writer.get_extra_info("peername")
  382. host, port, udp_mode = await self._handshake(reader, writer, peer)
  383. if udp_mode:
  384. return
  385. links = self._selected_links()
  386. session = TcpRaceSession(
  387. session_id=next(self.session_ids),
  388. stream_id=0,
  389. target_host=host,
  390. target_port=port,
  391. local_reader=reader,
  392. local_writer=writer,
  393. links=links,
  394. warmup_bytes=self.config.tcp_warmup_bytes,
  395. )
  396. await session.start()
  397. except Exception:
  398. writer.close()
  399. with contextlib.suppress(Exception):
  400. await writer.wait_closed()
  401. def _selected_links(self) -> list[RelayLink]:
  402. chosen = {node.name for node in self.scheduler.choose()}
  403. links = [link for link in self.links if link.node.name in chosen and not link.closed]
  404. return links or [link for link in self.links if not link.closed][:1]
  405. def _selected_udp_links(self) -> list[RelayLink]:
  406. online = [link for link in self.links if not link.closed and link.writer is not None]
  407. if not online:
  408. return []
  409. ordered = sorted(online, key=lambda link: self.scheduler.scores.get(link.node.name).score if link.node.name in self.scheduler.scores else 999999.0)
  410. return ordered
  411. def _udp_direct_redundancy_for_target(self, target_host: str) -> int:
  412. base = self.config.udp_direct_redundancy
  413. if ":" in target_host and self.config.udp_direct_redundancy_v6 is not None:
  414. base = self.config.udp_direct_redundancy_v6
  415. elif ":" not in target_host and self.config.udp_direct_redundancy_v4 is not None:
  416. base = self.config.udp_direct_redundancy_v4
  417. return max(1, base)
  418. async def _ensure_udp_direct_paths(self, flow: UdpFlowState, udp_server: UdpAssociateServer) -> None:
  419. target_count = self._udp_direct_redundancy_for_target(flow.target_host)
  420. for index in range(target_count):
  421. name = f"direct-{index + 1}" if target_count > 1 else "direct"
  422. if name in flow.direct_sockets or name in flow.direct_failures:
  423. continue
  424. try:
  425. family = socket.AF_INET6 if ":" in flow.target_host else socket.AF_INET
  426. sock = socket.socket(family, socket.SOCK_DGRAM)
  427. sock.setblocking(False)
  428. await asyncio.get_running_loop().sock_connect(sock, (flow.target_host, flow.target_port))
  429. flow.direct_sockets[name] = sock
  430. flow.direct_tasks[name] = asyncio.create_task(self._pump_udp_direct(flow, name, sock, udp_server))
  431. except Exception as exc:
  432. flow.direct_failures.add(name)
  433. print(f"[edge] udp direct open error flow={flow.flow_id} path={name} target={flow.target_host}:{flow.target_port} error={exc!r}")
  434. async def _pump_udp_direct(self, flow: UdpFlowState, path_name: str, sock: socket.socket, udp_server: UdpAssociateServer) -> None:
  435. loop = asyncio.get_running_loop()
  436. try:
  437. while True:
  438. data = await loop.sock_recv(sock, 65535)
  439. if not data:
  440. break
  441. await udp_server.handle_from_direct(flow, path_name, data)
  442. except Exception:
  443. pass
  444. finally:
  445. flow.direct_tasks.pop(path_name, None)
  446. flow.direct_sockets.pop(path_name, None)
  447. with contextlib.suppress(Exception):
  448. sock.close()
  449. async def forward_udp(self, flow: UdpFlowState, payload: bytes, packet_id: int, udp_server: UdpAssociateServer) -> None:
  450. await self._ensure_udp_direct_paths(flow, udp_server)
  451. meta = encode_json({"host": flow.target_host, "port": flow.target_port})
  452. links = self._selected_udp_links()
  453. direct_names = tuple(name for name in sorted(flow.direct_sockets))
  454. relay_names = tuple(link.node.name for link in links)
  455. candidate_names = direct_names + relay_names
  456. udp_server.set_flow_candidates(flow, candidate_names)
  457. if not candidate_names:
  458. udp_server.note_unsent(flow, packet_id)
  459. return
  460. active_direct_names = list(direct_names)
  461. active_links = links
  462. if not (self.config.udp_always_broadcast or flow.winner_name is None):
  463. active_direct_names = [name for name in active_direct_names if name == flow.winner_name]
  464. active_links = [link for link in active_links if link.node.name == flow.winner_name]
  465. if not active_direct_names and not active_links:
  466. if direct_names:
  467. active_direct_names = [direct_names[0]]
  468. elif links:
  469. active_links = links[:1]
  470. copies = max(1, self.config.udp_redundancy + 1)
  471. sent_any = False
  472. for attempt in range(copies):
  473. for path_name in active_direct_names:
  474. sock = flow.direct_sockets.get(path_name)
  475. if sock is None:
  476. continue
  477. try:
  478. await asyncio.get_running_loop().sock_sendall(sock, payload)
  479. sent_any = True
  480. except Exception as exc:
  481. flow.direct_failures.add(path_name)
  482. flow.direct_sockets.pop(path_name, None)
  483. task = flow.direct_tasks.pop(path_name, None)
  484. if task is not None:
  485. task.cancel()
  486. with contextlib.suppress(Exception):
  487. sock.close()
  488. flow.relay_failures[path_name] = flow.relay_failures.get(path_name, 0) + 1
  489. if path_name not in flow.relay_error_seen:
  490. flow.relay_error_seen.add(path_name)
  491. print(
  492. f"[edge] udp relay error flow={flow.flow_id} relay={path_name} error={exc!r}"
  493. )
  494. for link in active_links:
  495. stream_id = flow.link_streams.get(link.node.name)
  496. if stream_id is None:
  497. stream_id = next(self.udp_stream_ids)
  498. flow.link_streams[link.node.name] = stream_id
  499. self.udp_flow_sessions[(flow.flow_id, stream_id)] = flow
  500. include_meta = link.node.name not in flow.initialized_links
  501. body = (meta + payload) if include_meta else payload
  502. meta_len = len(meta) if include_meta else 0
  503. try:
  504. await link.send(Frame(UDP_SEND, flow.flow_id, stream_id, 0, meta_len, body))
  505. flow.initialized_links.add(link.node.name)
  506. sent_any = True
  507. except Exception as exc:
  508. flow.link_streams.pop(link.node.name, None)
  509. self.udp_flow_sessions.pop((flow.flow_id, stream_id), None)
  510. flow.relay_failures[link.node.name] = flow.relay_failures.get(link.node.name, 0) + 1
  511. if link.node.name not in flow.relay_error_seen:
  512. flow.relay_error_seen.add(link.node.name)
  513. print(
  514. f"[edge] udp relay error flow={flow.flow_id} relay={link.node.name} error={exc!r}"
  515. )
  516. if attempt + 1 < copies and self.config.udp_copy_interval_ms > 0:
  517. await asyncio.sleep(self.config.udp_copy_interval_ms / 1000)
  518. if not sent_any:
  519. udp_server.note_unsent(flow, packet_id)
  520. async def _handshake(self, reader: asyncio.StreamReader, writer: asyncio.StreamWriter, peer) -> tuple[str, int, bool]:
  521. version, methods_len = (await read_exact(reader, 2))
  522. if version != SOCKS_VERSION:
  523. raise ValueError("unsupported socks version")
  524. await read_exact(reader, methods_len)
  525. writer.write(b"\x05\x00")
  526. await writer.drain()
  527. version, command, _, atyp = await read_exact(reader, 4)
  528. if version != SOCKS_VERSION:
  529. raise ValueError("unsupported socks version")
  530. if atyp == 1:
  531. host = socket.inet_ntoa(await read_exact(reader, 4))
  532. elif atyp == 3:
  533. size = (await read_exact(reader, 1))[0]
  534. host = (await read_exact(reader, size)).decode()
  535. else:
  536. raise ValueError("unsupported atyp")
  537. port = struct.unpack("!H", await read_exact(reader, 2))[0]
  538. peer_text = f"{peer[0]}:{peer[1]}" if isinstance(peer, tuple) and len(peer) >= 2 else str(peer)
  539. if command == 1:
  540. print(f"[edge] socks handshake peer={peer_text} command=connect target={host}:{port}")
  541. writer.write(b"\x05\x00\x00\x01\x00\x00\x00\x00\x00\x00")
  542. await writer.drain()
  543. return host, port, False
  544. if command == 3 and self.udp_server and self.udp_server.transport:
  545. bind_host, bind_port = self.udp_server.transport.get_extra_info("sockname")[:2]
  546. self.udp_server.register_associate(peer)
  547. print(f"[edge] socks handshake peer={peer_text} command=udp_associate target={host}:{port} bind={bind_host}:{bind_port}")
  548. writer.write(b"\x05\x00\x00\x01" + socket.inet_aton(bind_host) + struct.pack("!H", bind_port))
  549. await writer.drain()
  550. return host, port, True
  551. raise ValueError("unsupported socks command")