Codex
83 lines
2.3 KiB
Python
83 lines
2.3 KiB
Python
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())
|