Enter usocksHEADmaster

SOCKS5 CONNECT proxy for Linux 6.1 or newer: C99, raw io_uring, libc
only. Each worker process runs one event loop with fixed 64 KB buffers
per direction; -w N workers share the port through SO_REUSEPORT.
IPv4, IPv6 and domain names, RFC 1929 authentication, TCP or unix
socket listener.

1 files changed, 323 insertions, 0 deletions

diff --git a/test b/test
new file mode 100755
index 0000000..5e3ef5a
--- /dev/null
+++ b/test
@@ -0,0 +1,323 @@
+#!/usr/bin/env python3
+import hashlib
+import os
+import signal
+import socket
+import struct
+import subprocess
+import sys
+import tempfile
+import threading
+import time
+
+failed = 0
+procs = []
+
+
+def check(name, cond):
+    global failed
+    print(("ok   " if cond else "FAIL ") + name)
+    failed += not cond
+
+
+def family(host):
+    return socket.AF_INET6 if ":" in host else socket.AF_INET
+
+
+def free_port(host="127.0.0.1"):
+    s = socket.socket(family(host))
+    s.bind((host, 0))
+    port = s.getsockname()[1]
+    s.close()
+    return port
+
+
+def start(tmp, listen, *args, lines=1):
+    log = os.path.join(tmp, listen.replace("/", "_") + ".log")
+    with open(log, "w") as f:
+        p = subprocess.Popen(["./usocks", "-l", listen, *args], stderr=f)
+    procs.append(p)
+    for _ in range(100):
+        with open(log) as f:
+            if f.read().count("listening on") >= lines:
+                return p
+        if p.poll() is not None:
+            break
+        time.sleep(0.05)
+    sys.exit(f"usocks -l {listen} did not start, see {log}")
+
+
+def echo_server(host="127.0.0.1"):
+    ls = socket.socket(family(host))
+    ls.bind((host, 0))
+    ls.listen(64)
+
+    def serve(c):
+        while b := c.recv(65536):
+            c.sendall(b)
+        c.shutdown(socket.SHUT_WR)
+        c.close()
+
+    def loop():
+        while True:
+            c, _ = ls.accept()
+            threading.Thread(target=serve, args=(c,), daemon=True).start()
+
+    threading.Thread(target=loop, daemon=True).start()
+    return ls.getsockname()[1]
+
+
+def recvn(s, n):
+    b = b""
+    while len(b) < n:
+        x = s.recv(n - len(b))
+        if not x:
+            raise EOFError
+        b += x
+    return b
+
+
+def recvall(s):
+    b = b""
+    while x := s.recv(65536):
+        b += x
+    return b
+
+
+def request(host, port, cmd=1):
+    for af, atyp in ((socket.AF_INET, 1), (socket.AF_INET6, 4)):
+        try:
+            a = bytes([atyp]) + socket.inet_pton(af, host)
+            break
+        except OSError:
+            pass
+    else:
+        a = bytes([3, len(host)]) + host.encode()
+    return bytes([5, cmd, 0]) + a + struct.pack(">H", port)
+
+
+def reply(s):
+    h = recvn(s, 4)
+    recvn(s, {1: 4, 4: 16}[h[3]] + 2)
+    return h[1]
+
+
+def proxy(addr):
+    if isinstance(addr, str):
+        s = socket.socket(socket.AF_UNIX)
+        s.connect(addr)
+    elif isinstance(addr, tuple):
+        s = socket.create_connection(addr)
+    else:
+        s = socket.create_connection(("127.0.0.1", addr))
+    s.settimeout(30)
+    return s
+
+
+def connect(addr, host, port, auth=None, extra=b"", cmd=1):
+    """Send the whole handshake in one write; return socket and reply code."""
+    s = proxy(addr)
+    if auth:
+        u, p = auth
+        msg = b"\5\1\2" + bytes([1, len(u)]) + u + bytes([len(p)]) + p
+    else:
+        msg = b"\5\1\0"
+    s.sendall(msg + request(host, port, cmd) + extra)
+    if recvn(s, 2) != (b"\5\2" if auth else b"\5\0"):
+        return s, -1
+    if auth and recvn(s, 2) != b"\1\0":
+        return s, -2
+    return s, reply(s)
+
+
+def nfds(pid):
+    return len(os.listdir(f"/proc/{pid}/fd"))
+
+
+def proc(pid):
+    try:
+        with open(f"/proc/{pid}/stat") as f:
+            s = f.read().rsplit(")", 1)[1].split()
+    except OSError:
+        return None
+    return s[0], int(s[1])
+
+
+def alive(pid):
+    p = proc(pid)
+    return p is not None and p[0] != "Z"
+
+
+def workers(pid):
+    ps = [int(d) for d in os.listdir("/proc") if d.isdigit()]
+    return [k for k in ps if alive(k) and proc(k)[1] == pid]
+
+
+def main():
+    os.chdir(os.path.dirname(os.path.abspath(__file__)))
+    echo = echo_server()
+    tmp = tempfile.mkdtemp()
+    unix = os.path.join(tmp, "usocks.sock")
+    port, aport, mport, wport = (free_port() for _ in range(4))
+    plain = start(tmp, f"127.0.0.1:{port}")
+    start(tmp, f"127.0.0.1:{aport}", "-u", "bob", "-p", "secret")
+    start(tmp, f"127.0.0.1:{mport}", "-m", "2")
+    start(tmp, unix)
+    fds = nfds(plain.pid)
+
+    s, r = connect(port, "127.0.0.1", echo)
+    s.sendall(b"hello")
+    check("connect ipv4", r == 0 and recvn(s, 5) == b"hello")
+    s.close()
+
+    try:
+        echo6, port6 = echo_server("::1"), free_port("::1")
+    except OSError:
+        print("skip ipv6: ::1 not available")
+    else:
+        start(tmp, f"[::1]:{port6}")
+        s, r = connect(("::1", port6), "::1", echo6)
+        s.sendall(b"six")
+        check("ipv6 listener and destination", r == 0 and recvn(s, 3) == b"six")
+        s.close()
+
+    s, r = connect(port, "localhost", echo, extra=b"early")
+    check("domain name with early data", r == 0 and recvn(s, 5) == b"early")
+    s.close()
+
+    s, r = connect(port, "nonexistent.invalid", echo)
+    check("unresolvable name gets 4", r == 4)
+    s, r = connect(port, "127.0.0.1", free_port())
+    check("refused gets 5", r == 5)
+    s, r = connect(port, "evil\nlog.example", echo)
+    check("non-printable domain name gets 1", r == 1)
+    s, r = connect(port, "127.0.0.1", echo, cmd=2)
+    check("bind gets 7", r == 7)
+    s = proxy(port)
+    s.sendall(b"\5\1\0\5\1\0\7")
+    recvn(s, 2)
+    check("bad address type gets 8", reply(s) == 8)
+    s = proxy(port)
+    s.sendall(b"\4\1")
+    check("bad version closes", recvall(s) == b"")
+    s = proxy(port)
+    s.sendall(b"\5\1\2")
+    check("no acceptable method", recvn(s, 2) == b"\5\377")
+
+    s = proxy(aport)
+    for b in b"\5\1\2\1\3bob\6secret" + request("127.0.0.1", echo) + b"x":
+        s.send(bytes([b]))
+        time.sleep(0.002)
+    ok = recvn(s, 4) == b"\5\2\1\0" and reply(s) == 0
+    check("byte by byte handshake with auth", ok and recvn(s, 1) == b"x")
+    s.close()
+    s, r = connect(aport, "127.0.0.1", echo, auth=(b"bob", b"secret"))
+    check("auth accepted", r == 0)
+    s, r = connect(aport, "127.0.0.1", echo, auth=(b"bob", b"secre"))
+    check("wrong password rejected", r == -2)
+    s, r = connect(aport, "127.0.0.1", echo)
+    check("no auth offered to auth server", r == -1)
+
+    s, r = connect(port, "127.0.0.1", echo)
+    s.sendall(b"abc")
+    s.shutdown(socket.SHUT_WR)
+    check("client half-close", recvall(s) == b"abc")
+    ls = socket.socket()
+    ls.bind(("127.0.0.1", 0))
+    ls.listen(1)
+    s, r = connect(port, "127.0.0.1", ls.getsockname()[1])
+    a, _ = ls.accept()
+    a.sendall(b"bye")
+    a.shutdown(socket.SHUT_WR)
+    got = recvall(s)
+    s.sendall(b"still open")
+    s.shutdown(socket.SHUT_WR)
+    check("remote half-close", got == b"bye" and recvall(a) == b"still open")
+
+    data = os.urandom(1 << 20) * 32
+    s, r = connect(port, "127.0.0.1", echo)
+
+    def write():
+        s.sendall(data)
+        s.shutdown(socket.SHUT_WR)
+
+    threading.Thread(target=write).start()
+    got = recvall(s)
+    check(
+        "32 MiB both ways",
+        hashlib.sha256(got).digest() == hashlib.sha256(data).digest(),
+    )
+
+    s, r = connect(unix, "127.0.0.1", echo)
+    s.sendall(b"unix")
+    check("unix socket listener", r == 0 and recvn(s, 4) == b"unix")
+    s.close()
+
+    stale = os.path.join(tmp, "stale.sock")
+    s = socket.socket(socket.AF_UNIX)
+    s.bind(stale)
+    s.close()
+    start(tmp, stale)
+    r = subprocess.run(
+        ["./usocks", "-l", stale], stderr=subprocess.DEVNULL, timeout=5, check=False
+    )
+    s, rc = connect(stale, "127.0.0.1", echo)
+    check("stale unix socket replaced, live one kept", r.returncode == 1 and rc == 0)
+    s.close()
+
+    a, _ = connect(mport, "127.0.0.1", echo)
+    b, _ = connect(mport, "127.0.0.1", echo)
+    x = proxy(mport)
+    x.sendall(b"\5\1\0")
+    x.settimeout(0.5)
+    try:
+        x.recv(2)
+        waited = False
+    except TimeoutError:
+        waited = True
+    a.close()
+    x.settimeout(5)
+    check("maxconn queues the third client", waited and recvn(x, 2) == b"\5\0")
+
+    w = start(tmp, f"127.0.0.1:{wport}", "-w", "2", lines=2)
+    kids = workers(w.pid)
+    base = [nfds(k) for k in kids]
+    held = [connect(wport, "127.0.0.1", echo) for _ in range(16)]
+    used = [nfds(k) > n for k, n in zip(kids, base)]
+    ok = all(r == 0 for _, r in held) and len(kids) == 2 and all(used)
+    check("workers share the port", ok)
+    os.kill(kids[0], signal.SIGKILL)
+    rc = w.wait(5)
+    t = time.time()
+    while alive(kids[1]) and time.time() - t < 5:
+        time.sleep(0.05)
+    check("a dead worker stops all", rc == 1 and not alive(kids[1]))
+
+    # A listener with a full accept queue drops SYNs: connect never finishes.
+    full = socket.socket()
+    full.bind(("127.0.0.1", 0))
+    full.listen(0)
+    fill = []
+    for _ in range(4):
+        f = socket.socket()
+        f.setblocking(False)
+        f.connect_ex(full.getsockname())
+        fill.append(f)
+    silent = proxy(port)
+    t = time.time()
+    s, r = connect(port, "127.0.0.1", full.getsockname()[1])
+    check("connect timeout gets 6", r == 6 and 9 < time.time() - t < 13)
+    check("handshake timeout closes", recvall(silent) == b"")
+
+    time.sleep(0.5)
+    check("no leaked descriptors", nfds(plain.pid) == fds)
+    print("FAILED" if failed else "all ok")
+    sys.exit(1 if failed else 0)
+
+
+try:
+    main()
+finally:
+    for p in procs:
+        p.terminate()
+        p.wait()