Update run.py

run.py

@@ -1,48 +1,27 @@
-from fastapi import FastAPI, Request
-from fastapi.responses import Response
+from contextlib import asynccontextmanager
+from typing import AsyncIterator
+
+from fastapi import FastAPI
 from fastapi_proxy_lib.core.http import ForwardHttpProxy
 from fastapi_proxy_lib.core.tool import default_proxy_filter
 from httpx import AsyncClient
-import re
+from starlette.requests import Request
 
 proxy = ForwardHttpProxy(AsyncClient(), proxy_filter=default_proxy_filter)
 
-app = FastAPI()
+@asynccontextmanager
+async def close_proxy_event(_: FastAPI) -> AsyncIterator[None]:
+    """Close proxy."""
+    yield
+    await proxy.aclose()
 
-# Handle HTTP(S) requests through the proxy
-@app.api_route("/{path:path}", methods=["GET", "POST", "PUT", "DELETE", "PATCH", "HEAD", "OPTIONS"])
-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)
+app = FastAPI(lifespan=close_proxy_event)
 
+@app.api_route("/{path:path}", methods=["GET", "POST", "PUT", "DELETE", "PATCH", "HEAD", "OPTIONS", "CONNECT"])
+async def _(request: Request, path: str = ""):
+    return await proxy.proxy(request=request, path=path)
 
-# Handle HTTPS requests with CONNECT method (to establish a tunnel)
-@app.api_route("/{path:path}", methods=["CONNECT"])
-async def handle_connect(request: Request):
-    """
-    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")