Phase 1: Auth + CORS + Rate Limiting

agent/.env.example

@@ -29,5 +29,10 @@ BRAVE_API_KEY=
 GITHUB_TOKEN=
 GITHUB_ORG=
 
+# === Auth & CORS ===
+VIBEDEPLOY_API_KEY=
+VIBEDEPLOY_OPS_TOKEN=
+VIBEDEPLOY_CORS_ORIGINS=http://localhost:3000
+
 # === Database ===
 DATABASE_URL=

agent/auth.py

@@ -0,0 +1,152 @@
+"""Authentication and rate-limiting middleware for vibeDeploy gateway."""
+
+import hmac
+import logging
+import os
+import time
+from collections import defaultdict
+
+from fastapi import HTTPException, Request, Security
+from fastapi.security import APIKeyHeader
+
+logger = logging.getLogger(__name__)
+
+_api_key_header = APIKeyHeader(name="X-API-Key", auto_error=False)
+
+_PUBLIC_PATHS: frozenset[str] = frozenset(
+    {
+        "/",
+        "/health",
+        "/cost-estimate",
+        "/api/cost-estimate",
+        "/models",
+        "/api/models",
+    }
+)
+
+_PUBLIC_PREFIXES: tuple[str, ...] = ("/test/",)
+
+
+def _get_api_key() -> str:
+    for key in ("VIBEDEPLOY_API_KEY", "VIBEDEPLOY_OPS_TOKEN", "DASHBOARD_ADMIN_TOKEN"):
+        value = os.getenv(key, "").strip()
+        if value:
+            return value
+    return ""
+
+
+def _is_public_path(path: str) -> bool:
+    if path in _PUBLIC_PATHS:
+        return True
+    return any(path.startswith(prefix) for prefix in _PUBLIC_PREFIXES)
+
+
+async def verify_api_key(
+    request: Request,
+    api_key: str | None = Security(_api_key_header),
+) -> str | None:
+    path = request.url.path
+
+    if _is_public_path(path):
+        return None
+
+    expected = _get_api_key()
+
+    # Auth disabled when no key is configured (dev mode)
+    if not expected:
+        return None
+
+    # SSE endpoints: fall back to query param (EventSource can't send headers)
+    if not api_key and "/events" in path:
+        api_key = request.query_params.get("api_key")
+
+    if not api_key:
+        raise HTTPException(
+            status_code=401,
+            detail="missing_api_key",
+            headers={"WWW-Authenticate": "ApiKey"},
+        )
+
+    if not hmac.compare_digest(api_key, expected):
+        logger.warning("Invalid API key from %s for %s", request.client.host if request.client else "unknown", path)
+        raise HTTPException(status_code=403, detail="invalid_api_key")
+
+    return api_key
+
+
+class _RateLimitBucket:
+    __slots__ = ("_requests",)
+
+    def __init__(self) -> None:
+        self._requests: list[float] = []
+
+    def hit(self, now: float, window_seconds: int, max_requests: int) -> bool:
+        cutoff = now - window_seconds
+        self._requests = [t for t in self._requests if t > cutoff]
+        if len(self._requests) >= max_requests:
+            return False
+        self._requests.append(now)
+        return True
+
+
+_rate_buckets: defaultdict[str, _RateLimitBucket] = defaultdict(_RateLimitBucket)
+
+_RATE_LIMITS: dict[str, tuple[int, int]] = {
+    "write": (10, 60),
+    "read": (120, 60),
+    "sse": (20, 60),
+}
+
+_WRITE_PATHS: frozenset[str] = frozenset(
+    {
+        "/run",
+        "/api/run",
+        "/resume",
+        "/api/resume",
+        "/brainstorm",
+        "/api/brainstorm",
+        "/zero-prompt/start",
+        "/api/zero-prompt/start",
+        "/zero-prompt/reset",
+        "/api/zero-prompt/reset",
+    }
+)
+
+_SSE_FRAGMENTS: tuple[str, ...] = ("/events", "/build/")
+
+
+def _classify_rate_tier(path: str, method: str) -> str:
+    if path in _WRITE_PATHS:
+        return "write"
+    for fragment in _SSE_FRAGMENTS:
+        if fragment in path:
+            return "sse"
+    if method == "POST" and "/actions" in path:
+        return "write"
+    return "read"
+
+
+async def rate_limit_check(request: Request) -> None:
+    path = request.url.path
+
+    if _is_public_path(path):
+        return
+
+    client_ip = request.client.host if request.client else "unknown"
+    method = request.method
+    tier = _classify_rate_tier(path, method)
+    max_requests, window_seconds = _RATE_LIMITS[tier]
+    bucket_key = f"{client_ip}:{tier}"
+    now = time.monotonic()
+
+    if not _rate_buckets[bucket_key].hit(now, window_seconds, max_requests):
+        logger.warning("Rate limit exceeded: %s %s from %s (tier=%s)", method, path, client_ip, tier)
+        raise HTTPException(
+            status_code=429,
+            detail="rate_limit_exceeded",
+            headers={
+                "Retry-After": str(window_seconds),
+                "X-RateLimit-Limit": str(max_requests),
+                "X-RateLimit-Window": str(window_seconds),
+            },
+        )

agent/server.py

@@ -27,7 +27,7 @@
 import httpx
 import uvicorn
 from dotenv import load_dotenv
-from fastapi import FastAPI, Header, HTTPException, Request
+from fastapi import Depends, FastAPI, Header, HTTPException, Request
 from fastapi.middleware.cors import CORSMiddleware
 from pydantic import BaseModel
 from starlette.responses import JSONResponse, StreamingResponse
@@ -346,7 +346,7 @@ def _meeting_store_payload(meeting: dict) -> dict:
 
 
 def _ops_token() -> str:
-    for key in ("VIBEDEPLOY_OPS_TOKEN", "DASHBOARD_ADMIN_TOKEN", "DIGITALOCEAN_INFERENCE_KEY"):
+    for key in ("VIBEDEPLOY_OPS_TOKEN", "DASHBOARD_ADMIN_TOKEN"):
         value = os.getenv(key, "").strip()
         if value:
             return value
@@ -610,13 +610,30 @@ async def lifespan(app: FastAPI):
     _store = None
 
 
-app = FastAPI(title="vibeDeploy Agent (local)", lifespan=lifespan)
+from .auth import rate_limit_check, verify_api_key
+
+app = FastAPI(
+    title="vibeDeploy Agent (local)",
+    lifespan=lifespan,
+    dependencies=[Depends(verify_api_key), Depends(rate_limit_check)],
+)
+
+_ALLOWED_ORIGINS = [
+    origin.strip()
+    for origin in os.getenv(
+        "VIBEDEPLOY_CORS_ORIGINS",
+        "https://vibedeploy-7tgzk.ondigitalocean.app,http://localhost:3000,http://localhost:9001",
+    ).split(",")
+    if origin.strip()
+]
 
 app.add_middleware(
     CORSMiddleware,
-    allow_origins=["*"],
-    allow_methods=["*"],
-    allow_headers=["*"],
+    allow_origins=_ALLOWED_ORIGINS,
+    allow_methods=["GET", "POST", "PUT", "OPTIONS"],
+    allow_headers=["Content-Type", "X-API-Key", "X-Vibedeploy-Ops-Token"],
+    allow_credentials=False,
+    max_age=600,
 )
 
 _NODE_EVENTS = NODE_EVENTS

agent/tests/conftest.py

@@ -7,9 +7,18 @@
 import pytest_asyncio
 from httpx import ASGITransport, AsyncClient
 
+from agent.auth import _rate_buckets
 from agent.db.store import ResultStore
 
 
+@pytest.fixture(autouse=True)
+def _clear_rate_limits():
+    """Clear rate limit buckets before each test to avoid cross-test pollution."""
+    _rate_buckets.clear()
+    yield
+    _rate_buckets.clear()
+
+
 @pytest_asyncio.fixture
 async def store() -> AsyncIterator[ResultStore]:
     s = ResultStore(":memory:")