Enable authentication for runtime environments (#4179)

Co-authored-by: Xingyao Wang <xingyao@all-hands.dev>

openhands/runtime/client/client.py

@@ -16,9 +16,10 @@ from contextlib import asynccontextmanager
 from pathlib import Path
 
 import pexpect
-from fastapi import FastAPI, HTTPException, Request, UploadFile
+from fastapi import Depends, FastAPI, HTTPException, Request, UploadFile
 from fastapi.exceptions import RequestValidationError
 from fastapi.responses import JSONResponse
+from fastapi.security import APIKeyHeader
 from pydantic import BaseModel
 from starlette.exceptions import HTTPException as StarletteHTTPException
 from uvicorn import run
@@ -63,6 +64,15 @@ INIT_COMMANDS = [
 ]
 SOFT_TIMEOUT_SECONDS = 5
 
+SESSION_API_KEY = os.environ.get('SESSION_API_KEY')
+api_key_header = APIKeyHeader(name='X-Session-API-Key', auto_error=False)
+
+
+def verify_api_key(api_key: str = Depends(api_key_header)):
+    if SESSION_API_KEY and api_key != SESSION_API_KEY:
+        raise HTTPException(status_code=403, detail='Invalid API Key')
+    return api_key
+
 
 class RuntimeClient:
     """RuntimeClient is running inside docker sandbox.
@@ -609,6 +619,16 @@ if __name__ == '__main__':
             response = await call_next(request)
         return response
 
+    @app.middleware('http')
+    async def authenticate_requests(request: Request, call_next):
+        if request.url.path != '/alive' and request.url.path != '/server_info':
+            try:
+                verify_api_key(request.headers.get('X-Session-API-Key'))
+            except HTTPException as e:
+                return e
+        response = await call_next(request)
+        return response
+
     @app.get('/server_info')
     async def get_server_info():
         assert client is not None

openhands/runtime/remote/runtime.py

@@ -176,6 +176,11 @@ class RemoteRuntime(Runtime):
             f'Sandbox started. Runtime ID: {self.runtime_id}, URL: {self.runtime_url}'
         )
 
+        if 'session_api_key' in start_response:
+            self.session.headers.update(
+                {'X-Session-API-Key': start_response['session_api_key']}
+            )
+
         # Initialize the eventstream and env vars
         super().__init__(
             config, event_stream, sid, plugins, env_vars, status_message_callback