122 lines
4.1 KiB
Python
122 lines
4.1 KiB
Python
#!/usr/bin/env python3
|
||
"""RRPP 裸 UoT v2 relay。
|
||
|
||
nginx 的 proxy_connect 模块负责 TLS、Basic Auth 和 CONNECT，relay 只监听
|
||
127.0.0.1:15201，收到的首字节就是 UoT v2 request，不再解析 HTTP。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import ipaddress
|
||
import logging
|
||
import socket
|
||
import struct
|
||
import threading
|
||
|
||
MAGIC = "sp.v2.udp-over-tcp.arpa"
|
||
LISTEN_HOST = "127.0.0.1"
|
||
LISTEN_PORT = 15201
|
||
MAX_PACKET = 65535
|
||
IDLE_TIMEOUT = 9 * 60
|
||
|
||
|
||
def read_exact(conn: socket.socket, size: int) -> bytes:
|
||
data = bytearray()
|
||
while len(data) < size:
|
||
chunk = conn.recv(size - len(data))
|
||
if not chunk:
|
||
raise ConnectionError("客户端连接已关闭")
|
||
data.extend(chunk)
|
||
return bytes(data)
|
||
|
||
|
||
def read_address(conn: socket.socket):
|
||
atyp = read_exact(conn, 1)[0]
|
||
if atyp == 0:
|
||
host = str(ipaddress.IPv4Address(read_exact(conn, 4)))
|
||
elif atyp == 1:
|
||
host = str(ipaddress.IPv6Address(read_exact(conn, 16)))
|
||
elif atyp == 2:
|
||
size = read_exact(conn, 1)[0]
|
||
host = read_exact(conn, size).decode("idna")
|
||
else:
|
||
raise ValueError(f"未知 UoT 地址类型: {atyp}")
|
||
return host, struct.unpack("!H", read_exact(conn, 2))[0]
|
||
|
||
|
||
def encode_address(host: str, port: int) -> bytes:
|
||
try:
|
||
addr = ipaddress.ip_address(host)
|
||
except ValueError:
|
||
raw = host.encode("idna")
|
||
return b"\x02" + bytes([len(raw)]) + raw + struct.pack("!H", port)
|
||
return (b"\x00" if addr.version == 4 else b"\x01") + addr.packed + struct.pack("!H", port)
|
||
|
||
|
||
def resolve_destination(host: str, port: int):
|
||
return socket.getaddrinfo(host, port, socket.AF_UNSPEC, socket.SOCK_DGRAM)[0][4]
|
||
|
||
|
||
def handle(conn: socket.socket, peer):
|
||
udp = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
|
||
udp.settimeout(0.2)
|
||
# 与 nginx 的 10 分钟读写超时错开，空闲连接由 relay 先回收。
|
||
conn.settimeout(IDLE_TIMEOUT)
|
||
send_lock = threading.Lock()
|
||
try:
|
||
is_connect = read_exact(conn, 1)[0]
|
||
first_target = read_address(conn)
|
||
logging.info("UoT client %s target=%s:%s connect=%s", peer, *first_target, is_connect)
|
||
|
||
def udp_reader():
|
||
while True:
|
||
try:
|
||
data, source = udp.recvfrom(MAX_PACKET)
|
||
except socket.timeout:
|
||
continue
|
||
except OSError:
|
||
return
|
||
try:
|
||
with send_lock:
|
||
if is_connect:
|
||
conn.sendall(struct.pack("!H", len(data)) + data)
|
||
else:
|
||
conn.sendall(encode_address(source[0], source[1]) +
|
||
struct.pack("!H", len(data)) + data)
|
||
except OSError:
|
||
return
|
||
|
||
threading.Thread(target=udp_reader, daemon=True).start()
|
||
if is_connect:
|
||
while True:
|
||
length = struct.unpack("!H", read_exact(conn, 2))[0]
|
||
payload = read_exact(conn, length)
|
||
udp.sendto(payload, resolve_destination(*first_target))
|
||
else:
|
||
while True:
|
||
destination = read_address(conn)
|
||
length = struct.unpack("!H", read_exact(conn, 2))[0]
|
||
payload = read_exact(conn, length)
|
||
udp.sendto(payload, resolve_destination(*destination))
|
||
except (ConnectionError, OSError, ValueError) as exc:
|
||
logging.info("UoT client %s closed: %s", peer, exc)
|
||
finally:
|
||
udp.close()
|
||
conn.close()
|
||
|
||
|
||
def main():
|
||
logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
|
||
listener = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
|
||
listener.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
|
||
listener.bind((LISTEN_HOST, LISTEN_PORT))
|
||
listener.listen(128)
|
||
logging.info("RRPP UoT relay listening on %s:%d", LISTEN_HOST, LISTEN_PORT)
|
||
while True:
|
||
conn, peer = listener.accept()
|
||
threading.Thread(target=handle, args=(conn, peer), daemon=True).start()
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|