Add local proxy relay helper

2026-05-07 01:34:45 +08:00
parent 3dc4214b78
commit 28f0a946a5
1 changed files with 82 additions and 0 deletions
@@ -0,0 +1,82 @@
+from pathlib import Path
+import os
+import socket
+import sys
+import threading
+
+
+ROOT = Path(__file__).resolve().parents[1]
+if str(ROOT) not in sys.path:
+    sys.path.insert(0, str(ROOT))
+
+
+DEFAULT_BIND_HOST = "0.0.0.0"
+DEFAULT_BIND_PORT = 18080
+DEFAULT_UPSTREAM_HOST = "127.0.0.1"
+DEFAULT_UPSTREAM_PORT = 10808
+
+
+def main(env=None) -> int:
+    env = os.environ if env is None else env
+    bind_host = str(env.get("AI_ICT_PROXY_RELAY_BIND_HOST") or DEFAULT_BIND_HOST)
+    bind_port = int(env.get("AI_ICT_PROXY_RELAY_BIND_PORT") or DEFAULT_BIND_PORT)
+    upstream_host = str(env.get("AI_ICT_PROXY_RELAY_UPSTREAM_HOST") or DEFAULT_UPSTREAM_HOST)
+    upstream_port = int(env.get("AI_ICT_PROXY_RELAY_UPSTREAM_PORT") or DEFAULT_UPSTREAM_PORT)
+
+    server = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+    server.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+    server.bind((bind_host, bind_port))
+    server.listen(128)
+    print(
+        f"proxy relay listening on {bind_host}:{bind_port} -> {upstream_host}:{upstream_port}",
+        flush=True,
+    )
+
+    while True:
+        client, _ = server.accept()
+        worker = threading.Thread(
+            target=_handle_connection,
+            args=(client, upstream_host, upstream_port),
+            daemon=True,
+        )
+        worker.start()
+
+
+def _handle_connection(client: socket.socket, upstream_host: str, upstream_port: int) -> None:
+    try:
+        upstream = socket.create_connection((upstream_host, upstream_port), timeout=10)
+    except Exception:
+        _close_socket(client)
+        return
+
+    threading.Thread(target=_pump, args=(client, upstream), daemon=True).start()
+    threading.Thread(target=_pump, args=(upstream, client), daemon=True).start()
+
+
+def _pump(source: socket.socket, target: socket.socket) -> None:
+    try:
+        while True:
+            data = source.recv(65536)
+            if not data:
+                break
+            target.sendall(data)
+    except Exception:
+        pass
+    finally:
+        _close_socket(source)
+        _close_socket(target)
+
+
+def _close_socket(sock: socket.socket) -> None:
+    try:
+        sock.shutdown(socket.SHUT_RDWR)
+    except Exception:
+        pass
+    try:
+        sock.close()
+    except Exception:
+        pass
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())