fixed findings from review

session-manager/main.py

@@ -16,10 +16,12 @@ from contextlib import asynccontextmanager
 
 import docker
 from docker.errors import DockerException, NotFound
-from fastapi import FastAPI, HTTPException, BackgroundTasks
-from fastapi.responses import JSONResponse
+from fastapi import FastAPI, HTTPException, BackgroundTasks, Request, Response
+from fastapi.responses import JSONResponse, StreamingResponse
 from pydantic import BaseModel
 import uvicorn
+import httpx
+import asyncio
 
 
 # Configuration
@@ -45,7 +47,16 @@ class SessionData(BaseModel):
 
 class SessionManager:
     def __init__(self):
-        self.docker_client = docker.from_env()
+        # Use TLS certificates for secure Docker communication
+        tls_config = docker.tls.TLSConfig(
+            client_cert=("/certs/client/cert.pem", "/certs/client/key.pem"),
+            ca_cert="/certs/client/ca.pem",
+            verify=True,
+        )
+        self.docker_client = docker.DockerClient(
+            base_url=os.getenv("DOCKER_HOST", "tcp://docker-daemon:2376"),
+            tls=tls_config,
+        )
         self.sessions: Dict[str, SessionData] = {}
         self._load_sessions()
 
@@ -316,6 +327,56 @@ async def trigger_cleanup():
     return {"message": "Cleanup completed"}
 
 
+@app.api_route(
+    "/{path:path}", methods=["GET", "POST", "PUT", "DELETE", "PATCH", "OPTIONS", "HEAD"]
+)
+async def proxy_to_session(request: Request, path: str):
+    """Proxy requests to session containers based on X-Session-ID header"""
+    session_id = request.headers.get("X-Session-ID")
+    if not session_id:
+        raise HTTPException(status_code=400, detail="Missing X-Session-ID header")
+
+    session = await session_manager.get_session(session_id)
+    if not session or session.status != "running":
+        raise HTTPException(status_code=404, detail="Session not found or not running")
+
+    # Proxy the request to the container
+    container_url = f"http://localhost:{session.port}"
+
+    # Prepare the request
+    url = f"{container_url}/{path}"
+    if request.url.query:
+        url += f"?{request.url.query}"
+
+    # Get request body
+    body = await request.body()
+
+    async with httpx.AsyncClient(timeout=30.0) as client:
+        try:
+            response = await client.request(
+                method=request.method,
+                url=url,
+                headers={
+                    k: v
+                    for k, v in request.headers.items()
+                    if k.lower() not in ["host", "x-session-id"]
+                },
+                content=body,
+                follow_redirects=False,
+            )
+
+            # Return the proxied response
+            return Response(
+                content=response.content,
+                status_code=response.status_code,
+                headers=dict(response.headers),
+            )
+        except httpx.RequestError as e:
+            raise HTTPException(
+                status_code=502, detail=f"Container proxy error: {str(e)}"
+            )
+
+
 @app.get("/health")
 async def health_check():
     """Health check endpoint"""