transparent_edge.py 40 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886
  1. from __future__ import annotations
  2. from pathlib import Path
  3. import asyncio
  4. import contextlib
  5. import itertools
  6. import os
  7. import socket
  8. import struct
  9. from dataclasses import dataclass, field
  10. from typing import Awaitable, Callable
  11. from .config import Config
  12. from .protocol import STATUS_OK, TCP_CLOSE, TCP_DATA, TCP_OPEN, TCP_STATUS, UDP_RECV, UDP_SEND, Frame, encode_json
  13. from .relay_client import RelayConnection, RelayManager
  14. SO_ORIGINAL_DST = 80
  15. IP6T_SO_ORIGINAL_DST = 80
  16. IP_RECVORIGDSTADDR = 20
  17. IPV6_RECVORIGDSTADDR = 74
  18. @dataclass(frozen=True)
  19. class TargetAddress:
  20. host: str
  21. port: int
  22. family: int
  23. @dataclass(frozen=True)
  24. class PeerAddress:
  25. host: str
  26. port: int
  27. family: int
  28. def parse_sockaddr(raw: bytes) -> TargetAddress:
  29. if len(raw) < 8:
  30. raise ValueError("invalid transparent destination payload")
  31. family = struct.unpack_from("=H", raw, 0)[0]
  32. port = struct.unpack_from("!H", raw, 2)[0]
  33. if family == socket.AF_INET:
  34. host = socket.inet_ntoa(raw[4:8])
  35. return TargetAddress(host=host, port=port, family=family)
  36. if family == socket.AF_INET6:
  37. if len(raw) < 28:
  38. raise ValueError("invalid IPv6 transparent destination payload")
  39. host = socket.inet_ntop(socket.AF_INET6, raw[8:24])
  40. return TargetAddress(host=host, port=port, family=family)
  41. raise ValueError(f"unsupported family={family}")
  42. def winner_group(name: str) -> str:
  43. return "direct" if name.startswith("direct") else name
  44. def grouped_total(stats: dict[str, int], group: str) -> int:
  45. return sum(count for name, count in stats.items() if winner_group(name) == group)
  46. class BasePath:
  47. def __init__(self, name: str, on_frame: Callable[["BasePath", str, bytes | None], Awaitable[None]]) -> None:
  48. self.name = name
  49. self.on_frame = on_frame
  50. self.opened = False
  51. self.closed = False
  52. async def open(self, target: TargetAddress) -> None:
  53. raise NotImplementedError
  54. async def send(self, data: bytes) -> None:
  55. raise NotImplementedError
  56. async def close(self) -> None:
  57. raise NotImplementedError
  58. class DirectTcpPath(BasePath):
  59. def __init__(self, name: str, on_frame: Callable[[BasePath, str, bytes | None], Awaitable[None]], open_timeout: float, happy_eyeballs_delay: float | None, tcp_nodelay: bool = True) -> None:
  60. super().__init__(name, on_frame)
  61. self.reader: asyncio.StreamReader | None = None
  62. self.writer: asyncio.StreamWriter | None = None
  63. self.pump_task: asyncio.Task | None = None
  64. self.open_timeout = open_timeout
  65. self.happy_eyeballs_delay = happy_eyeballs_delay
  66. self.tcp_nodelay = tcp_nodelay
  67. async def open(self, target: TargetAddress) -> None:
  68. try:
  69. family = socket.AF_INET6 if target.family == socket.AF_INET6 else socket.AF_INET
  70. kwargs = {"host": target.host, "port": target.port, "family": family}
  71. if self.happy_eyeballs_delay is not None:
  72. kwargs["happy_eyeballs_delay"] = self.happy_eyeballs_delay
  73. self.reader, self.writer = await asyncio.wait_for(asyncio.open_connection(**kwargs), timeout=self.open_timeout)
  74. sock = self.writer.get_extra_info("socket")
  75. if sock is not None and self.tcp_nodelay:
  76. with contextlib.suppress(OSError):
  77. sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
  78. self.opened = True
  79. self.pump_task = asyncio.create_task(self._pump())
  80. await self.on_frame(self, "status", b"ok")
  81. except Exception as exc:
  82. await self.on_frame(self, "status", str(exc).encode())
  83. async def _pump(self) -> None:
  84. assert self.reader is not None
  85. try:
  86. while True:
  87. chunk = await self.reader.read(65536)
  88. if not chunk:
  89. break
  90. await self.on_frame(self, "data", chunk)
  91. except Exception:
  92. pass
  93. finally:
  94. await self.on_frame(self, "close", None)
  95. async def send(self, data: bytes) -> None:
  96. if self.closed or self.writer is None:
  97. return
  98. try:
  99. self.writer.write(data)
  100. await self.writer.drain()
  101. except (BrokenPipeError, ConnectionResetError, RuntimeError, OSError, asyncio.CancelledError) as exc:
  102. await self.close()
  103. raise ConnectionError("relay closed") from exc
  104. async def close(self) -> None:
  105. if self.closed:
  106. return
  107. self.closed = True
  108. if self.pump_task and self.pump_task is not asyncio.current_task():
  109. self.pump_task.cancel()
  110. with contextlib.suppress(Exception):
  111. await self.pump_task
  112. if self.writer:
  113. self.writer.close()
  114. with contextlib.suppress(Exception):
  115. await self.writer.wait_closed()
  116. class RelayTcpPath(BasePath):
  117. def __init__(self, name: str, on_frame: Callable[[BasePath, str, bytes | None], Awaitable[None]], connection: RelayConnection, session_id: int, stream_id: int) -> None:
  118. super().__init__(name, on_frame)
  119. self.connection = connection
  120. self.session_id = session_id
  121. self.stream_id = stream_id
  122. self.unbind_task: asyncio.Task | None = None
  123. async def open(self, target: TargetAddress) -> None:
  124. if self.connection.closed:
  125. await self.on_frame(self, "status", b"relay unavailable")
  126. return
  127. self.connection.bind(self.session_id, self.stream_id, self._handle_frame)
  128. try:
  129. await self.connection.send(Frame(TCP_OPEN, self.session_id, self.stream_id, 0, 0, encode_json({"host": target.host, "port": target.port, "family": target.family})))
  130. except Exception as exc:
  131. self.connection.unbind(self.session_id, self.stream_id)
  132. await self.on_frame(self, "status", str(exc).encode())
  133. async def _handle_frame(self, _conn: RelayConnection, frame: Frame) -> None:
  134. if frame.kind == TCP_STATUS:
  135. if frame.packet_id == STATUS_OK:
  136. self.opened = True
  137. await self.on_frame(self, "status", b"ok")
  138. else:
  139. await self.on_frame(self, "status", frame.payload)
  140. return
  141. if frame.kind == TCP_DATA:
  142. await self.on_frame(self, "data", frame.payload)
  143. return
  144. if frame.kind == TCP_CLOSE:
  145. await self.on_frame(self, "close", None)
  146. async def send(self, data: bytes) -> None:
  147. if self.closed or self.connection.closed:
  148. return
  149. await self.connection.send(Frame(TCP_DATA, self.session_id, self.stream_id, 0, 0, data))
  150. async def close(self) -> None:
  151. if self.closed:
  152. return
  153. self.closed = True
  154. if self.unbind_task is None or self.unbind_task.done():
  155. self.unbind_task = asyncio.create_task(self._delayed_unbind())
  156. if not self.connection.closed:
  157. with contextlib.suppress(Exception):
  158. await self.connection.send(Frame(TCP_CLOSE, self.session_id, self.stream_id, 0, 0, b""))
  159. async def _delayed_unbind(self) -> None:
  160. await asyncio.sleep(0.5)
  161. self.connection.unbind(self.session_id, self.stream_id)
  162. @dataclass
  163. class TransparentSession:
  164. session_id: int
  165. target: TargetAddress
  166. reader: asyncio.StreamReader
  167. writer: asyncio.StreamWriter
  168. paths: list[BasePath]
  169. warmup_bytes: int
  170. loser_grace_ms: int
  171. tcp_failover_idle_ms: int
  172. stats: dict[str, int]
  173. target_stats: dict[tuple[str, int], dict[str, int]]
  174. family_stats: dict[str, dict[str, int]]
  175. opened_count: int = 0
  176. status_count: int = 0
  177. errors: list[str] = field(default_factory=list)
  178. winner: BasePath | None = None
  179. uplink_bytes: int = 0
  180. open_event: asyncio.Event = field(default_factory=asyncio.Event)
  181. winner_event: asyncio.Event = field(default_factory=asyncio.Event)
  182. closed: bool = False
  183. pump_task: asyncio.Task | None = None
  184. loser_close_task: asyncio.Task | None = None
  185. open_tasks: list[asyncio.Task] = field(default_factory=list)
  186. backup_path: BasePath | None = None
  187. last_winner_data_at: float = 0.0
  188. converged: bool = False
  189. def _select_backup_path(self, winner: BasePath) -> BasePath | None:
  190. candidates = [path for path in self.paths if path is not winner and path.opened and not path.closed]
  191. if not candidates:
  192. return None
  193. winner_is_direct = winner_group(winner.name) == "direct"
  194. # Prefer the opposite group to increase failover diversity.
  195. opposite = [path for path in candidates if (winner_group(path.name) == "direct") != winner_is_direct]
  196. pool = opposite or candidates
  197. # Keep the first eligible path as a synchronized backup.
  198. return pool[0]
  199. def _record_win(self, winner: BasePath) -> None:
  200. self.stats[winner.name] = self.stats.get(winner.name, 0) + 1
  201. key = (self.target.host, self.target.port)
  202. target_stats = self.target_stats.setdefault(key, {})
  203. target_stats[winner.name] = target_stats.get(winner.name, 0) + 1
  204. family_key = "ipv6" if self.target.family == socket.AF_INET6 else "ipv4"
  205. family_stats = self.family_stats.setdefault(family_key, {})
  206. family_stats[winner.name] = family_stats.get(winner.name, 0) + 1
  207. direct_wins = grouped_total(self.stats, "direct")
  208. relay_wins = sum(count for name, count in self.stats.items() if winner_group(name) != "direct")
  209. target_direct = grouped_total(target_stats, "direct")
  210. target_relay = sum(count for name, count in target_stats.items() if winner_group(name) != "direct")
  211. family_direct = grouped_total(family_stats, "direct")
  212. family_relay = sum(count for name, count in family_stats.items() if winner_group(name) != "direct")
  213. relay_detail = ", ".join(f"{name}={count}" for name, count in sorted(self.stats.items()) if winner_group(name) != "direct") or "none"
  214. target_detail = ", ".join(f"{name}={count}" for name, count in sorted(target_stats.items()) if winner_group(name) != "direct") or "none"
  215. target_pref = "relay" if target_relay > target_direct else "direct"
  216. family_pref = "relay" if family_relay > family_direct else "direct"
  217. print(f"[edge] tcp win session={self.session_id} target={self.target.host}:{self.target.port} winner={winner.name} direct={direct_wins} relay={relay_wins} relay_breakdown={relay_detail} target_pref={target_pref} target_direct={target_direct} target_relay={target_relay} target_breakdown={target_detail} family_pref={family_pref} family={family_key} family_direct={family_direct} family_relay={family_relay}")
  218. async def start(self) -> None:
  219. self.open_tasks = [asyncio.create_task(path.open(self.target)) for path in self.paths]
  220. await asyncio.wait_for(self.open_event.wait(), timeout=8)
  221. if self.opened_count == 0:
  222. raise ConnectionError(self.errors[0] if self.errors else "all paths failed")
  223. self.pump_task = asyncio.create_task(self._pump_local())
  224. async def _pump_local(self) -> None:
  225. try:
  226. while True:
  227. chunk = await self.reader.read(65536)
  228. if not chunk:
  229. break
  230. self.uplink_bytes += len(chunk)
  231. active = [path for path in self.paths if path.opened and not path.closed]
  232. if not active:
  233. break
  234. if self.winner is None and self.uplink_bytes <= self.warmup_bytes:
  235. await asyncio.gather(*(path.send(chunk) for path in active), return_exceptions=True)
  236. else:
  237. if self.winner is None:
  238. await self.winner_event.wait()
  239. if self.winner:
  240. send_targets = [self.winner]
  241. if (
  242. not self.converged
  243. and self.backup_path
  244. and self.backup_path.opened
  245. and not self.backup_path.closed
  246. and self.backup_path is not self.winner
  247. ):
  248. send_targets.append(self.backup_path)
  249. await asyncio.gather(*(path.send(chunk) for path in send_targets), return_exceptions=True)
  250. self.converged = True
  251. except Exception:
  252. pass
  253. finally:
  254. await self.close()
  255. async def handle_path(self, path: BasePath, event: str, payload: bytes | None) -> None:
  256. if self.closed:
  257. return
  258. if event == "status":
  259. self.status_count += 1
  260. if payload == b"ok":
  261. self.opened_count += 1
  262. elif payload is not None:
  263. self.errors.append(payload.decode("utf-8", errors="replace"))
  264. if self.opened_count > 0 or self.status_count == len(self.paths):
  265. self.open_event.set()
  266. return
  267. if event == "data":
  268. if self.winner is None:
  269. self.winner = path
  270. self._record_win(path)
  271. self.backup_path = self._select_backup_path(path)
  272. self.winner_event.set()
  273. self.converged = False
  274. if self.loser_grace_ms > 0:
  275. self.loser_close_task = asyncio.create_task(self._close_losers_after_grace(path))
  276. else:
  277. await self._close_losers(path)
  278. self.last_winner_data_at = asyncio.get_running_loop().time()
  279. if path is self.winner and payload is not None:
  280. self.writer.write(payload)
  281. await self.writer.drain()
  282. return
  283. if event == "close":
  284. path.closed = True
  285. if self.winner is None:
  286. remaining = [candidate for candidate in self.paths if candidate.opened and not candidate.closed]
  287. if not remaining:
  288. await self.close()
  289. elif path is self.winner:
  290. await self.close()
  291. async def _close_losers(self, winner: BasePath) -> None:
  292. await asyncio.gather(*(path.close() for path in self.paths if path is not winner and path is not self.backup_path), return_exceptions=True)
  293. async def _close_losers_after_grace(self, winner: BasePath) -> None:
  294. await asyncio.sleep(self.loser_grace_ms / 1000)
  295. if not self.closed:
  296. await self._close_losers(winner)
  297. async def close(self) -> None:
  298. if self.closed:
  299. return
  300. self.closed = True
  301. if self.errors:
  302. detail = ", ".join(self.errors[:3])
  303. print(
  304. f"[edge] session={self.session_id} closed target={self.target.host}:{self.target.port} "
  305. f"errors={len(self.errors)} detail={detail}"
  306. )
  307. if self.pump_task and self.pump_task is not asyncio.current_task():
  308. self.pump_task.cancel()
  309. with contextlib.suppress(Exception):
  310. await self.pump_task
  311. if self.loser_close_task and self.loser_close_task is not asyncio.current_task():
  312. self.loser_close_task.cancel()
  313. with contextlib.suppress(Exception):
  314. await self.loser_close_task
  315. for task in self.open_tasks:
  316. if task is not asyncio.current_task():
  317. task.cancel()
  318. for task in self.open_tasks:
  319. if task is not asyncio.current_task():
  320. with contextlib.suppress(Exception):
  321. await task
  322. await asyncio.gather(*(path.close() for path in self.paths), return_exceptions=True)
  323. self.writer.close()
  324. with contextlib.suppress(Exception):
  325. await self.writer.wait_closed()
  326. class DirectUdpPath(BasePath):
  327. def __init__(self, name: str, on_frame: Callable[[BasePath, str, bytes | None], Awaitable[None]], target: TargetAddress) -> None:
  328. super().__init__(name, on_frame)
  329. self.target = target
  330. self.socket: socket.socket | None = None
  331. self.read_task: asyncio.Task | None = None
  332. async def open(self, _target: TargetAddress) -> None:
  333. try:
  334. family = socket.AF_INET6 if self.target.family == socket.AF_INET6 else socket.AF_INET
  335. self.socket = socket.socket(family, socket.SOCK_DGRAM)
  336. self.socket.setblocking(False)
  337. await asyncio.get_running_loop().sock_connect(self.socket, (self.target.host, self.target.port))
  338. self.opened = True
  339. self.read_task = asyncio.create_task(self._pump())
  340. await self.on_frame(self, "status", b"ok")
  341. except Exception as exc:
  342. await self.on_frame(self, "status", str(exc).encode())
  343. async def _pump(self) -> None:
  344. assert self.socket is not None
  345. loop = asyncio.get_running_loop()
  346. try:
  347. while True:
  348. data = await loop.sock_recv(self.socket, 65535)
  349. if not data:
  350. break
  351. await self.on_frame(self, "data", data)
  352. except Exception:
  353. pass
  354. finally:
  355. await self.on_frame(self, "close", None)
  356. async def send(self, data: bytes) -> None:
  357. if self.closed or self.socket is None:
  358. return
  359. await asyncio.get_running_loop().sock_sendall(self.socket, data)
  360. async def close(self) -> None:
  361. if self.closed:
  362. return
  363. self.closed = True
  364. if self.read_task and self.read_task is not asyncio.current_task():
  365. self.read_task.cancel()
  366. with contextlib.suppress(Exception):
  367. await self.read_task
  368. if self.socket:
  369. self.socket.close()
  370. class RelayUdpPath(BasePath):
  371. def __init__(self, name: str, on_frame: Callable[[BasePath, str, bytes | None], Awaitable[None]], connection: RelayConnection, session_id: int, stream_id: int, target: TargetAddress) -> None:
  372. super().__init__(name, on_frame)
  373. self.connection = connection
  374. self.session_id = session_id
  375. self.stream_id = stream_id
  376. self.target = target
  377. self.unbind_task: asyncio.Task | None = None
  378. async def open(self, _target: TargetAddress) -> None:
  379. if self.connection.closed:
  380. await self.on_frame(self, "status", b"relay unavailable")
  381. return
  382. self.connection.bind(self.session_id, self.stream_id, self._handle_frame)
  383. try:
  384. self.opened = True
  385. await self.on_frame(self, "status", b"ok")
  386. except Exception:
  387. self.connection.unbind(self.session_id, self.stream_id)
  388. self.closed = True
  389. raise
  390. async def _handle_frame(self, _conn: RelayConnection, frame: Frame) -> None:
  391. if frame.kind == UDP_RECV:
  392. await self.on_frame(self, "data", frame.payload)
  393. async def send(self, data: bytes) -> None:
  394. if self.closed or self.connection.closed:
  395. return
  396. meta = encode_json({"host": self.target.host, "port": self.target.port, "family": self.target.family})
  397. payload = meta + data
  398. try:
  399. await self.connection.send(Frame(UDP_SEND, self.session_id, self.stream_id, 0, len(meta), payload))
  400. except Exception:
  401. self.closed = True
  402. raise
  403. async def close(self) -> None:
  404. if self.closed:
  405. return
  406. self.closed = True
  407. if self.unbind_task is None or self.unbind_task.done():
  408. self.unbind_task = asyncio.create_task(self._delayed_unbind())
  409. async def _delayed_unbind(self) -> None:
  410. await asyncio.sleep(0.5)
  411. self.connection.unbind(self.session_id, self.stream_id)
  412. @dataclass
  413. class UdpFlow:
  414. flow_id: int
  415. source: PeerAddress
  416. target: TargetAddress
  417. send_response: Callable[[PeerAddress, bytes], Awaitable[None]]
  418. paths: list[BasePath]
  419. redundancy: int = 0
  420. always_broadcast: bool = True
  421. copy_interval_ms: int = 0
  422. winner: BasePath | None = None
  423. closed: bool = False
  424. last_activity: float = 0.0
  425. packets_sent: int = 0
  426. packets_received: int = 0
  427. duplicate_responses: int = 0
  428. send_task: asyncio.Task | None = None
  429. winner_burst_sent: int = 0
  430. converged: bool = False
  431. async def start(self) -> None:
  432. await asyncio.gather(*(path.open(self.target) for path in self.paths), return_exceptions=True)
  433. async def send(self, payload: bytes) -> None:
  434. self.last_activity = asyncio.get_running_loop().time()
  435. self.packets_sent += 1
  436. active = [path for path in self.paths if path.opened and not path.closed]
  437. if not active:
  438. return
  439. copies = max(1, self.redundancy + 1)
  440. if self.winner is None or self.winner.closed:
  441. self.converged = False
  442. self.winner_burst_sent = 0
  443. targets = active
  444. elif not self.converged:
  445. # 先并发、后收敛:winner 刚出现时保留短暂重叠,随后快速收敛到单路径。
  446. self.winner_burst_sent += 1
  447. backup = [path for path in active if path is not self.winner][:1]
  448. targets = [self.winner, *backup] if self.winner_burst_sent <= 2 else [self.winner]
  449. if self.winner_burst_sent > 2:
  450. self.converged = True
  451. else:
  452. targets = [self.winner]
  453. for attempt in range(copies):
  454. await asyncio.gather(*(path.send(payload) for path in targets), return_exceptions=True)
  455. if attempt + 1 < copies and self.copy_interval_ms > 0:
  456. await asyncio.sleep(self.copy_interval_ms / 1000)
  457. async def handle_path(self, path: BasePath, event: str, payload: bytes | None) -> None:
  458. self.last_activity = asyncio.get_running_loop().time()
  459. if event == "data" and payload is not None:
  460. self.packets_received += 1
  461. if self.winner is None:
  462. self.winner = path
  463. self.converged = False
  464. self.winner_burst_sent = 0
  465. mode = "redundant" if self.redundancy > 0 else "single"
  466. print(f"[edge] udp flow={self.flow_id} winner={path.name} target={self.target.host}:{self.target.port} mode={mode} candidates={len(self.paths)}")
  467. elif path is not self.winner:
  468. self.duplicate_responses += 1
  469. if path is self.winner:
  470. await self.send_response(self.source, payload)
  471. if event == "close":
  472. path.closed = True
  473. if path is self.winner:
  474. remaining = [candidate for candidate in self.paths if candidate.opened and not candidate.closed]
  475. self.winner = remaining[0] if remaining else None
  476. self.converged = False
  477. self.winner_burst_sent = 0
  478. async def close(self) -> None:
  479. if self.closed:
  480. return
  481. self.closed = True
  482. if self.send_task and self.send_task is not asyncio.current_task():
  483. self.send_task.cancel()
  484. with contextlib.suppress(Exception):
  485. await self.send_task
  486. await asyncio.gather(*(path.close() for path in self.paths), return_exceptions=True)
  487. class TransparentUdpListener:
  488. def __init__(self, edge: "TransparentEdge", family: int, bind_host: str, port: int) -> None:
  489. self.edge = edge
  490. self.family = family
  491. self.bind_host = bind_host
  492. self.port = port
  493. self.socket: socket.socket | None = None
  494. self.udp_packets_received = 0
  495. self.udp_recv_errors = 0
  496. self.udp_parse_errors = 0
  497. self.udp_missing_original = 0
  498. self.udp_self_loop_skipped = 0
  499. self.udp_flows_created = 0
  500. self.last_summary_at = 0.0
  501. def start(self) -> None:
  502. sock = socket.socket(self.family, socket.SOCK_DGRAM)
  503. sock.setblocking(False)
  504. if self.family == socket.AF_INET:
  505. sock.setsockopt(socket.SOL_IP, IP_RECVORIGDSTADDR, 1)
  506. sock.bind((self.bind_host, self.port))
  507. else:
  508. sock.setsockopt(socket.IPPROTO_IPV6, socket.IPV6_V6ONLY, 1)
  509. sock.setsockopt(socket.IPPROTO_IPV6, IPV6_RECVORIGDSTADDR, 1)
  510. sock.bind((self.bind_host, self.port, 0, 0))
  511. self.socket = sock
  512. asyncio.get_running_loop().add_reader(sock.fileno(), self._on_readable)
  513. print(f"[edge] transparent udp listening on {sock.getsockname()}")
  514. def _log_udp_summary(self, force: bool = False) -> None:
  515. now = asyncio.get_running_loop().time()
  516. if not force and now - self.last_summary_at < 10:
  517. return
  518. self.last_summary_at = now
  519. print(
  520. f"[edge] udp summary family={self.family} bind={self.bind_host}:{self.port} "
  521. f"received={self.udp_packets_received} flows={self.udp_flows_created} "
  522. f"self_loop={self.udp_self_loop_skipped} missing_original={self.udp_missing_original} "
  523. f"parse_error={self.udp_parse_errors} recv_error={self.udp_recv_errors}"
  524. )
  525. def _on_readable(self) -> None:
  526. assert self.socket is not None
  527. try:
  528. data, ancdata, _flags, src = self.socket.recvmsg(65535, 512)
  529. except BlockingIOError:
  530. return
  531. except Exception as exc:
  532. self.udp_recv_errors += 1
  533. print(f"[edge] udp recvmsg error family={self.family} error={exc!r}")
  534. self._log_udp_summary(force=True)
  535. return
  536. self.udp_packets_received += 1
  537. original = None
  538. for level, ctype, cdata in ancdata:
  539. if self.family == socket.AF_INET and level == socket.SOL_IP and ctype == IP_RECVORIGDSTADDR:
  540. try:
  541. original = parse_sockaddr(cdata)
  542. except Exception as exc:
  543. self.udp_parse_errors += 1
  544. print(f"[edge] udp parse original dst error family={self.family} src={src} error={exc!r} raw_len={len(cdata)}")
  545. self._log_udp_summary(force=True)
  546. return
  547. break
  548. if self.family == socket.AF_INET6 and level == socket.IPPROTO_IPV6 and ctype == IPV6_RECVORIGDSTADDR:
  549. try:
  550. original = parse_sockaddr(cdata)
  551. except Exception as exc:
  552. self.udp_parse_errors += 1
  553. print(f"[edge] udp parse original dst error family={self.family} src={src} error={exc!r} raw_len={len(cdata)}")
  554. self._log_udp_summary(force=True)
  555. return
  556. break
  557. if original is None:
  558. self.udp_missing_original += 1
  559. self._log_udp_summary()
  560. return
  561. if self.family == socket.AF_INET:
  562. source = PeerAddress(host=src[0], port=src[1], family=socket.AF_INET)
  563. else:
  564. source = PeerAddress(host=src[0], port=src[1], family=socket.AF_INET6)
  565. if original.port == self.port and (original.host in ("127.0.0.1", "::1") or original.host == self.bind_host):
  566. self.udp_self_loop_skipped += 1
  567. print(
  568. f"[edge] udp self_loop family={self.family} src={source.host}:{source.port} "
  569. f"original={original.host}:{original.port} size={len(data)}"
  570. )
  571. self._log_udp_summary()
  572. return
  573. asyncio.create_task(self.edge.handle_udp_datagram(source, original, data, self))
  574. async def send_response(self, source: PeerAddress, payload: bytes) -> None:
  575. assert self.socket is not None
  576. if source.family == socket.AF_INET:
  577. self.socket.sendto(payload, (source.host, source.port))
  578. else:
  579. self.socket.sendto(payload, (source.host, source.port, 0, 0))
  580. async def close(self) -> None:
  581. if self.socket is None:
  582. return
  583. asyncio.get_running_loop().remove_reader(self.socket.fileno())
  584. self.socket.close()
  585. self.socket = None
  586. class TransparentEdge:
  587. def __init__(self, listen_host: str, listen_port: int, config: Config, enable_udp: bool = False, kernel_mode: str = "auto") -> None:
  588. self.listen_host = listen_host
  589. self.listen_port = listen_port
  590. self.config = config
  591. self.enable_udp = enable_udp
  592. self.kernel_mode = self._resolve_kernel_mode(kernel_mode, config.kernel_mode)
  593. self.manager = RelayManager(config)
  594. self.session_ids = itertools.count(1)
  595. self.stream_ids = itertools.count(1)
  596. self.udp_listeners: list[TransparentUdpListener] = []
  597. self.udp_flows: dict[tuple[PeerAddress, TargetAddress], UdpFlow] = {}
  598. self.udp_flow_ids = itertools.count(1)
  599. self.udp_gc_task: asyncio.Task | None = None
  600. self.tcp_win_counts: dict[str, int] = {}
  601. self.tcp_target_wins: dict[tuple[str, int], dict[str, int]] = {}
  602. self.tcp_family_wins: dict[str, dict[str, int]] = {"ipv4": {}, "ipv6": {}}
  603. def _resolve_kernel_mode(self, cli_kernel_mode: str, config_kernel_mode: str) -> str:
  604. mode = cli_kernel_mode if cli_kernel_mode != "auto" else config_kernel_mode
  605. if mode != "auto":
  606. return mode
  607. try:
  608. if Path("/etc/os-release").exists() and 'VERSION_ID="24' in Path("/etc/os-release").read_text(errors="ignore"):
  609. return "24"
  610. except Exception:
  611. pass
  612. try:
  613. release = os.uname().release
  614. if release.startswith("6."):
  615. return "24"
  616. except Exception:
  617. pass
  618. return "20"
  619. async def start(self) -> None:
  620. if self.kernel_mode == "24":
  621. if self.config.direct_open_timeout == 10.0:
  622. self.config.direct_open_timeout = 6.0
  623. if self.config.relay_open_timeout == 10.0:
  624. self.config.relay_open_timeout = 6.0
  625. if self.config.tcp_connect_happy_eyeballs_delay is None:
  626. self.config.tcp_connect_happy_eyeballs_delay = 0.25
  627. await self.manager.start()
  628. print(f"[edge] kernel_mode={self.kernel_mode} relay snapshot: {self.manager.snapshot()}")
  629. server4 = await asyncio.start_server(self._accept, self.listen_host, self.listen_port, family=socket.AF_INET)
  630. sockets = [str(sock.getsockname()) for sock in server4.sockets or []]
  631. server6 = None
  632. if self.listen_host in ("::", "::1", "0.0.0.0", "127.0.0.1"):
  633. host6 = "::1" if self.listen_host == "127.0.0.1" else "::"
  634. try:
  635. server6 = await asyncio.start_server(self._accept, host6, self.listen_port, family=socket.AF_INET6)
  636. sockets.extend(str(sock.getsockname()) for sock in server6.sockets or [])
  637. except Exception as exc:
  638. print(f"[edge] ipv6 tcp listener skipped: {exc!r}")
  639. if self.enable_udp:
  640. self._start_udp_listeners()
  641. self.udp_gc_task = asyncio.create_task(self._gc_udp_flows())
  642. print(f"[edge] transparent tcp listening on {', '.join(sockets)}")
  643. if server6 is None:
  644. async with server4:
  645. await server4.serve_forever()
  646. else:
  647. async with server4, server6:
  648. await asyncio.gather(server4.serve_forever(), server6.serve_forever())
  649. def _direct_redundancy_for_target(self, target: TargetAddress) -> int:
  650. if target.family == socket.AF_INET6 and not self.config.direct_ipv6_enabled:
  651. return 0
  652. base = self.config.direct_redundancy
  653. if target.family == socket.AF_INET6 and self.config.direct_redundancy_v6 is not None:
  654. base = self.config.direct_redundancy_v6
  655. elif target.family == socket.AF_INET and self.config.direct_redundancy_v4 is not None:
  656. base = self.config.direct_redundancy_v4
  657. base = max(1, min(base, self.config.direct_max_redundancy))
  658. target_stats = self.tcp_target_wins.get((target.host, target.port), {})
  659. family_key = "ipv6" if target.family == socket.AF_INET6 else "ipv4"
  660. family_stats = self.tcp_family_wins.get(family_key, {})
  661. target_total = sum(target_stats.values())
  662. family_total = sum(family_stats.values())
  663. target_relay = sum(count for name, count in target_stats.items() if winner_group(name) != "direct")
  664. family_relay = sum(count for name, count in family_stats.items() if winner_group(name) != "direct")
  665. target_prefers_relay = target_total >= 4 and target_relay > grouped_total(target_stats, "direct")
  666. family_prefers_relay = family_total >= 8 and family_relay > grouped_total(family_stats, "direct")
  667. target_prefers_direct = target_total >= 4 and grouped_total(target_stats, "direct") > target_relay
  668. family_prefers_direct = family_total >= 8 and grouped_total(family_stats, "direct") > family_relay
  669. if target_prefers_relay or family_prefers_relay:
  670. return max(1, base - 1)
  671. if target_prefers_direct or family_prefers_direct:
  672. if base > 2:
  673. return base - 1
  674. return base
  675. def _build_direct_paths(self, session: TransparentSession) -> list[BasePath]:
  676. count = self._direct_redundancy_for_target(session.target)
  677. if count <= 0:
  678. return []
  679. return [
  680. DirectTcpPath(
  681. name=f"direct-{index + 1}" if count > 1 else "direct",
  682. on_frame=lambda path, event, payload, s=session: self._handle_tcp_session(s, path, event, payload),
  683. open_timeout=self.config.direct_open_timeout,
  684. happy_eyeballs_delay=self.config.tcp_connect_happy_eyeballs_delay,
  685. tcp_nodelay=self.config.relay_tcp_nodelay,
  686. )
  687. for index in range(count)
  688. ]
  689. def _build_udp_direct_paths(self, target: TargetAddress, flow_id: int) -> list[BasePath]:
  690. if target.family == socket.AF_INET6 and not self.config.direct_ipv6_enabled:
  691. return []
  692. count = max(1, self.config.udp_direct_redundancy)
  693. if target.family == socket.AF_INET6 and self.config.udp_direct_redundancy_v6 is not None:
  694. count = max(1, self.config.udp_direct_redundancy_v6)
  695. elif target.family == socket.AF_INET and self.config.udp_direct_redundancy_v4 is not None:
  696. count = max(1, self.config.udp_direct_redundancy_v4)
  697. return [
  698. DirectUdpPath(
  699. name=f"direct-{index + 1}" if count > 1 else "direct",
  700. on_frame=lambda path, event, data, fid=flow_id: self._handle_udp_path(fid, path, event, data),
  701. target=target,
  702. )
  703. for index in range(count)
  704. ]
  705. def _tcp_relay_connections(self) -> list[RelayConnection]:
  706. return [connection for connection in self.manager.available() if connection.supports_tcp]
  707. def _udp_relay_connections(self) -> list[RelayConnection]:
  708. if hasattr(self.manager, "available_udp"):
  709. return [connection for connection in self.manager.available_udp() if connection.supports_udp]
  710. return [connection for connection in self.manager.available() if connection.supports_udp]
  711. def _start_udp_listeners(self) -> None:
  712. binds = []
  713. if self.listen_host == "127.0.0.1":
  714. binds = [(socket.AF_INET, "127.0.0.1"), (socket.AF_INET6, "::1")]
  715. elif self.listen_host == "0.0.0.0":
  716. binds = [(socket.AF_INET, "0.0.0.0"), (socket.AF_INET6, "::")]
  717. else:
  718. family = socket.AF_INET6 if ":" in self.listen_host else socket.AF_INET
  719. binds = [(family, self.listen_host)]
  720. for family, host in binds:
  721. try:
  722. listener = TransparentUdpListener(self, family, host, self.listen_port)
  723. listener.start()
  724. self.udp_listeners.append(listener)
  725. except Exception as exc:
  726. print(f"[edge] udp listener skipped family={family} host={host} error={exc!r}")
  727. async def _accept(self, reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None:
  728. peer = writer.get_extra_info("peername")
  729. try:
  730. target = self._get_original_dst(writer)
  731. session_id = next(self.session_ids)
  732. session = TransparentSession(session_id=session_id, target=target, reader=reader, writer=writer, paths=[], warmup_bytes=self.config.tcp_warmup_bytes, loser_grace_ms=self.config.tcp_loser_grace_ms, tcp_failover_idle_ms=self.config.tcp_failover_idle_ms, stats=self.tcp_win_counts, target_stats=self.tcp_target_wins, family_stats=self.tcp_family_wins)
  733. paths: list[BasePath] = self._build_direct_paths(session)
  734. for connection in self._tcp_relay_connections():
  735. stream_id = next(self.stream_ids)
  736. paths.append(RelayTcpPath(name=connection.node.name, on_frame=lambda path, event, payload, s=session: self._handle_tcp_session(s, path, event, payload), connection=connection, session_id=session_id, stream_id=stream_id))
  737. session.paths = paths
  738. print(f"[edge] accept peer={peer} session={session_id} target={target.host}:{target.port} candidates={[path.name for path in paths]}")
  739. await session.start()
  740. except Exception as exc:
  741. print(f"[edge] accept failed peer={peer} error={exc!r}")
  742. writer.close()
  743. with contextlib.suppress(Exception):
  744. await writer.wait_closed()
  745. async def _handle_tcp_session(self, session: TransparentSession, path: BasePath, event: str, payload: bytes | None) -> None:
  746. await session.handle_path(path, event, payload)
  747. def _get_original_dst(self, writer: asyncio.StreamWriter) -> TargetAddress:
  748. sock = writer.get_extra_info("socket")
  749. if sock is None:
  750. raise RuntimeError("socket unavailable")
  751. family = sock.family
  752. if family == socket.AF_INET:
  753. raw = sock.getsockopt(socket.SOL_IP, SO_ORIGINAL_DST, 16)
  754. return parse_sockaddr(raw)
  755. if family == socket.AF_INET6:
  756. raw = sock.getsockopt(socket.IPPROTO_IPV6, IP6T_SO_ORIGINAL_DST, 128)
  757. return parse_sockaddr(raw)
  758. raise RuntimeError(f"unsupported socket family={family}")
  759. async def handle_udp_datagram(self, source: PeerAddress, target: TargetAddress, payload: bytes, listener: TransparentUdpListener) -> None:
  760. if not self.enable_udp:
  761. return
  762. if target.port == self.listen_port and target.host in ("127.0.0.1", "::1", self.listen_host):
  763. return
  764. key = (source, target)
  765. flow = self.udp_flows.get(key)
  766. if flow is None:
  767. flow_id = next(self.udp_flow_ids)
  768. paths: list[BasePath] = self._build_udp_direct_paths(target, flow_id)
  769. for connection in self._udp_relay_connections():
  770. stream_id = next(self.stream_ids)
  771. paths.append(RelayUdpPath(name=connection.node.name, on_frame=lambda path, event, data, fid=flow_id: self._handle_udp_path(fid, path, event, data), connection=connection, session_id=flow_id, stream_id=stream_id, target=target))
  772. flow = UdpFlow(
  773. flow_id=flow_id,
  774. source=source,
  775. target=target,
  776. send_response=listener.send_response,
  777. paths=paths,
  778. redundancy=self.config.udp_redundancy,
  779. always_broadcast=self.config.udp_always_broadcast,
  780. copy_interval_ms=self.config.udp_copy_interval_ms,
  781. )
  782. self.udp_flows[key] = flow
  783. listener.udp_flows_created += 1
  784. listener._log_udp_summary(force=True)
  785. print(f"[edge] udp flow={flow_id} source={source.host}:{source.port} target={target.host}:{target.port} redundancy={self.config.udp_redundancy} direct_redundancy={self.config.udp_direct_redundancy} always_broadcast={self.config.udp_always_broadcast} candidates={[path.name for path in paths]}")
  786. await flow.start()
  787. await flow.send(payload)
  788. async def _handle_udp_path(self, flow_id: int, path: BasePath, event: str, payload: bytes | None) -> None:
  789. for flow in list(self.udp_flows.values()):
  790. if flow.flow_id == flow_id:
  791. await flow.handle_path(path, event, payload)
  792. break
  793. async def _gc_udp_flows(self) -> None:
  794. loop = asyncio.get_running_loop()
  795. while True:
  796. await asyncio.sleep(30)
  797. now = loop.time()
  798. stale = [key for key, flow in self.udp_flows.items() if flow.last_activity and now - flow.last_activity > 120]
  799. for key in stale:
  800. flow = self.udp_flows.pop(key, None)
  801. if flow:
  802. await flow.close()