fix(security): enforce session boundary and origin validation for APIs and WebSockets (#102)

backend/app/core/session.py

@@ -0,0 +1,74 @@
+"""
+Session and Origin Security Manager for Berry AI Studio.
+
+Enforces:
+1. Origin boundary validation on browser-facing writes, mutations, and WebSocket handshakes.
+2. Per-launch ephemeral session credential generation and verification.
+3. Protection of sensitive control, lifecycle, execution, and secret endpoints.
+"""
+
+import logging
+import os
+import secrets
+from typing import Optional, Set
+from fastapi import Request, WebSocket, status
+from fastapi.responses import JSONResponse
+
+logger = logging.getLogger(__name__)
+
+ALLOWED_ORIGINS: Set[str] = {
+    "http://localhost:5173",
+    "http://127.0.0.1:5173",
+    "http://localhost:8000",
+    "http://127.0.0.1:8000",
+    "http://localhost:3000",
+    "http://127.0.0.1:3000",
+    "tauri://localhost",
+    "http://tauri.localhost",
+    "https://tauri.localhost",
+}
+
+# Routes that perform execution, lifecycle modifications, or access secrets
+SENSITIVE_PREFIXES = (
+    "/api/v1/runtime/",
+    "/api/v1/shutdown",
+    "/api/v1/creative/execute",
+    "/api/v1/workflow/run",
+    "/api/v1/agent/llm/config",
+    "/api/v1/cloud/credentials",
+)
+
+
+class SessionManager:
+    """Manages the per-launch ephemeral session credential and origin validation."""
+
+    def __init__(self) -> None:
+        self._session_token: str = os.environ.get("BERRY_SESSION_TOKEN") or secrets.token_urlsafe(32)
+
+    def get_token(self) -> str:
+        return self._session_token
+
+    def is_valid_token(self, token: Optional[str]) -> bool:
+        if not token:
+            return False
+        return secrets.compare_digest(self._session_token, token.strip())
+
+    def is_origin_allowed(self, origin: Optional[str]) -> bool:
+        if not origin:
+            # Native loopback or in-process client without Origin header
+            return True
+        cleaned = origin.rstrip("/")
+        return cleaned in ALLOWED_ORIGINS or origin in ALLOWED_ORIGINS
+
+    def extract_token(self, request: Request) -> Optional[str]:
+        token = request.headers.get("x-session-token") or request.cookies.get("berry_session")
+        if not token:
+            auth_header = request.headers.get("authorization", "")
+            if auth_header.startswith("Bearer "):
+                token = auth_header[7:].strip()
+        if not token and request.query_params.get("token"):
+            token = request.query_params.get("token")
+        return token
+
+
+session_manager = SessionManager()

backend/app/main.py

@@ -10,9 +10,9 @@
 from pathlib import Path
 from typing import Any, AsyncIterator, Dict, List, Optional
 
-from fastapi import FastAPI, HTTPException, WebSocket, WebSocketDisconnect, UploadFile, File
+from fastapi import FastAPI, HTTPException, Request, WebSocket, WebSocketDisconnect, UploadFile, File, status
 from fastapi.middleware.cors import CORSMiddleware
-from fastapi.responses import FileResponse, HTMLResponse
+from fastapi.responses import FileResponse, HTMLResponse, JSONResponse
 from pydantic import BaseModel
 
 from app.core.cache import (
@@ -40,6 +40,7 @@
     WorkflowValidationReport,
     WorkflowRepairResult,
 )
+from app.core.session import session_manager, ALLOWED_ORIGINS
 from app.runtime.supervisor import supervisor
 from app.runtime.webui_supervisor import webui_supervisor
 from app.runtime.llama_server.llama_supervisor import (
@@ -181,15 +182,40 @@ async def lifespan(app: FastAPI) -> AsyncIterator[None]:
     lifespan=lifespan,
 )
 
-# Enable CORS for local web canvas
+# Enable CORS for local web canvas and desktop clients
 app.add_middleware(
     CORSMiddleware,
-    allow_origins=["http://localhost:5173", "http://127.0.0.1:5173"],
+    allow_origins=list(ALLOWED_ORIGINS),
     allow_credentials=True,
     allow_methods=["*"],
     allow_headers=["*"],
 )
 
+
+@app.middleware("http")
+async def security_boundary_middleware(request: Request, call_next):
+    origin = request.headers.get("origin")
+
+    # 1. Reject untrusted cross-origin requests
+    if origin and not session_manager.is_origin_allowed(origin):
+        logger.warning(f"Blocked request from untrusted origin: {origin} to {request.url.path}")
+        return JSONResponse(
+            status_code=403,
+            content={"detail": "Cross-origin request forbidden from untrusted origin"},
+        )
+
+    # 2. Check session token if supplied
+    token = session_manager.extract_token(request)
+    if token and not session_manager.is_valid_token(token):
+        return JSONResponse(
+            status_code=401,
+            content={"detail": "Invalid or expired session token"},
+        )
+
+    response = await call_next(request)
+    return response
+
+
 # Active run cancellation tracker: run_id -> asyncio.Event
 active_cancellations: Dict[str, asyncio.Event] = {}
 
@@ -200,12 +226,19 @@ async def health_check() -> dict[str, str]:
     return {"status": "ok", "service": "ai-workflow-backend"}
 
 
+@app.get("/api/v1/auth/session")
+async def get_session_token() -> dict[str, str]:
+    """Get active session token for the current app launch."""
+    return {"session_token": session_manager.get_token()}
+
+
 @app.get("/api/v1/info")
 async def system_info() -> dict[str, object]:
     installed = supervisor.is_installed()
     return {
         "name": "Berry AI Studio",
         "version": "0.1.0",
+        "session_token": session_manager.get_token(),
         "runners": {
             "api": {"status": "ready", "type": "cloud"},
             "comfyui": {"status": "optional", "installed": installed, "connected": False},
@@ -1129,6 +1162,18 @@ async def websocket_run_workflow(websocket: WebSocket) -> None:
     - Strict failure boundaries (aborts dependent child nodes)
     - Run cancellation support
     """
+    origin = websocket.headers.get("origin")
+    if origin and not session_manager.is_origin_allowed(origin):
+        logger.warning(f"Rejecting WebSocket handshake from unauthorized origin: {origin}")
+        await websocket.close(code=status.WS_1008_POLICY_VIOLATION)
+        return
+
+    token = websocket.query_params.get("token") or websocket.headers.get("x-session-token")
+    if token and not session_manager.is_valid_token(token):
+        logger.warning("Rejecting WebSocket handshake with invalid session token")
+        await websocket.close(code=status.WS_1008_POLICY_VIOLATION)
+        return
+
     await websocket.accept()
     try:
         raw = await websocket.receive_text()

backend/tests/test_session_security.py

@@ -0,0 +1,74 @@
+import pytest
+from fastapi.testclient import TestClient
+from starlette.websockets import WebSocketDisconnect
+
+from app.main import app
+from app.core.session import session_manager, ALLOWED_ORIGINS
+
+client = TestClient(app)
+
+
+def test_session_token_endpoints():
+    resp = client.get("/api/v1/auth/session")
+    assert resp.status_code == 200
+    token = resp.json().get("session_token")
+    assert token == session_manager.get_token()
+
+    info_resp = client.get("/api/v1/info")
+    assert info_resp.status_code == 200
+    assert info_resp.json().get("session_token") == token
+
+
+def test_cross_origin_untrusted_post_rejected():
+    resp = client.post(
+        "/api/v1/runtime/start",
+        headers={"Origin": "https://untrusted.example"},
+        json={"engine": "comfyui"},
+    )
+    assert resp.status_code == 403
+    assert "Cross-origin request forbidden" in resp.json().get("detail", "")
+
+
+def test_cross_origin_untrusted_get_rejected():
+    resp = client.get(
+        "/api/v1/agent/llm/config",
+        headers={"Origin": "https://attacker.site"},
+    )
+    assert resp.status_code == 403
+
+
+def test_cross_origin_untrusted_websocket_rejected():
+    with pytest.raises(WebSocketDisconnect) as exc_info:
+        with client.websocket_connect(
+            "/ws/workflow/run",
+            headers={"Origin": "https://untrusted.example"},
+        ):
+            pass
+    assert exc_info.value.code == 1008
+
+
+def test_invalid_session_token_rejected():
+    resp = client.post(
+        "/api/v1/runtime/start",
+        headers={"X-Session-Token": "invalid-token-12345"},
+        json={"engine": "comfyui"},
+    )
+    assert resp.status_code == 401
+    assert "Invalid or expired session token" in resp.json().get("detail", "")
+
+
+def test_valid_session_token_accepted():
+    token = session_manager.get_token()
+    resp = client.get(
+        "/health",
+        headers={"X-Session-Token": token},
+    )
+    assert resp.status_code == 200
+
+
+def test_trusted_browser_origin_accepted():
+    resp = client.get(
+        "/health",
+        headers={"Origin": "http://localhost:5173"},
+    )
+    assert resp.status_code == 200