Gogs hai 1 semana
pai
achega
c5a531fb1b
Modificáronse 6 ficheiros con 65 adicións e 12 borrados
  1. 2 1
      cli.py
  2. 1 0
      config.json
  3. 24 3
      relay_client.py
  4. 15 4
      relay_server.py
  5. 13 2
      socks_edge.py
  6. 10 2
      transparent_edge.py

+ 2 - 1
cli.py

@@ -122,6 +122,7 @@ def build_parser() -> argparse.ArgumentParser:
     relay.add_argument("--listen-host", default="0.0.0.0")
     relay.add_argument("--listen-host", default="0.0.0.0")
     relay.add_argument("--listen-port", type=int, default=9009)
     relay.add_argument("--listen-port", type=int, default=9009)
     relay.add_argument("--token", required=True)
     relay.add_argument("--token", required=True)
+    relay.add_argument("--udp-only", action="store_true", help="仅支持 UDP relay,不处理 TCP 连接")
     relay.set_defaults(handler=handle_relay)
     relay.set_defaults(handler=handle_relay)
 
 
     edge = sub.add_parser("edge", help="在当前主 VPS 上启动透明 direct 出站加速")
     edge = sub.add_parser("edge", help="在当前主 VPS 上启动透明 direct 出站加速")
@@ -152,7 +153,7 @@ def build_parser() -> argparse.ArgumentParser:
 
 
 
 
 def handle_relay(args: argparse.Namespace) -> int:
 def handle_relay(args: argparse.Namespace) -> int:
-    asyncio.run(RelayServer(args.token).start(args.listen_host, args.listen_port))
+    asyncio.run(RelayServer(args.token).start(args.listen_host, args.listen_port, udp_only=args.udp_only))
     return 0
     return 0
 
 
 
 

+ 1 - 0
config.json

@@ -13,5 +13,6 @@
   "socks_host": "127.0.0.1",
   "socks_host": "127.0.0.1",
   "socks_port": 19180,
   "socks_port": 19180,
   "relays": [
   "relays": [
+    {"name": "hk2", "host": "23.95.134.159", "port": 9009, "token": "130", "weight": 100}
   ]
   ]
 }
 }

+ 24 - 3
relay_client.py

@@ -2,6 +2,7 @@ from __future__ import annotations
 
 
 import asyncio
 import asyncio
 import contextlib
 import contextlib
+import json
 import random
 import random
 import socket
 import socket
 from dataclasses import dataclass
 from dataclasses import dataclass
@@ -31,6 +32,8 @@ class RelayConnection:
     closed_event: asyncio.Event | None = None
     closed_event: asyncio.Event | None = None
     dropped_frames: Dict[int, int] = None
     dropped_frames: Dict[int, int] = None
     dropped_report_task: asyncio.Task | None = None
     dropped_report_task: asyncio.Task | None = None
+    supports_tcp: bool = True
+    supports_udp: bool = True
 
 
     def __post_init__(self) -> None:
     def __post_init__(self) -> None:
         if self.handlers is None:
         if self.handlers is None:
@@ -50,10 +53,21 @@ class RelayConnection:
         frame = await read_frame(self.reader)
         frame = await read_frame(self.reader)
         if frame.kind != AUTH or frame.packet_id != STATUS_OK:
         if frame.kind != AUTH or frame.packet_id != STATUS_OK:
             raise ConnectionError(f"relay auth failed: {self.node.name}")
             raise ConnectionError(f"relay auth failed: {self.node.name}")
+        ack = {}
+        if frame.payload:
+            try:
+                ack = json.loads(frame.payload.decode("utf-8"))
+            except Exception:
+                ack = {}
+        self.supports_tcp = not bool(ack.get("udp_only", False))
+        self.supports_udp = True
         self.last_pong_at = time.monotonic()
         self.last_pong_at = time.monotonic()
         self.keepalive_task = asyncio.create_task(self._keepalive())
         self.keepalive_task = asyncio.create_task(self._keepalive())
         self.pump_task = asyncio.create_task(self._pump())
         self.pump_task = asyncio.create_task(self._pump())
-        print(f"[edge] relay connected name={self.node.name} addr={self.node.host}:{self.node.port}")
+        print(
+            f"[edge] relay connected name={self.node.name} addr={self.node.host}:{self.node.port} "
+            f"supports_tcp={self.supports_tcp} supports_udp={self.supports_udp}"
+        )
 
 
     async def _keepalive(self) -> None:
     async def _keepalive(self) -> None:
         try:
         try:
@@ -236,10 +250,17 @@ class RelayManager:
 
 
     def available(self) -> list[RelayConnection]:
     def available(self) -> list[RelayConnection]:
         chosen = {node.name for node in self.scheduler.choose()}
         chosen = {node.name for node in self.scheduler.choose()}
-        preferred = [self.connections[name] for name in chosen if name in self.connections and not self.connections[name].closed]
+        preferred = [self.connections[name] for name in chosen if name in self.connections and not self.connections[name].closed and self.connections[name].supports_tcp]
         if preferred:
         if preferred:
             return preferred
             return preferred
-        return [conn for conn in self.connections.values() if not conn.closed]
+        return [conn for conn in self.connections.values() if not conn.closed and conn.supports_tcp]
+
+    def available_udp(self) -> list[RelayConnection]:
+        chosen = {node.name for node in self.scheduler.choose()}
+        preferred = [self.connections[name] for name in chosen if name in self.connections and not self.connections[name].closed and self.connections[name].supports_udp]
+        if preferred:
+            return preferred
+        return [conn for conn in self.connections.values() if not conn.closed and conn.supports_udp]
 
 
     def snapshot(self) -> list[dict[str, object]]:
     def snapshot(self) -> list[dict[str, object]]:
         data = self.scheduler.snapshot()
         data = self.scheduler.snapshot()

+ 15 - 4
relay_server.py

@@ -69,6 +69,7 @@ class RelayChannel:
     authed_at: float = 0.0
     authed_at: float = 0.0
     frame_count: int = 0
     frame_count: int = 0
     authed_kind: str = "normal"
     authed_kind: str = "normal"
+    udp_only: bool = False
 
 
     async def run(self) -> None:
     async def run(self) -> None:
         peer = self.writer.get_extra_info("peername")
         peer = self.writer.get_extra_info("peername")
@@ -86,7 +87,7 @@ class RelayChannel:
             authed = True
             authed = True
             self.authed_at = time.monotonic()
             self.authed_at = time.monotonic()
             self.authed_kind = payload.get("purpose", "normal")
             self.authed_kind = payload.get("purpose", "normal")
-            ack_payload = {"status": "ok", "kind": self.authed_kind}
+            ack_payload = {"status": "ok", "kind": self.authed_kind, "udp_only": self.udp_only}
             await self.safe_send(Frame(AUTH, 0, 0, 0, STATUS_OK, encode_json(ack_payload)))
             await self.safe_send(Frame(AUTH, 0, 0, 0, STATUS_OK, encode_json(ack_payload)))
             while True:
             while True:
                 frame = await read_frame(self.reader)
                 frame = await read_frame(self.reader)
@@ -126,6 +127,9 @@ class RelayChannel:
         if frame.kind == AUTH:
         if frame.kind == AUTH:
             return
             return
         if frame.kind == TCP_OPEN:
         if frame.kind == TCP_OPEN:
+            if self.udp_only:
+                await self.safe_send(Frame(TCP_STATUS, frame.session_id, frame.stream_id, 0, STATUS_ERR, b"tcp disabled on udp-only relay"))
+                return
             try:
             try:
                 meta = decode_json(frame.payload) if frame.payload else {}
                 meta = decode_json(frame.payload) if frame.payload else {}
                 family = int(meta.get("family", 0)) or 0
                 family = int(meta.get("family", 0)) or 0
@@ -137,6 +141,8 @@ class RelayChannel:
                 await self.safe_send(Frame(TCP_STATUS, frame.session_id, frame.stream_id, 0, STATUS_ERR, str(exc).encode()))
                 await self.safe_send(Frame(TCP_STATUS, frame.session_id, frame.stream_id, 0, STATUS_ERR, str(exc).encode()))
             return
             return
         if frame.kind == TCP_DATA:
         if frame.kind == TCP_DATA:
+            if self.udp_only:
+                return
             session = self.tcp_sessions.get(key)
             session = self.tcp_sessions.get(key)
             if session:
             if session:
                 try:
                 try:
@@ -146,6 +152,8 @@ class RelayChannel:
                     await self._close_tcp(key)
                     await self._close_tcp(key)
             return
             return
         if frame.kind == TCP_CLOSE:
         if frame.kind == TCP_CLOSE:
+            if self.udp_only:
+                return
             await self._close_tcp(key)
             await self._close_tcp(key)
             return
             return
         if frame.kind == UDP_SEND:
         if frame.kind == UDP_SEND:
@@ -249,13 +257,16 @@ class RelayChannel:
 class RelayServer:
 class RelayServer:
     def __init__(self, token: str) -> None:
     def __init__(self, token: str) -> None:
         self.token = token
         self.token = token
+        self.udp_only = False
 
 
-    async def start(self, host: str, port: int) -> None:
+    async def start(self, host: str, port: int, udp_only: bool = False) -> None:
+        self.udp_only = udp_only
         server = await asyncio.start_server(self._accept, host, port)
         server = await asyncio.start_server(self._accept, host, port)
         sockets = ", ".join(str(sock.getsockname()) for sock in server.sockets or [])
         sockets = ", ".join(str(sock.getsockname()) for sock in server.sockets or [])
-        print(f"[relay] listening on {sockets}")
+        mode = "udp-only" if udp_only else "normal"
+        print(f"[relay] listening on {sockets} mode={mode}")
         async with server:
         async with server:
             await server.serve_forever()
             await server.serve_forever()
 
 
     async def _accept(self, reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None:
     async def _accept(self, reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None:
-        await RelayChannel(reader, writer, self.token).run()
+        await RelayChannel(reader, writer, self.token, udp_only=self.udp_only).run()

+ 13 - 2
socks_edge.py

@@ -4,6 +4,7 @@ import asyncio
 import contextlib
 import contextlib
 import itertools
 import itertools
 from collections import deque
 from collections import deque
+import json
 import socket
 import socket
 import struct
 import struct
 from dataclasses import dataclass, field
 from dataclasses import dataclass, field
@@ -34,12 +35,22 @@ class RelayLink:
     tcp_sessions: Dict[tuple[int, int], "TcpRaceSession"] = field(default_factory=dict)
     tcp_sessions: Dict[tuple[int, int], "TcpRaceSession"] = field(default_factory=dict)
     udp_server: "UdpAssociateServer | None" = None
     udp_server: "UdpAssociateServer | None" = None
     closed: bool = False
     closed: bool = False
+    supports_tcp: bool = True
+    supports_udp: bool = True
 
 
     async def start(self) -> None:
     async def start(self) -> None:
         await write_frame(self.writer, Frame(AUTH, 0, 0, 0, 0, encode_json({"token": self.node.token})))
         await write_frame(self.writer, Frame(AUTH, 0, 0, 0, 0, encode_json({"token": self.node.token})))
         frame = await read_frame(self.reader)
         frame = await read_frame(self.reader)
         if frame.kind != AUTH or frame.packet_id != STATUS_OK:
         if frame.kind != AUTH or frame.packet_id != STATUS_OK:
             raise ConnectionError(f"relay auth failed: {self.node.name}")
             raise ConnectionError(f"relay auth failed: {self.node.name}")
+        ack = {}
+        if frame.payload:
+            try:
+                ack = json.loads(frame.payload.decode("utf-8"))
+            except Exception:
+                ack = {}
+        self.supports_tcp = not bool(ack.get("udp_only", False))
+        self.supports_udp = True
         self.closed = False
         self.closed = False
         self.closed_event.clear()
         self.closed_event.clear()
         self.pump = asyncio.create_task(self._pump())
         self.pump = asyncio.create_task(self._pump())
@@ -480,11 +491,11 @@ class SocksEdge:
 
 
     def _selected_links(self) -> list[RelayLink]:
     def _selected_links(self) -> list[RelayLink]:
         chosen = {node.name for node in self.scheduler.choose()}
         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]
+        links = [link for link in self.links if link.node.name in chosen and not link.closed and link.supports_tcp]
         return links or [link for link in self.links if not link.closed][:1]
         return links or [link for link in self.links if not link.closed][:1]
 
 
     def _selected_udp_links(self) -> list[RelayLink]:
     def _selected_udp_links(self) -> list[RelayLink]:
-        online = [link for link in self.links if not link.closed and link.writer is not None]
+        online = [link for link in self.links if not link.closed and link.writer is not None and link.supports_udp]
         if not online:
         if not online:
             return []
             return []
         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)
         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)

+ 10 - 2
transparent_edge.py

@@ -799,6 +799,14 @@ class TransparentEdge:
             for index in range(count)
             for index in range(count)
         ]
         ]
 
 
+    def _tcp_relay_connections(self) -> list[RelayConnection]:
+        return [connection for connection in self.manager.available() if connection.supports_tcp]
+
+    def _udp_relay_connections(self) -> list[RelayConnection]:
+        if hasattr(self.manager, "available_udp"):
+            return [connection for connection in self.manager.available_udp() if connection.supports_udp]
+        return [connection for connection in self.manager.available() if connection.supports_udp]
+
     def _start_udp_listeners(self) -> None:
     def _start_udp_listeners(self) -> None:
         binds = []
         binds = []
         if self.listen_host == "127.0.0.1":
         if self.listen_host == "127.0.0.1":
@@ -823,7 +831,7 @@ class TransparentEdge:
             session_id = next(self.session_ids)
             session_id = next(self.session_ids)
             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)
             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)
             paths: list[BasePath] = self._build_direct_paths(session)
             paths: list[BasePath] = self._build_direct_paths(session)
-            for connection in self.manager.available():
+            for connection in self._tcp_relay_connections():
                 stream_id = next(self.stream_ids)
                 stream_id = next(self.stream_ids)
                 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))
                 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))
             session.paths = paths
             session.paths = paths
@@ -861,7 +869,7 @@ class TransparentEdge:
         if flow is None:
         if flow is None:
             flow_id = next(self.udp_flow_ids)
             flow_id = next(self.udp_flow_ids)
             paths: list[BasePath] = self._build_udp_direct_paths(target, flow_id)
             paths: list[BasePath] = self._build_udp_direct_paths(target, flow_id)
-            for connection in self.manager.available():
+            for connection in self._udp_relay_connections():
                 stream_id = next(self.stream_ids)
                 stream_id = next(self.stream_ids)
                 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))
                 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))
             flow = UdpFlow(
             flow = UdpFlow(