|
|
@@ -66,6 +66,24 @@ class RelayLink:
|
|
|
await self.writer.wait_closed()
|
|
|
|
|
|
|
|
|
+@dataclass
|
|
|
+class UdpFlowState:
|
|
|
+ flow_id: int
|
|
|
+ client_addr: tuple[str, int]
|
|
|
+ target_host: str
|
|
|
+ target_port: int
|
|
|
+ created_at: float
|
|
|
+ last_activity: float
|
|
|
+ packets_sent: int = 0
|
|
|
+ packets_received: int = 0
|
|
|
+ duplicate_responses: int = 0
|
|
|
+ winner_name: str | None = None
|
|
|
+ candidate_names: tuple[str, ...] = ()
|
|
|
+
|
|
|
+ def touch(self, now: float) -> None:
|
|
|
+ self.last_activity = now
|
|
|
+
|
|
|
+
|
|
|
@dataclass
|
|
|
class TcpRaceSession:
|
|
|
session_id: int
|
|
|
@@ -174,6 +192,11 @@ class UdpAssociateServer(asyncio.DatagramProtocol):
|
|
|
self.client_addr = None
|
|
|
self.packet_counter = itertools.count(1)
|
|
|
self.pending: set[int] = set()
|
|
|
+ self.packet_flows: dict[int, int] = {}
|
|
|
+ self.client_flows: dict[tuple[tuple[str, int], str, int], UdpFlowState] = {}
|
|
|
+ self.flow_counter = itertools.count(1)
|
|
|
+ self.last_summary_at = 0.0
|
|
|
+ self.win_counts: Dict[str, int] = {}
|
|
|
|
|
|
def connection_made(self, transport) -> None:
|
|
|
self.transport = transport
|
|
|
@@ -183,21 +206,105 @@ class UdpAssociateServer(asyncio.DatagramProtocol):
|
|
|
return
|
|
|
if self.client_addr is None:
|
|
|
self.client_addr = addr
|
|
|
+ print(f"[edge] udp client bound addr={addr[0]}:{addr[1]}")
|
|
|
if addr != self.client_addr:
|
|
|
return
|
|
|
host, port, payload = self._parse_socks_udp(data)
|
|
|
+ loop = asyncio.get_running_loop()
|
|
|
+ now = loop.time()
|
|
|
+ flow_key = ((addr[0], addr[1]), host, port)
|
|
|
+ flow = self.client_flows.get(flow_key)
|
|
|
+ if flow is None:
|
|
|
+ flow = UdpFlowState(
|
|
|
+ flow_id=next(self.flow_counter),
|
|
|
+ client_addr=(addr[0], addr[1]),
|
|
|
+ target_host=host,
|
|
|
+ target_port=port,
|
|
|
+ created_at=now,
|
|
|
+ last_activity=now,
|
|
|
+ )
|
|
|
+ self.client_flows[flow_key] = flow
|
|
|
+ flow.touch(now)
|
|
|
+ flow.packets_sent += 1
|
|
|
packet_id = next(self.packet_counter)
|
|
|
self.pending.add(packet_id)
|
|
|
+ self.packet_flows[packet_id] = flow.flow_id
|
|
|
+ print(f"[edge] udp recv flow={flow.flow_id} packet_id={packet_id} target={host}:{port} size={len(payload)}")
|
|
|
asyncio.create_task(self.edge.forward_udp(host, port, payload, packet_id, self))
|
|
|
+ self._log_udp_summary()
|
|
|
|
|
|
- async def handle_from_relay(self, frame: Frame, _link: RelayLink) -> None:
|
|
|
+ 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)
|
|
|
+ flow_id = self.packet_flows.pop(frame.packet_id, 0)
|
|
|
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)
|
|
|
+ winner_log = ""
|
|
|
+ flow = self._find_flow(flow_id)
|
|
|
+ if flow is not None:
|
|
|
+ now = asyncio.get_running_loop().time()
|
|
|
+ flow.touch(now)
|
|
|
+ flow.packets_received += 1
|
|
|
+ if flow.winner_name is None:
|
|
|
+ flow.winner_name = link.node.name
|
|
|
+ self.win_counts[link.node.name] = self.win_counts.get(link.node.name, 0) + 1
|
|
|
+ relay_detail = ", ".join(f"{name}={count}" for name, count in sorted(self.win_counts.items())) or "none"
|
|
|
+ print(
|
|
|
+ f"[edge] udp flow={flow.flow_id} winner={link.node.name} "
|
|
|
+ f"target={flow.target_host}:{flow.target_port} mode=single candidates={len(flow.candidate_names) or len(self.edge.links)}"
|
|
|
+ )
|
|
|
+ print(f"[edge] udp win relay_breakdown={relay_detail}")
|
|
|
+ elif flow.winner_name != link.node.name:
|
|
|
+ flow.duplicate_responses += 1
|
|
|
+ winner_log = f" duplicate=1 winner={flow.winner_name} from={link.node.name}"
|
|
|
+ print(
|
|
|
+ f"[edge] udp send flow={flow_id or 'unknown'} packet_id={frame.packet_id} "
|
|
|
+ f"target={host}:{port} size={len(frame.payload)} relay={link.node.name}{winner_log}"
|
|
|
+ )
|
|
|
self.transport.sendto(packet, self.client_addr)
|
|
|
+ self._log_udp_summary()
|
|
|
+
|
|
|
+ def set_flow_candidates(self, packet_id: int, candidate_names: tuple[str, ...]) -> None:
|
|
|
+ flow_id = self.packet_flows.get(packet_id)
|
|
|
+ flow = self._find_flow(flow_id)
|
|
|
+ if flow is not None and not flow.candidate_names:
|
|
|
+ flow.candidate_names = candidate_names
|
|
|
+
|
|
|
+ def note_unsent(self, packet_id: int) -> None:
|
|
|
+ flow_id = self.packet_flows.pop(packet_id, 0)
|
|
|
+ self.pending.discard(packet_id)
|
|
|
+ flow = self._find_flow(flow_id)
|
|
|
+ if flow is not None:
|
|
|
+ flow.touch(asyncio.get_running_loop().time())
|
|
|
+ print(f"[edge] udp drop flow={flow_id or 'unknown'} packet_id={packet_id} reason=no_available_links")
|
|
|
+ self._log_udp_summary(force=True)
|
|
|
+
|
|
|
+ def _find_flow(self, flow_id: int | None) -> UdpFlowState | None:
|
|
|
+ if not flow_id:
|
|
|
+ return None
|
|
|
+ for flow in self.client_flows.values():
|
|
|
+ if flow.flow_id == flow_id:
|
|
|
+ return flow
|
|
|
+ return None
|
|
|
+
|
|
|
+ def _log_udp_summary(self, force: bool = False) -> None:
|
|
|
+ now = asyncio.get_running_loop().time()
|
|
|
+ if not force and now - self.last_summary_at < 10:
|
|
|
+ return
|
|
|
+ self.last_summary_at = now
|
|
|
+ active_flows = len(self.client_flows)
|
|
|
+ winners = sum(1 for flow in self.client_flows.values() if flow.winner_name)
|
|
|
+ packets_sent = sum(flow.packets_sent for flow in self.client_flows.values())
|
|
|
+ packets_received = sum(flow.packets_received for flow in self.client_flows.values())
|
|
|
+ duplicates = sum(flow.duplicate_responses for flow in self.client_flows.values())
|
|
|
+ print(
|
|
|
+ f"[edge] udp summary bind={self.client_addr[0]}:{self.client_addr[1]} active_flows={active_flows} "
|
|
|
+ f"winner_flows={winners} packets_sent={packets_sent} packets_received={packets_received} dup={duplicates}"
|
|
|
+ if self.client_addr
|
|
|
+ else f"[edge] udp summary bind=unbound active_flows={active_flows} winner_flows={winners} packets_sent={packets_sent} packets_received={packets_received} dup={duplicates}"
|
|
|
+ )
|
|
|
|
|
|
def _parse_socks_udp(self, packet: bytes) -> tuple[str, int, bytes]:
|
|
|
atyp = packet[3]
|
|
|
@@ -261,7 +368,8 @@ class SocksEdge:
|
|
|
|
|
|
async def _accept(self, reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None:
|
|
|
try:
|
|
|
- host, port, udp_mode = await self._handshake(reader, writer)
|
|
|
+ peer = writer.get_extra_info("peername")
|
|
|
+ host, port, udp_mode = await self._handshake(reader, writer, peer)
|
|
|
if udp_mode:
|
|
|
return
|
|
|
links = self._selected_links()
|
|
|
@@ -290,11 +398,17 @@ class SocksEdge:
|
|
|
self.udp_targets[packet_id] = (host, port)
|
|
|
meta = encode_json({"host": host, "port": port})
|
|
|
links = self._selected_links()
|
|
|
+ link_names = ",".join(link.node.name for link in links) or "none"
|
|
|
+ udp_server.set_flow_candidates(packet_id, tuple(link.node.name for link in links))
|
|
|
+ print(f"[edge] udp forward packet_id={packet_id} target={host}:{port} size={len(payload)} links={link_names}")
|
|
|
+ if not links:
|
|
|
+ udp_server.note_unsent(packet_id)
|
|
|
+ return
|
|
|
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]:
|
|
|
+ async def _handshake(self, reader: asyncio.StreamReader, writer: asyncio.StreamWriter, peer) -> tuple[str, int, bool]:
|
|
|
version, methods_len = (await read_exact(reader, 2))
|
|
|
if version != SOCKS_VERSION:
|
|
|
raise ValueError("unsupported socks version")
|
|
|
@@ -312,12 +426,15 @@ class SocksEdge:
|
|
|
else:
|
|
|
raise ValueError("unsupported atyp")
|
|
|
port = struct.unpack("!H", await read_exact(reader, 2))[0]
|
|
|
+ peer_text = f"{peer[0]}:{peer[1]}" if isinstance(peer, tuple) and len(peer) >= 2 else str(peer)
|
|
|
if command == 1:
|
|
|
+ print(f"[edge] socks handshake peer={peer_text} command=connect target={host}:{port}")
|
|
|
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]
|
|
|
+ print(f"[edge] socks handshake peer={peer_text} command=udp_associate target={host}:{port} bind={bind_host}:{bind_port}")
|
|
|
writer.write(b"\x05\x00\x00\x01" + socket.inet_aton(bind_host) + struct.pack("!H", bind_port))
|
|
|
await writer.drain()
|
|
|
return host, port, True
|