refactor: 改用HTTP CONNECT代理替代wstunnel, 适配HF Space代理层

server.py

@@ -2,48 +2,70 @@ import asyncio
 import os
 import signal
 import sys
+from urllib.parse import urlparse
 
-WSTUNNEL_PORT = 7861
 PROXY_PORT = 7860
-WSTUNNEL_BIN = "/usr/local/bin/wstunnel"
 
 
-async def health_check(reader, writer):
+async def health_check(writer):
     writer.write(b"HTTP/1.1 200 OK\r\nContent-Length: 3\r\n\r\nOK\n")
     await writer.drain()
     writer.close()
 
 
-async def tcp_proxy(reader, writer, initial_data=None):
+async def error(writer, status, message):
+    body = message.encode()
+    writer.write(f"HTTP/1.1 {status} {message}\r\nContent-Length: {len(body)}\r\n\r\n{message}\n".encode())
+    await writer.drain()
+    writer.close()
+
+
+async def forward(src, dst):
+    try:
+        while True:
+            data = await src.read(65536)
+            if not data:
+                break
+            dst.write(data)
+            await dst.drain()
+    except (ConnectionResetError, BrokenPipeError, OSError):
+        pass
+    finally:
+        try:
+            dst.close()
+        except OSError:
+            pass
+
+
+async def handle_tunnel(target_host, target_port, reader, writer):
     try:
         backend_reader, backend_writer = await asyncio.wait_for(
-            asyncio.open_connection("127.0.0.1", WSTUNNEL_PORT), timeout=5
+            asyncio.open_connection(target_host, target_port), timeout=15
         )
     except (OSError, asyncio.TimeoutError):
-        writer.write(b"HTTP/1.1 502 Bad Gateway\r\nContent-Length: 15\r\n\r\nBackend unavailable\n")
-        await writer.drain()
-        writer.close()
+        await error(writer, 502, "Cannot connect to target")
         return
 
-    if initial_data:
-        backend_writer.write(initial_data)
-        await backend_writer.drain()
+    writer.write(b"HTTP/1.1 200 Connection Established\r\n\r\n")
+    await writer.drain()
 
-    async def forward(src, dst):
-        try:
-            while True:
-                data = await src.read(65536)
-                if not data:
-                    break
-                dst.write(data)
-                await dst.drain()
-        except (ConnectionResetError, BrokenPipeError, OSError):
-            pass
-        finally:
-            try:
-                dst.close()
-            except OSError:
-                pass
+    await asyncio.gather(
+        forward(reader, backend_writer),
+        forward(backend_reader, writer),
+    )
+
+
+async def forward_request(host, port, data, reader, writer):
+    try:
+        backend_reader, backend_writer = await asyncio.wait_for(
+            asyncio.open_connection(host, port), timeout=15
+        )
+    except (OSError, asyncio.TimeoutError):
+        await error(writer, 502, "Cannot connect to target")
+        return
+
+    backend_writer.write(data)
+    await backend_writer.drain()
 
     await asyncio.gather(
         forward(reader, backend_writer),
@@ -57,12 +79,38 @@ async def handle_client(reader, writer):
         if not data:
             writer.close()
             return
-        path = data.split(b" ", 2)[1] if data.count(b" ") >= 2 else b""
-        path = path.split(b"?")[0]
+
+        parts = data.split(b" ", 2)
+        if len(parts) < 2:
+            writer.close()
+            return
+
+        method = parts[0]
+        raw_path = parts[1]
+
+        if method == b"CONNECT":
+            target = raw_path.split(b":")
+            if len(target) == 2:
+                try:
+                    await handle_tunnel(target[0].decode(), int(target[1]), reader, writer)
+                    return
+                except (ValueError, UnicodeDecodeError):
+                    pass
+
+        path = raw_path.split(b"?")[0]
         if path in (b"/", b"/health"):
-            await health_check(reader, writer)
-        else:
-            await tcp_proxy(reader, writer, initial_data=data)
+            await health_check(writer)
+            return
+
+        if raw_path.startswith(b"http://") or raw_path.startswith(b"https://"):
+            parsed = urlparse(raw_path.decode())
+            host = parsed.hostname
+            port = parsed.port or (443 if parsed.scheme == "https" else 80)
+            if host:
+                await forward_request(host, port, data, reader, writer)
+                return
+
+        await error(writer, 400, "Bad Request")
     except (ConnectionResetError, BrokenPipeError, OSError, asyncio.TimeoutError):
         pass
     finally:
@@ -72,33 +120,12 @@ async def handle_client(reader, writer):
             pass
 
 
-async def run_wstunnel():
-    token = os.environ.get("PROXY_TOKEN", "")
-    args = [WSTUNNEL_BIN, "server"]
-    if token:
-        args += ["--restrict-http-upgrade-path-prefix", token]
-    args += [f"ws://0.0.0.0:{WSTUNNEL_PORT}"]
-
-    proc = await asyncio.create_subprocess_exec(*args, stdout=sys.stderr, stderr=sys.stderr)
-    return proc
-
-
 async def main():
-    print(f"starting wstunnel on port {WSTUNNEL_PORT}...", flush=True)
-    wstunnel_proc = await run_wstunnel()
-    await asyncio.sleep(0.5)
-
-    if wstunnel_proc.returncode is not None:
-        print(f"ERROR: wstunnel exited with code {wstunnel_proc.returncode}", flush=True)
-        sys.exit(1)
-
-    print(f"starting TCP proxy on port {PROXY_PORT}...", flush=True)
+    print(f"starting proxy on port {PROXY_PORT}...", flush=True)
     server = await asyncio.start_server(handle_client, "0.0.0.0", PROXY_PORT)
 
     async def shutdown():
         server.close()
-        wstunnel_proc.terminate()
-        await wstunnel_proc.wait()
 
     loop = asyncio.get_event_loop()
     for sig in (signal.SIGTERM, signal.SIGINT):