Update run.py

run.py

@@ -1,78 +1,48 @@
-import socket
-import asyncio
-from contextlib import asynccontextmanager
-from typing import AsyncIterator
-
-from fastapi import FastAPI, Request, Response
+from fastapi import FastAPI, Request
+from fastapi.responses import Response
 from fastapi_proxy_lib.core.http import ForwardHttpProxy
 from fastapi_proxy_lib.core.tool import default_proxy_filter
 from httpx import AsyncClient
+import re
 
 proxy = ForwardHttpProxy(AsyncClient(), proxy_filter=default_proxy_filter)
 
-@asynccontextmanager
-async def close_proxy_event(_: FastAPI) -> AsyncIterator[None]:
-    """Close proxy."""
-    yield
-    await proxy.aclose()
-
-app = FastAPI(lifespan=close_proxy_event)
+app = FastAPI()
 
-# Handle all regular HTTP methods via the proxy
+# Handle HTTP(S) requests through the proxy
 @app.api_route("/{path:path}", methods=["GET", "POST", "PUT", "DELETE", "PATCH", "HEAD", "OPTIONS"])
-async def proxy_request(request: Request, path: str = ""):
-    """Handle HTTP requests."""
-    return await proxy.proxy(request=request, path=path)
-
-# Manually handle the CONNECT method for HTTPS tunneling
+async def handle_proxy(request: Request, path: str = ""):
+    """
+    Handles HTTP requests and forwards them to the destination URL.
+    """
+    # Extract the destination URL from the request headers
+    target_url = request.headers.get("host", "")
+    
+    if target_url:
+        scheme = "https" if request.url.port == 443 else "http"
+        # Reconstruct the full URL
+        forward_url = f"{scheme}://{target_url}{request.url.path}"
+        
+        # Forward the request to the actual destination using ForwardHttpProxy
+        return await proxy.proxy(request=request, path=forward_url)
+    return Response(content="Host header not found", status_code=400)
+
+
+# Handle HTTPS requests with CONNECT method (to establish a tunnel)
 @app.api_route("/{path:path}", methods=["CONNECT"])
 async def handle_connect(request: Request):
-    """Handle the CONNECT method for HTTPS tunneling."""
-    client_host = request.url.hostname
-    client_port = request.url.port or 443  # Use default HTTPS port if not specified
-
-    # Establish a connection to the target server
-    client_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
-    client_socket.settimeout(10)
-    client_socket.connect((client_host, client_port))
-
-    # Send back a 200 Connection Established response to the client
-    writer = request.scope["client"]
-    await writer.write(b"HTTP/1.1 200 Connection Established\r\n\r\n")
-    await writer.drain()
-
-    # Setup forwarding between the client and the target server
-    loop = asyncio.get_event_loop()
-
-    async def forward_data(reader, sock):
-        try:
-            while True:
-                data = await reader.read(1024)
-                if not data:
-                    break
-                sock.sendall(data)
-        finally:
-            sock.close()
-
-    async def reverse_data(sock, writer):
-        try:
-            while True:
-                data = sock.recv(1024)
-                if not data:
-                    break
-                await writer.write(data)
-                await writer.drain()
-        finally:
-            writer.close()
-
-    # Create bidirectional data forwarding
-    client_reader, client_writer = await asyncio.open_connection(sock=client_socket)
-    forward_task = loop.create_task(forward_data(client_reader, client_socket))
-    reverse_task = loop.create_task(reverse_data(client_socket, client_writer))
-
-    await asyncio.gather(forward_task, reverse_task)
-    return Response(status_code=200)
+    """
+    Handles CONNECT method for HTTPS proxying.
+    This is for forwarding HTTPS requests via curl -x.
+    """
+    # Extract the target host and port from the request's URL
+    match = re.match(r"^([^:]+):(\d+)$", request.url.path)
+    if match:
+        host, port = match.groups()
+        return await proxy.proxy(request=request, path=f"https://{host}:{port}")
+    
+    return Response(content="Invalid CONNECT request", status_code=400)
 
 if __name__ == '__main__':
     import uvicorn
-    uvicorn.run("run:app", host="0.0.0.0", port=8080)
+    uvicorn.run("run:app", host="0.0.0.0", port=8080)