feat: add proxy authentication, limit concurrent connections, and allow configurable proxy port.

proxy_server.py

@@ -1,5 +1,8 @@
 #!/usr/bin/env python3
 from __future__ import annotations
+import base64
+import os
+import secrets
 import select
 import socket
 import threading
@@ -7,6 +10,15 @@ import urllib.parse
 import time
 from typing import Any
 
+def parse_positive_int(value: str | None, default: int) -> int:
+    try:
+        return max(1, int(value or default))
+    except (TypeError, ValueError):
+        return default
+
+MAX_PROXY_CONNECTIONS = parse_positive_int(os.environ.get("LOCAL_PROXY_MAX_CONNECTIONS"), 256)
+proxy_connection_sem = threading.BoundedSemaphore(MAX_PROXY_CONNECTIONS)
+
 def parse_int(value: Any) -> int:
     try:
         return int(value)
@@ -22,18 +34,57 @@ def recv_exact(sock: socket.socket, size: int) -> bytes:
         data += chunk
     return data
 
-def resolve_dns_over_tun0(host: str, dns_server: str = "8.8.8.8", timeout: float = 3.0) -> str | None:
-    try:
-        socket.inet_aton(host)
-        return host
-    except OSError:
-        pass
-    try:
-        socket.inet_pton(socket.AF_INET6, host)
-        return host
-    except OSError:
-        pass
+def parse_host_port(authority: str, default_port: int) -> tuple[str, int]:
+    authority = authority.strip()
+    if authority.startswith("["):
+        host_part, sep, rest = authority.partition("]")
+        host = host_part.lstrip("[")
+        port = default_port
+        if sep and rest.startswith(":"):
+            port_text = rest[1:]
+            port = parse_int(port_text) or default_port
+        return host, port
+    if authority.count(":") == 1:
+        host, _, port_text = authority.rpartition(":")
+        return host, parse_int(port_text) or default_port
+    return authority, default_port
 
+def get_proxy_credentials() -> tuple[str | None, str | None]:
+    user = os.environ.get("LOCAL_PROXY_USER") or os.environ.get("LOCAL_PROXY_USERNAME")
+    password = os.environ.get("LOCAL_PROXY_PASS") or os.environ.get("LOCAL_PROXY_PASSWORD")
+    if user is None and password is None:
+        return None, None
+    return user or "", password or ""
+
+def proxy_auth_enabled() -> bool:
+    user, password = get_proxy_credentials()
+    return user is not None and password is not None
+
+def parse_http_basic_auth(lines: list[str]) -> tuple[str | None, str | None]:
+    for line in lines:
+        name, sep, value = line.partition(":")
+        if not sep or name.strip().lower() != "proxy-authorization":
+            continue
+        scheme, _, token = value.strip().partition(" ")
+        if scheme.lower() != "basic" or not token:
+            return None, None
+        try:
+            decoded = base64.b64decode(token, validate=True).decode("utf-8", errors="replace")
+        except Exception:
+            return None, None
+        username, sep, password = decoded.partition(":")
+        if not sep:
+            return None, None
+        return username, password
+    return None, None
+
+def check_credentials(username: str | None, password: str | None) -> bool:
+    expected_user, expected_pass = get_proxy_credentials()
+    if expected_user is None or expected_pass is None:
+        return True
+    return secrets.compare_digest(username or "", expected_user) and secrets.compare_digest(password or "", expected_pass)
+
+def dns_query_over_tun0(host: str, qtype: int, dns_server: str, timeout: float) -> str | None:
     import random
     sock = None
     try:
@@ -52,7 +103,7 @@ def resolve_dns_over_tun0(host: str, dns_server: str = "8.8.8.8", timeout: float
             qname += len(part_bytes).to_bytes(1, "big") + part_bytes
         qname += b"\x00"
 
-        qtype_qclass = b"\x00\x01\x00\x01"
+        qtype_qclass = qtype.to_bytes(2, "big") + b"\x00\x01"
         packet = tx_id + flags + questions + rrs + qname + qtype_qclass
 
         sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
@@ -66,7 +117,7 @@ def resolve_dns_over_tun0(host: str, dns_server: str = "8.8.8.8", timeout: float
                 print("[DNS 绑定失败] [错误代码 3004] DNS 解析绑定 tun0 失败，网卡设备不存在，请检查 VPN 连接！", flush=True)
             return None
         sock.sendto(packet, (dns_server, 53))
-        resp, _ = sock.recvfrom(2048)
+        resp, _ = sock.recvfrom(4096)
     except Exception:
         return None
     finally:
@@ -77,11 +128,8 @@ def resolve_dns_over_tun0(host: str, dns_server: str = "8.8.8.8", timeout: float
                 pass
 
     try:
-        if len(resp) < 12:
+        if len(resp) < 12 or resp[:2] != tx_id:
             return None
-        if resp[:2] != tx_id:
-            return None
-
         rcode = resp[3] & 0x0F
         if rcode != 0:
             return None
@@ -92,17 +140,13 @@ def resolve_dns_over_tun0(host: str, dns_server: str = "8.8.8.8", timeout: float
             if length == 0:
                 offset += 1
                 break
-            elif (length & 0xC0) == 0xC0:
+            if (length & 0xC0) == 0xC0:
                 offset += 2
                 break
-            else:
-                offset += 1 + length
+            offset += 1 + length
 
         offset += 4
         answers_count = int.from_bytes(resp[6:8], "big")
-        if answers_count == 0:
-            return None
-
         for _ in range(answers_count):
             if offset >= len(resp):
                 break
@@ -111,11 +155,10 @@ def resolve_dns_over_tun0(host: str, dns_server: str = "8.8.8.8", timeout: float
                 if length == 0:
                     offset += 1
                     break
-                elif (length & 0xC0) == 0xC0:
+                if (length & 0xC0) == 0xC0:
                     offset += 2
                     break
-                else:
-                    offset += 1 + length
+                offset += 1 + length
             if offset + 10 > len(resp):
                 break
             atype = int.from_bytes(resp[offset : offset + 2], "big")
@@ -124,14 +167,30 @@ def resolve_dns_over_tun0(host: str, dns_server: str = "8.8.8.8", timeout: float
             offset += 10
             if offset + rdlength > len(resp):
                 break
-            if atype == 1 and aclass == 1 and rdlength == 4:
-                ip_bytes = resp[offset : offset + 4]
-                return socket.inet_ntoa(ip_bytes)
+            record = resp[offset : offset + rdlength]
+            if atype == qtype and aclass == 1:
+                if qtype == 1 and rdlength == 4:
+                    return socket.inet_ntoa(record)
+                if qtype == 28 and rdlength == 16:
+                    return socket.inet_ntop(socket.AF_INET6, record)
             offset += rdlength
     except Exception:
         return None
     return None
 
+def resolve_dns_over_tun0(host: str, dns_server: str = "8.8.8.8", timeout: float = 3.0) -> str | None:
+    try:
+        socket.inet_aton(host)
+        return host
+    except OSError:
+        pass
+    try:
+        socket.inet_pton(socket.AF_INET6, host)
+        return host
+    except OSError:
+        pass
+    return dns_query_over_tun0(host, 1, dns_server, timeout) or dns_query_over_tun0(host, 28, dns_server, timeout)
+
 def create_connection(address: tuple[str, int], timeout: float = 20) -> socket.socket:
     host, port = address
     resolved_ip = resolve_dns_over_tun0(host)
@@ -178,8 +237,24 @@ def socks5_client(client: socket.socket, first_byte: bytes) -> None:
     upstream = None
     try:
         methods_count = recv_exact(client, 1)[0]
-        recv_exact(client, methods_count)
-        client.sendall(b"\x05\x00")
+        methods = recv_exact(client, methods_count)
+        if proxy_auth_enabled():
+            if 2 not in methods:
+                client.sendall(b"\x05\xff")
+                return
+            client.sendall(b"\x05\x02")
+            auth_version = recv_exact(client, 1)[0]
+            if auth_version != 1:
+                client.sendall(b"\x01\x01")
+                return
+            username = recv_exact(client, recv_exact(client, 1)[0]).decode("utf-8", errors="replace")
+            password = recv_exact(client, recv_exact(client, 1)[0]).decode("utf-8", errors="replace")
+            if not check_credentials(username, password):
+                client.sendall(b"\x01\x01")
+                return
+            client.sendall(b"\x01\x00")
+        else:
+            client.sendall(b"\x05\x00")
         version, command, _, address_type = recv_exact(client, 4)
         if version != 5 or command != 1:
             client.sendall(b"\x05\x07\x00\x01\x00\x00\x00\x00\x00\x00")
@@ -223,12 +298,30 @@ def http_client(client: socket.socket, first_byte: bytes) -> None:
     upstream = None
     try:
         header = read_http_header(client, first_byte)
+        if b"\r\n\r\n" not in header:
+            client.sendall(b"HTTP/1.1 400 Bad Request\r\nContent-Length: 0\r\n\r\n")
+            return
         head, rest = header.split(b"\r\n\r\n", 1)
         lines = head.decode("iso-8859-1", errors="replace").split("\r\n")
-        method, target, version = lines[0].split(" ", 2)
+        try:
+            method, target, version = lines[0].split(" ", 2)
+        except ValueError:
+            client.sendall(b"HTTP/1.1 400 Bad Request\r\nContent-Length: 0\r\n\r\n")
+            return
+        if not version.startswith("HTTP/"):
+            client.sendall(b"HTTP/1.1 400 Bad Request\r\nContent-Length: 0\r\n\r\n")
+            return
+        if proxy_auth_enabled():
+            username, password = parse_http_basic_auth(lines[1:])
+            if not check_credentials(username, password):
+                client.sendall(
+                    b"HTTP/1.1 407 Proxy Authentication Required\r\n"
+                    b"Proxy-Authenticate: Basic realm=\"AimiliVPN Proxy\"\r\n"
+                    b"Content-Length: 0\r\n\r\n"
+                )
+                return
         if method.upper() == "CONNECT":
-            host, _, port_text = target.partition(":")
-            port = parse_int(port_text) or 443
+            host, port = parse_host_port(target, 443)
             upstream = create_connection((host, port), timeout=20)
             client.sendall(b"HTTP/1.1 200 Connection Established\r\n\r\n")
             if rest:
@@ -236,7 +329,11 @@ def http_client(client: socket.socket, first_byte: bytes) -> None:
             relay(client, upstream)
             return
 
-        parsed = urllib.parse.urlsplit(target)
+        try:
+            parsed = urllib.parse.urlsplit(target)
+        except ValueError:
+            client.sendall(b"HTTP/1.1 400 Bad Request\r\nContent-Length: 0\r\n\r\n")
+            return
         hostname = parsed.hostname
         port = parsed.port
         scheme = parsed.scheme
@@ -254,16 +351,15 @@ def http_client(client: socket.socket, first_byte: bytes) -> None:
                         else:
                             port = None
                     else:
-                        host_part, _, port_part = host_val.partition(":")
-                        hostname = host_part
-                        port = int(port_part) if port_part.isdigit() else None
+                        hostname, parsed_port = parse_host_port(host_val, 0)
+                        port = parsed_port or None
                     break
         if not hostname:
             client.sendall(b"HTTP/1.1 400 Bad Request\r\nContent-Length: 0\r\n\r\n")
             return
         port = port or (443 if scheme == "https" else 80)
         path = urllib.parse.urlunsplit(("", "", parsed.path or "/", parsed.query, ""))
-        headers = [line for line in lines[1:] if not line.lower().startswith(("proxy-connection:", "connection:"))]
+        headers = [line for line in lines[1:] if not line.lower().startswith(("proxy-connection:", "connection:", "proxy-authorization:"))]
         request = f"{method} {path} {version}\r\n" + "\r\n".join(headers) + "\r\nConnection: close\r\n\r\n"
         upstream = create_connection((hostname, port), timeout=20)
         upstream.sendall(request.encode("iso-8859-1") + rest)
@@ -355,7 +451,21 @@ def start_proxy_server(host: str, port: int) -> None:
     while True:
         try:
             client, address = server.accept()
-            threading.Thread(target=proxy_client, args=(client, address), daemon=True).start()
+            if not proxy_connection_sem.acquire(blocking=False):
+                print(f"[代理限流] 当前连接数已达到上限 {MAX_PROXY_CONNECTIONS}，拒绝客户端 {address}", flush=True)
+                try:
+                    client.close()
+                except OSError:
+                    pass
+                continue
+
+            def run_client() -> None:
+                try:
+                    proxy_client(client, address)
+                finally:
+                    proxy_connection_sem.release()
+
+            threading.Thread(target=run_client, daemon=True).start()
         except Exception as e:
             print(f"[ERROR] Proxy accept failed: {e}", flush=True)
-            time.sleep(0.5)
+            time.sleep(0.5)