#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
RRPP UoT v2 relay — UDP over TCP 解封装服务（sing-box UoT v2 协议）

协议（SagerNet 专有，见 sing-box 文档 configuration/shared/udp-over-tcp）：

  请求头:  | isConnect u8 | ATYP u8 | address | port u16be |
    isConnect=1: 固定目标模式，后续帧全部是 length+data，发往请求头里的目标
    isConnect=0: 每帧自带目标，帧格式同 v1: | ATYP | address | port | length u16be | data |

  connect 帧:      | length u16be | data |
  non-connect 帧:  | ATYP | address | port | length u16be | data |

监听 127.0.0.1，仅供本机 nginx proxy_connect 转发，不对公网开放。

健壮性设计（2026-09 修复，实测复现过隧道被误杀）：
  * 一律使用「未连接」UDP socket（sendto/recvfrom）——连接式 socket 会把目标侧
    ICMP（port unreachable 等）转成 ConnectionRefusedError，在 send/recv 上抛异常，
    旧实现因此把整条 UoT 隧道关掉（客户端表现为 "UoT TCP 流已关闭"）。
  * 数据报级错误（BlockingIOError/EAGAIN、ENETUNREACH、EHOSTUNREACH、EMSGSIZE
    等）只丢当前包并继续；EAGAIN 先等可写再重试（有总时限）。
  * 只有 TCP 侧 EOF、UoT 协议错误或空闲超时才关闭隧道。
  * 按地址族惰性建 socket，non-connect 模式可混用 IPv4/IPv6 目标。
"""
import errno
import select
import socket
import struct
import threading
import time

LISTEN_HOST = "127.0.0.1"
LISTEN_PORT = 15201
IDLE_TIMEOUT = 540          # 9 分钟无流量断开（须 < nginx proxy_connect 隧道超时 10m）
MAX_DGRAM = 65507           # UDP over IPv4 单包上限
SEND_RETRY_TOTAL = 1.0      # EAGAIN 时最多等待 1 秒可写，超时则丢弃该包
RECV_ERROR_STREAK_LIMIT = 200   # 连续 recv 错误上限（防病态自旋），超过视为隧道异常
LOG_PREFIX = "[uot-relay]"


def log(msg):
    print(f"{LOG_PREFIX} {time.strftime('%Y-%m-%d %H:%M:%S')} {msg}", flush=True)


def read_exact(f, n):
    buf = b""
    while len(buf) < n:
        chunk = f.read(n - len(buf))
        if not chunk:
            raise EOFError("connection closed")
        buf += chunk
    return buf


def parse_addr(f):
    """从流中读取 SOCKS 格式地址（ATYP + address + port）。"""
    atyp = read_exact(f, 1)[0]
    if atyp == 0x00:
        addr = socket.inet_ntop(socket.AF_INET, read_exact(f, 4))
    elif atyp == 0x01:
        addr = socket.inet_ntop(socket.AF_INET6, read_exact(f, 16))
    elif atyp == 0x02:
        ln = read_exact(f, 1)[0]
        addr = read_exact(f, ln).decode("utf-8", "replace")
    else:
        raise ValueError(f"bad ATYP {atyp:#x}")
    port = struct.unpack(">H", read_exact(f, 2))[0]
    return addr, port


def encode_addr(addr, port):
    """SOCKS 格式地址编码（用于 non-connect 模式的回包）。"""
    try:
        return b"\x00" + socket.inet_pton(socket.AF_INET, addr) + struct.pack(">H", port)
    except OSError:
        pass
    try:
        return b"\x01" + socket.inet_pton(socket.AF_INET6, addr) + struct.pack(">H", port)
    except OSError:
        pass
    d = addr.encode("utf-8")
    if len(d) > 255:
        raise ValueError("domain too long")
    return b"\x02" + bytes([len(d)]) + d + struct.pack(">H", port)


def norm_src(src):
    """recvfrom 地址归一化：IPv4 为 (host, port)，IPv6 为 4 元组。"""
    return src[0], src[1]


class UdpPool:
    """按地址族惰性创建「未连接」UDP socket（v4/v6 可混用）。"""

    def __init__(self):
        self.socks = {}

    def get(self, addr):
        fam = socket.AF_INET6 if ":" in addr else socket.AF_INET
        s = self.socks.get(fam)
        if s is None:
            s = socket.socket(fam, socket.SOCK_DGRAM)
            s.setblocking(False)
            self.socks[fam] = s
        return s

    def all(self):
        return list(self.socks.values())

    def close(self):
        for s in self.socks.values():
            try:
                s.close()
            except Exception:
                pass
        self.socks.clear()


def send_dgram(sock, data, dst):
    """发送一个 UDP 数据报。

    返回 True=已发出，False=本包丢弃（**不是**连接级错误，调用方继续跑）。
    EAGAIN 时等可写重试（有总时限）；其余 OSError 一律丢弃本包。
    """
    deadline = time.time() + SEND_RETRY_TOTAL
    while True:
        try:
            sock.sendto(data, dst)
            return True
        except BlockingIOError:
            if time.time() >= deadline:
                return False
            try:
                select.select([], [sock], [], 0.1)
            except Exception:
                return False
        except InterruptedError:
            continue
        except OSError:
            # ENETUNREACH / EHOSTUNREACH / EMSGSIZE / EACCES ...
            return False


def handle(conn, peer):
    conn.settimeout(IDLE_TIMEOUT)
    f = conn.makefile("rb")
    pool = UdpPool()
    stop = threading.Event()
    fixed_dst = None
    try:
        is_connect = read_exact(f, 1)[0]
        req_dst = parse_addr(f)
        log(f"conn from nginx: isConnect={is_connect} dst={req_dst}")
        if is_connect:
            fixed_dst = req_dst
            pool.get(fixed_dst[0])   # 预建 socket：主循环首轮才有可 select 的对象

        write_lock = threading.Lock()

        def tcp_write(data):
            with write_lock:
                conn.sendall(data)

        def uplink():
            """TCP -> UDP：逐帧解析并发送；单包错误只丢包，不关隧道。"""
            try:
                while not stop.is_set():
                    if is_connect:
                        ln = struct.unpack(">H", read_exact(f, 2))[0]
                        if ln > MAX_DGRAM:
                            raise ValueError(f"datagram too large: {ln}")
                        data = read_exact(f, ln)
                        send_dgram(pool.get(fixed_dst[0]), data, fixed_dst)
                    else:
                        dst = parse_addr(f)
                        ln = struct.unpack(">H", read_exact(f, 2))[0]
                        if ln > MAX_DGRAM:
                            raise ValueError(f"datagram too large: {ln}")
                        data = read_exact(f, ln)
                        send_dgram(pool.get(dst[0]), data, dst)
            except Exception as e:
                if not stop.is_set():
                    log(f"uplink end: {type(e).__name__}: {e}")
            finally:
                stop.set()

        threading.Thread(target=uplink, daemon=True).start()

        # 主循环：UDP -> TCP（带空闲超时）
        last = time.time()
        err_streak = 0
        while not stop.is_set():
            socks = pool.all()
            if not socks:
                # non-connect 模式：socket 由第一帧触发创建，空池时等待而非关连接
                if time.time() - last > IDLE_TIMEOUT:
                    log("idle timeout (no datagram), closing")
                    break
                time.sleep(0.2)
                continue
            try:
                r, _, _ = select.select(socks, [], [], 2.0)
            except OSError:
                break
            got = False
            for s in r:
                try:
                    data, src = s.recvfrom(MAX_DGRAM)
                except (BlockingIOError, InterruptedError):
                    continue
                except OSError as e:
                    # 单包级错误（含 ICMP 触发的错误）：丢包继续，不关隧道
                    err_streak += 1
                    if err_streak % 50 == 1:
                        log(f"recv error (datagram dropped, streak={err_streak}): "
                            f"{errno.errorcode.get(e.errno, e.errno)} {e}")
                    if err_streak >= RECV_ERROR_STREAK_LIMIT:
                        log("recv error streak limit reached, closing tunnel")
                        stop.set()
                        break
                    continue
                err_streak = 0
                a, p = norm_src(src)
                if is_connect:
                    if (a, p) != fixed_dst:
                        continue      # 固定目标模式：只回传目标自身的包
                    tcp_write(struct.pack(">H", len(data)) + data)
                else:
                    tcp_write(encode_addr(a, p) + struct.pack(">H", len(data)) + data)
                got = True
            if got:
                last = time.time()
            elif time.time() - last > IDLE_TIMEOUT:
                log("idle timeout, closing")
                break
    except Exception as e:
        log(f"conn error: {type(e).__name__}: {e}")
    finally:
        stop.set()
        try:
            f.close()
        except Exception:
            pass
        try:
            conn.close()
        except Exception:
            pass
        pool.close()


def main():
    srv = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    srv.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    srv.bind((LISTEN_HOST, LISTEN_PORT))
    srv.listen(128)
    log(f"listening on {LISTEN_HOST}:{LISTEN_PORT}")
    while True:
        conn, peer = srv.accept()
        threading.Thread(target=handle, args=(conn, peer), daemon=True).start()


if __name__ == "__main__":
    main()