diff --git a/server/src/api/lifecycle.py b/server/src/api/lifecycle.py index 807d57fd..729e22c1 100644 --- a/server/src/api/lifecycle.py +++ b/server/src/api/lifecycle.py @@ -440,12 +440,20 @@ async def proxy_sandbox_endpoint_request(request: Request, sandbox_id: str, port try: # Filter headers + hop_by_hop = set(HOP_BY_HOP_HEADERS) + connection_header = request.headers.get("connection") + if connection_header: + hop_by_hop.update( + header.strip().lower() + for header in connection_header.split(",") + if header.strip() + ) headers = {} for key, value in request.headers.items(): key_lower = key.lower() if ( key_lower != "host" - and key_lower not in HOP_BY_HOP_HEADERS + and key_lower not in hop_by_hop and key_lower not in SENSITIVE_HEADERS ): headers[key] = value @@ -466,10 +474,24 @@ async def proxy_sandbox_endpoint_request(request: Request, sandbox_id: str, port resp = await client.send(req, stream=True) + hop_by_hop = set(HOP_BY_HOP_HEADERS) + connection_header = resp.headers.get("connection") + if connection_header: + hop_by_hop.update( + header.strip().lower() + for header in connection_header.split(",") + if header.strip() + ) + response_headers = { + key: value + for key, value in resp.headers.items() + if key.lower() not in hop_by_hop + } + return StreamingResponse( content=resp.aiter_bytes(), status_code=resp.status_code, - headers=resp.headers, + headers=response_headers, ) except httpx.ConnectError as e: raise HTTPException(