Update app.py

app.py

@@ -9,23 +9,19 @@ from starlette.background import BackgroundTask
 from contextlib import asynccontextmanager
 
 # --- CONFIGURATION ---
-TARGET_URL = os.environ.get("TARGET_URL")
-HF_TOKEN = os.environ.get("HF_TOKEN")
+TARGET_URL = os.environ.get("TARGET_URL") # Your Private Direct URL
+HF_TOKEN = os.environ.get("HF_TOKEN")     # Your HF Read Token
 
 if not TARGET_URL:
-    raise ValueError("❌ TARGET_URL is missing. Set it in Space Settings > Secrets.")
+    raise ValueError("❌ TARGET_URL is missing.")
 
-# Global HTTP Client
+# Global Client
 http_client = None
 
 @asynccontextmanager
 async def lifespan(app: FastAPI):
     global http_client
-    http_client = httpx.AsyncClient(
-        timeout=httpx.Timeout(60.0, connect=10.0),
-        follow_redirects=True,
-        limits=httpx.Limits(max_keepalive_connections=20, max_connections=40)
-    )
+    http_client = httpx.AsyncClient(timeout=60.0, follow_redirects=True)
     yield
     await http_client.aclose()
 
@@ -33,96 +29,68 @@ app = FastAPI(lifespan=lifespan)
 
 @app.middleware("http")
 async def proxy_http(request: Request, call_next):
-    # 1. Build Target URL
-    # We append the HF_TOKEN as a query parameter '?token=...'
-    # This unlocks the Private Space Gate without using the Authorization header
-    path = request.url.path
-    query = request.url.query
+    url = f"{TARGET_URL}{request.url.path}"
+    if request.url.query:
+        url += f"?{request.url.query}"
+
+    # 1. Prepare Headers
+    excluded = ['host', 'content-length', 'connection', 'upgrade', 'accept-encoding']
+    proxy_headers = {k: v for k, v in request.headers.items() if k.lower() not in excluded}
+
+    # --- THE HEADER SWAP ---
+    # A. If the user sent a token, hide it in 'X-User-Auth'
+    if "authorization" in proxy_headers:
+        proxy_headers["X-User-Auth"] = proxy_headers["authorization"]
     
-    # Logic to merge existing queries with the token
-    if query:
-        target_query = f"{query}&token={HF_TOKEN}"
-    else:
-        target_query = f"token={HF_TOKEN}"
-        
-    url = f"{TARGET_URL}{path}?{target_query}"
-
-    # 2. Filter Headers
-    # We allow 'authorization' to pass through untouched (for your User Login)
-    # We allow 'content-type' (for the Login Form)
-    excluded_headers = ['host', 'content-length', 'connection', 'upgrade', 'accept-encoding']
-    proxy_headers = {
-        k: v for k, v in request.headers.items() 
-        if k.lower() not in excluded_headers
-    }
+    # B. Put the HF Token in the main 'Authorization' slot to pass the Gate
+    proxy_headers["Authorization"] = f"Bearer {HF_TOKEN}"
+    # -----------------------
 
     try:
-        # 3. Build Request
         rp_req = http_client.build_request(
             method=request.method,
             url=url,
-            headers=proxy_headers,
+            headers=proxy_headers, # Sending modified headers
             content=request.stream(),
             cookies=request.cookies
         )
-        
-        # 4. Send Request
         rp_resp = await http_client.send(rp_req, stream=True)
-
-        # 5. Filter Response Headers
+        
         res_excluded = ['content-encoding', 'content-length', 'transfer-encoding', 'connection']
-        res_headers = {
-            k: v for k, v in rp_resp.headers.items() 
-            if k.lower() not in res_excluded
-        }
+        res_headers = {k: v for k, v in rp_resp.headers.items() if k.lower() not in res_excluded}
 
-        # 6. Return Response
         return StreamingResponse(
             rp_resp.aiter_raw(),
             status_code=rp_resp.status_code,
             headers=res_headers,
             background=BackgroundTask(rp_resp.aclose)
         )
-
     except Exception as e:
-        print(f"❌ Proxy Error: {e}")
         return Response(f"Proxy Error: {str(e)}", status_code=502)
 
 @app.websocket_route("/{path:path}")
 async def proxy_ws(websocket: WebSocket):
     await websocket.accept()
-    
-    # Same logic for Websockets: Pass token in URL
     ws_url = TARGET_URL.replace("https://", "wss://").replace("http://", "ws://")
-    target_ws_url = f"{ws_url}/{websocket.path_params['path']}?token={HF_TOKEN}"
+    target_ws_url = f"{ws_url}/{websocket.path_params['path']}"
     
+    # Same Swap for Websockets
     headers = dict(websocket.headers)
-    
+    if "authorization" in headers:
+        headers["X-User-Auth"] = headers["authorization"]
+    headers["Authorization"] = f"Bearer {HF_TOKEN}"
+
     try:
         async with websockets.connect(target_ws_url, additional_headers=headers) as target_ws:
-            async def forward_to_target():
-                try:
-                    while True:
-                        data = await websocket.receive_text()
-                        await target_ws.send(data)
-                except Exception:
-                    pass
-
-            async def forward_to_client():
+            async def forward(src, dst):
                 try:
                     while True:
-                        data = await target_ws.recv()
-                        await websocket.send_text(data)
-                except Exception:
-                    pass
-
-            await asyncio.gather(forward_to_target(), forward_to_client())
-
-    except Exception as e:
-        try:
-            await websocket.close()
-        except:
-            pass
+                        data = await src.recv() if hasattr(src, 'recv') else await src.receive_text()
+                        await dst.send(data) if hasattr(dst, 'send') else await dst.send_text(data)
+                except: pass
+            await asyncio.gather(forward(websocket, target_ws), forward(target_ws, websocket))
+    except:
+        await websocket.close()
 
 if __name__ == '__main__':
     uvicorn.run(app, host='0.0.0.0', port=7860)