diff --git a/app/config.py b/app/config.py index d4fddb2..c3e697a 100644 --- a/app/config.py +++ b/app/config.py @@ -52,9 +52,16 @@ class Settings(BaseSettings): port: int = 8000 log_level: str = "info" - # ── Encryption (auto-generated on first run if empty) ── + # ── Encryption ── + # NOTE: credential_encryption_key is NOT auto-generated. If empty, the + # app seals provider keys with a publicly-known dev key (with a loud + # warning), and packages.db.guards refuses to boot once real + # credentials are at stake. Generate one: `openssl rand -hex 32`. credential_encryption_key: str = "" api_key_pepper: str = "" + # Explicit opt-out from the startup guard that refuses to run with the + # publicly-known dev encryption key when real credentials are at stake. + allow_insecure_dev_key: bool = False # ── Provider keys via env (alternative to UI-stored keys) ── # Keep in sync with `_PROVIDERS_FROM_ENV` above. Pydantic-settings reads diff --git a/app/main.py b/app/main.py index caec121..3ab4aad 100644 --- a/app/main.py +++ b/app/main.py @@ -22,7 +22,7 @@ async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]: from app.config import get_settings from packages.db import session as session_mod - from packages.db.engine import dispose_engine, get_engine + from packages.db.engine import dispose_engine, get_engine, redacted_url from packages.db.models.base import Base settings = get_settings() @@ -39,13 +39,27 @@ async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]: ), ) log = structlog.get_logger() - log.info("lite_starting", database_url=settings.database_url) + # Redact: on the documented Postgres path the raw URL carries the DB + # password, and structured logs are retained by hosted aggregators. + log.info("lite_starting", database_url=redacted_url(settings.database_url)) engine = get_engine(settings.database_url) async with engine.begin() as conn: await conn.run_sync(Base.metadata.create_all) + # Fail closed before any traffic can be served: refuse to boot when + # provider credentials are (or would be) sealed with the publicly-known + # dev encryption key. Runs after create_all so a fresh database's empty + # provider_keys table counts as "no credentials at risk". + from packages.db.guards import assert_credential_encryption_ready + + await assert_credential_encryption_ready( + make_session=async_sessionmaker(engine, expire_on_commit=False), + database_url=settings.database_url, + allow_insecure_dev_key=settings.allow_insecure_dev_key, + ) + session_mod._session_factory = async_sessionmaker(engine, expire_on_commit=False) from app.seed import seed_initial_state @@ -53,7 +67,9 @@ async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]: async with session_mod._session_factory() as s: seed = await seed_initial_state(s) if seed.created and seed.api_key: - log.info("seed_complete", api_key=seed.api_key) + # No key material in the structured event: logs are retained by + # aggregators. The print() below is the one-time delivery channel. + log.info("seed_complete", workspace_id=seed.workspace_id) print(f"\n ✓ orcarouter-lite ready. API key: {seed.api_key}\n") from app import cache_invalidation_bus @@ -68,6 +84,36 @@ async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]: await dispose_engine() +async def unhandled_exception_handler(request, exc: Exception): + """Last-resort handler: log everything, return an opaque envelope. + + Declared `async` on purpose: Starlette runs async exception handlers on + the event loop inside the `except` block where `sys.exc_info()` still + holds the active exception. A sync handler would be dispatched via + `run_in_threadpool` in a worker thread, where `exc_info` is empty and the + traceback would be lost — defeating the whole point of this handler. + + The traceback MUST be recorded here — this is the only place an + arbitrary exception surfaces, and without it every production 500 is + undebuggable. + """ + structlog.get_logger().exception( + "unhandled_exception", + path=str(request.url.path), + error=str(exc), + exc_info=exc, + ) + return JSONResponse( + status_code=500, + content={ + "error": { + "message": "Internal server error", + "type": "server_error", + } + }, + ) + + def create_app() -> FastAPI: app = FastAPI( title="OrcaRouter Lite", @@ -106,17 +152,7 @@ async def val_exc_handler(_req, exc: RequestValidationError): content={"error": {"message": msg, "type": "validation_error"}}, ) - @app.exception_handler(Exception) - async def unhandled(_req, exc: Exception): - return JSONResponse( - status_code=500, - content={ - "error": { - "message": "Internal server error", - "type": "server_error", - } - }, - ) + app.add_exception_handler(Exception, unhandled_exception_handler) from app.middleware.auth import AuthMiddleware diff --git a/app/middleware/auth.py b/app/middleware/auth.py index ffa7bac..8939099 100644 --- a/app/middleware/auth.py +++ b/app/middleware/auth.py @@ -10,9 +10,13 @@ import json +import structlog + from packages.auth.key_validator import AuthError, validate_api_key from packages.db import session as session_mod +logger = structlog.get_logger() + SKIP_AUTH_PATHS: set[str] = { "/health", "/health/ready", @@ -95,6 +99,9 @@ async def __call__(self, scope, receive, send): await _send_error(send, e.status_code, e.message, "auth_error") return except Exception: + # A DB failure during key validation must be visible — without + # this log the operator only sees unexplained 503s. + logger.exception("auth_middleware_error", path=path) await _send_error(send, 503, "Service temporarily unavailable", "server_error") return diff --git a/app/routes/chat.py b/app/routes/chat.py index 750ea85..4fba9f2 100644 --- a/app/routes/chat.py +++ b/app/routes/chat.py @@ -30,6 +30,7 @@ from app.deps import get_db, get_key_context from app.quality_scores import resolve_model_metrics from app.schemas import ChatCompletionRequest +from packages.auth.spend import budget_exceeded, get_lifetime_spend_microcents from packages.auth.types import KeyContext from packages.db.models.request_log import RequestLog from packages.litellm_adapter.catalog import CATALOG, CATALOG_BY_ID @@ -267,6 +268,21 @@ async def chat_completions( detail=f"Model '{body.model}' is not allowed for this API key", ) + # Budget enforcement: `budget_limit_cents` is a lifetime cap on this + # key's billable (status < 400) spend. Checked before any routing, + # resolution, or cache work so an exhausted key costs the operator + # nothing — no upstream attempt, no cache fill. + if kc.budget_limit_cents is not None: + spend = await get_lifetime_spend_microcents(db, str(kc.key_id)) + if budget_exceeded(spend, kc.budget_limit_cents): + raise HTTPException( + status_code=429, + detail=( + "API key budget exhausted " + f"({spend} of {kc.budget_limit_cents * 10_000} microcents spent)." + ), + ) + client = await router_cache.get_router(db) raw_strategy = getattr(client, "strategy", None) strategy = raw_strategy if isinstance(raw_strategy, str) and raw_strategy else "balanced" diff --git a/app/routes/keys.py b/app/routes/keys.py index 2d1ca86..380b845 100644 --- a/app/routes/keys.py +++ b/app/routes/keys.py @@ -5,7 +5,7 @@ from datetime import datetime, timezone from fastapi import APIRouter, Depends, HTTPException, Response -from pydantic import BaseModel +from pydantic import BaseModel, Field from sqlalchemy import select from sqlalchemy.ext.asyncio import AsyncSession @@ -20,16 +20,46 @@ class CreateKey(BaseModel): name: str + # Optional restrictions for child keys. Only reachable by unrestricted + # callers (require_unrestricted above), so a restricted key can never + # mint a sibling with looser limits than its own — it can't mint at all. + model_allowlist: list[str] | None = None + budget_limit_cents: int | None = Field(default=None, gt=0) + + +def require_unrestricted(kc: KeyContext) -> None: + """Key management is reserved for unrestricted keys. + + A key that carries any restriction (`model_allowlist` or + `budget_limit_cents`) must not be able to mint, list, or revoke other + keys — otherwise it could create a sibling with no restrictions and + trivially bypass its own allowlist/budget. Unrestricted keys already + hold the maximum privilege this single-workspace edition exposes + (same trust level as PUT /v1/providers/*), so denying restricted keys + here grants nothing to anyone; it only closes the escalation path. + """ + if kc.model_allowlist is not None or kc.budget_limit_cents is not None: + raise HTTPException( + status_code=403, + detail=( + "Restricted API keys cannot manage keys. " + "Use an unrestricted key." + ), + ) @router.get("") async def list_keys( - _kc: KeyContext = Depends(get_key_context), + kc: KeyContext = Depends(get_key_context), db: AsyncSession = Depends(get_db), ) -> dict: + require_unrestricted(kc) rows = ( await db.execute( - select(ApiKey).where(ApiKey.is_deleted == 0).order_by(ApiKey.created_at) + select(ApiKey).where( + ApiKey.workspace_id == kc.workspace_id, + ApiKey.is_deleted == 0, + ).order_by(ApiKey.created_at) ) ).scalars().all() return { @@ -54,12 +84,15 @@ async def create_key( kc: KeyContext = Depends(get_key_context), db: AsyncSession = Depends(get_db), ) -> dict: + require_unrestricted(kc) full_key, key_hash, key_prefix = generate_api_key() row = ApiKey( workspace_id=kc.workspace_id, name=body.name, key_hash=key_hash, key_prefix=key_prefix, + model_allowlist=body.model_allowlist, + budget_limit_cents=body.budget_limit_cents, ) db.add(row) await db.commit() @@ -70,18 +103,28 @@ async def create_key( "name": row.name, "key_prefix": row.key_prefix, "api_key": full_key, # plaintext shown ONCE + "model_allowlist": row.model_allowlist, + "budget_limit_cents": row.budget_limit_cents, } @router.delete("/{key_id}", status_code=204) async def revoke_key( key_id: str, - _kc: KeyContext = Depends(get_key_context), + kc: KeyContext = Depends(get_key_context), db: AsyncSession = Depends(get_db), ) -> Response: + require_unrestricted(kc) row = ( await db.execute( - select(ApiKey).where(ApiKey.id == key_id, ApiKey.is_deleted == 0) + select(ApiKey).where( + ApiKey.id == key_id, + # Workspace scoping: without this, any key could revoke any + # other workspace's keys (the write path has always been + # scoped; the read/delete paths were not). + ApiKey.workspace_id == kc.workspace_id, + ApiKey.is_deleted == 0, + ) ) ).scalar_one_or_none() if row is None: diff --git a/packages/auth/encryption.py b/packages/auth/encryption.py index 4a7fdd2..673e55d 100644 --- a/packages/auth/encryption.py +++ b/packages/auth/encryption.py @@ -5,21 +5,69 @@ test fixtures that set the env var directly. If neither yields a key, derives a deterministic dev key from a fixed seed -so local development doesn't require any setup. **In production, set a real -64-char hex string** — the dev fallback is publicly known via the source -code, so anyone with read access to the SQLite file could decrypt provider -keys with it. +so local development doesn't require any setup. The dev fallback is +publicly known via the source code, so anyone with read access to the +SQLite file could decrypt provider keys with it: + +- Every use logs a prominent WARNING (`insecure_dev_encryption_key`). +- `packages.db.guards.assert_credential_encryption_ready` fail-closes at + startup when the fallback would protect real credentials (existing + provider rows, or any non-SQLite database) unless + ORCA_ALLOW_INSECURE_DEV_KEY=1 is set explicitly. + +Ciphertext format: + +- v1 (current): ``b"\\x01" + nonce(12) + ciphertext+tag`` +- legacy: ``nonce(12) + ciphertext+tag`` (no version byte) + +Decrypt auto-detects. A legacy blob whose first nonce byte happens to be +``0x01`` (~0.4% of legacy blobs) is attempted as v1 first and falls back +to the legacy parse when authentication fails, so upgrades never brick +stored credentials. """ from __future__ import annotations import hashlib +import logging import os +from cryptography.exceptions import InvalidTag from cryptography.hazmat.primitives.ciphers.aead import AESGCM +logger = logging.getLogger("orca.encryption") + +VERSION_BYTE = b"\x01" +_NONCE_LEN = 12 +_TAG_LEN = 16 + +_dev_fallback_warned = False + + +def _warn_dev_fallback_once() -> None: + global _dev_fallback_warned + if _dev_fallback_warned: + return + _dev_fallback_warned = True + logger.warning( + "insecure_dev_encryption_key: CREDENTIAL_ENCRYPTION_KEY is not set; " + "provider credentials are being sealed with a PUBLICLY-KNOWN dev " + "key. Anyone with read access to the database can decrypt them. " + "Generate one with `openssl rand -hex 32` (or set " + "ORCA_ALLOW_INSECURE_DEV_KEY=1 to silence this check)." + ) + def _get_encryption_key() -> bytes: + key, _source = _resolve_key_material() + return key + + +def _resolve_key_material() -> tuple[bytes, str]: + """Return (key_bytes, source) where source names how the key was obtained. + + Sources: "config" (Settings/.env), "env" (os.environ), "dev-fallback". + """ # Prefer Settings (which loads .env) over raw os.environ, because # pydantic-settings does NOT propagate .env values into os.environ. # Without this lookup, a user who follows the README and writes @@ -33,27 +81,56 @@ def _get_encryption_key() -> bytes: # Settings may not be importable in some isolated test contexts; # fall through to env-only behavior. pass - if not key_hex: - key_hex = os.environ.get("CREDENTIAL_ENCRYPTION_KEY", "") if key_hex: try: raw = bytes.fromhex(key_hex) if len(raw) >= 32: - return raw[:32] + return raw[:32], "config" except ValueError: pass - return hashlib.sha256(key_hex.encode()).digest() + return hashlib.sha256(key_hex.encode()).digest(), "config" + key_hex = os.environ.get("CREDENTIAL_ENCRYPTION_KEY", "") + if key_hex: + try: + raw = bytes.fromhex(key_hex) + if len(raw) >= 32: + return raw[:32], "env" + except ValueError: + pass + return hashlib.sha256(key_hex.encode()).digest(), "env" # Dev fallback so test fixtures and `docker compose up` Just Work. - return hashlib.sha256(b"orcarouter-lite-dev-key").digest() + _warn_dev_fallback_once() + return hashlib.sha256(b"orcarouter-lite-dev-key").digest(), "dev-fallback" + + +def is_using_insecure_dev_key() -> bool: + try: + return _resolve_key_material()[1] == "dev-fallback" + except Exception: + return False def encrypt_credential(plaintext: str) -> bytes: aes = AESGCM(_get_encryption_key()) - nonce = os.urandom(12) - return nonce + aes.encrypt(nonce, plaintext.encode("utf-8"), None) + nonce = os.urandom(_NONCE_LEN) + return VERSION_BYTE + nonce + aes.encrypt(nonce, plaintext.encode("utf-8"), None) def decrypt_credential(blob: bytes) -> str: - aes = AESGCM(_get_encryption_key()) - nonce, ciphertext = blob[:12], blob[12:] + key = _get_encryption_key() + aes = AESGCM(key) + + if blob[:1] == VERSION_BYTE and len(blob) >= 1 + _NONCE_LEN + _TAG_LEN: + try: + return aes.decrypt( + blob[1:1 + _NONCE_LEN], blob[1 + _NONCE_LEN:], None + ).decode("utf-8") + except InvalidTag: + # Could be a LEGACY blob whose first nonce byte happens to be + # 0x01 (~0.4%). Fall through and try the unversioned layout + # before giving up. + pass + + # Legacy unversioned blob: nonce(12) || ciphertext+tag. + nonce, ciphertext = blob[:_NONCE_LEN], blob[_NONCE_LEN:] return aes.decrypt(nonce, ciphertext, None).decode("utf-8") diff --git a/packages/auth/spend.py b/packages/auth/spend.py new file mode 100644 index 0000000..ded59ed --- /dev/null +++ b/packages/auth/spend.py @@ -0,0 +1,34 @@ +"""Per-key spend lookup used to enforce `ApiKey.budget_limit_cents`. + +Semantics: `budget_limit_cents` is a lifetime cap on the key's billable +spend — request-log rows with `status_code < 400`. 1 cent = 10,000 +microcents (1 USD = 1,000,000 microcents, matching chat.py's cost math). + +Kept free of FastAPI imports so it stays unit-testable and reusable from +non-HTTP contexts (background jobs, CLI minting tools). +""" + +from __future__ import annotations + +from sqlalchemy import func, select +from sqlalchemy.ext.asyncio import AsyncSession + +from packages.db.models.request_log import RequestLog + +MICROCENTS_PER_CENT = 10_000 + + +async def get_lifetime_spend_microcents( + session: AsyncSession, api_key_id: str +) -> int: + """Sum of billable (status < 400) spend ever recorded for this key.""" + stmt = select(func.coalesce(func.sum(RequestLog.cost_microcents), 0)).where( + RequestLog.api_key_id == api_key_id, + RequestLog.is_deleted == 0, + RequestLog.status_code < 400, + ) + return int((await session.execute(stmt)).scalar_one()) + + +def budget_exceeded(spend_microcents: int, budget_limit_cents: int) -> bool: + return spend_microcents >= budget_limit_cents * MICROCENTS_PER_CENT diff --git a/packages/db/engine.py b/packages/db/engine.py index 01db3fc..aa72b36 100644 --- a/packages/db/engine.py +++ b/packages/db/engine.py @@ -5,9 +5,24 @@ from __future__ import annotations +from sqlalchemy.engine import make_url from sqlalchemy.ext.asyncio import AsyncEngine, create_async_engine +def redacted_url(database_url: str) -> str: + """URL safe for logging: password replaced, garbage input never echoed. + + Uses SQLAlchemy's own renderer so every driver scheme is handled the + same way (`postgresql+asyncpg://user:***@host/db`). Unparseable input + degrades to a fixed placeholder instead of being reflected back into + logs. + """ + try: + return make_url(database_url).render_as_string(hide_password=True) + except Exception: + return "" + + def build_engine(database_url: str) -> AsyncEngine: """Build an async engine for the given URL. diff --git a/packages/db/guards.py b/packages/db/guards.py new file mode 100644 index 0000000..ec85154 --- /dev/null +++ b/packages/db/guards.py @@ -0,0 +1,77 @@ +"""Startup safety guards that need DB access. + +`assert_credential_encryption_ready` fail-closes boot when provider +credentials would be (or are) protected by the publicly-known dev +encryption key. Kept separate from `packages.auth.encryption` so the +crypto module stays free of SQLAlchemy imports. +""" + +from __future__ import annotations + +from sqlalchemy import func, select +from sqlalchemy.ext.asyncio import AsyncSession + +from packages.auth.encryption import is_using_insecure_dev_key + +_ALLOW_FLAG_ENV = "ORCA_ALLOW_INSECURE_DEV_KEY" + + +def _allow_flag_enabled(settings_value: bool, os_environ) -> bool: + if settings_value: + return True + return str(os_environ.get(_ALLOW_FLAG_ENV, "")).lower() in ("1", "true", "yes") + + +async def _count_provider_keys(session: AsyncSession) -> int: + from packages.db.models.provider_key import ProviderKey + + return int( + (await session.execute(select(func.count()).select_from(ProviderKey))).scalar_one() + ) + + +async def assert_credential_encryption_ready( + *, + make_session, + database_url: str, + allow_insecure_dev_key: bool = False, + os_environ=None, +) -> None: + """Refuse to start when the dev encryption key would guard real secrets. + + - SQLite + zero stored provider keys -> allowed (fresh dev install), + `encryption.py` warns loudly at first use. + - Anything else without a configured key -> RuntimeError with remediation. + """ + import os as _os + + environ = os_environ if os_environ is not None else _os.environ + if not is_using_insecure_dev_key(): + return + if _allow_flag_enabled(allow_insecure_dev_key, environ): + return + + is_sqlite = database_url.startswith("sqlite") + try: + async with make_session() as session: + key_rows = await _count_provider_keys(session) + except Exception: + # Table missing (pre-migration fresh DB) counts as zero rows. + key_rows = 0 + + if is_sqlite and key_rows == 0: + return + + raise RuntimeError( + "CREDENTIAL_ENCRYPTION_KEY is not set, so provider API keys would be " + "sealed with a publicly-known development key. " + + ( + f"{key_rows} provider key(s) already exist in this database." + if key_rows + else "A non-SQLite database requires an explicit encryption key." + ) + + " Generate one with `openssl rand -hex 32`, set it as " + "CREDENTIAL_ENCRYPTION_KEY, and re-save your provider keys. " + "(If you knowingly want to keep using the insecure dev key, set " + "ORCA_ALLOW_INSECURE_DEV_KEY=1.)" + ) diff --git a/tests/integration/test_budget_enforcement.py b/tests/integration/test_budget_enforcement.py new file mode 100644 index 0000000..5eff090 --- /dev/null +++ b/tests/integration/test_budget_enforcement.py @@ -0,0 +1,235 @@ +"""Budget enforcement on /v1/chat/completions. + +`budget_limit_cents` was loaded into KeyContext but never enforced anywhere — +a leaked key meant unbounded spend. These tests pin the new behavior: an +exhausted key gets 429 before any routing / cache / upstream work, +unbudgeted keys are unaffected, and the keys API can provision +budgeted/allowlisted child keys. +""" + +from __future__ import annotations + +import time +from unittest.mock import AsyncMock + +import pytest + + +@pytest.fixture +async def budget_env(tmp_sqlite_url, monkeypatch): + """Full app + seeded root key, with the router client mocked out. + + Yields (make_client, fake_client, session_factory, root_key). + """ + monkeypatch.setenv("DATABASE_URL", tmp_sqlite_url) + monkeypatch.setenv("OPENAI_API_KEY", "sk-test-openai") + + from app import config as cfg + cfg.get_settings.cache_clear() + + from packages.db.engine import build_engine + from packages.db.models.base import Base + + engine = build_engine(tmp_sqlite_url) + async with engine.begin() as conn: + await conn.run_sync(Base.metadata.create_all) + + from sqlalchemy.ext.asyncio import async_sessionmaker + + from packages.db import session as session_mod + factory = async_sessionmaker(engine, expire_on_commit=False) + session_mod._session_factory = factory + + from app.seed import seed_initial_state + async with factory() as s: + seed = await seed_initial_state(s) + + fake_client = AsyncMock() + fake_client.acompletion = AsyncMock( + return_value={ + "id": "chatcmpl-budget-test", + "model": "gpt-4o-mini", + "object": "chat.completion", + "created": int(time.time()), + "choices": [{ + "index": 0, + "message": {"role": "assistant", "content": "Hello!"}, + "finish_reason": "stop", + }], + "usage": {"prompt_tokens": 5, "completion_tokens": 2, "total_tokens": 7}, + "_orca_meta": { + "provider": "openai", + "litellm_model": "openai/gpt-4o-mini", + "latency_ms": 42, + }, + } + ) + + from app import router_cache + router_cache.invalidate_router() + + async def _fake_get_router(_session): + return fake_client + + monkeypatch.setattr(router_cache, "get_router", _fake_get_router) + + from httpx import ASGITransport, AsyncClient + + from app.main import create_app + app = create_app() + + async def make_client(api_key: str): + return AsyncClient( + transport=ASGITransport(app=app), + base_url="http://t", + headers={"Authorization": f"Bearer {api_key}"}, + ) + + yield make_client, fake_client, factory, seed.api_key + + await engine.dispose() + session_mod._session_factory = None + + +async def _make_budgeted_key( + factory, *, budget_limit_cents: int | None +) -> tuple[str, str]: + """Insert a budgeted child key; return (plaintext_key, key_id).""" + from packages.auth.hashing import generate_api_key + from packages.db.models.api_key import ApiKey + + full_key, key_hash, key_prefix = generate_api_key() + async with factory() as s: + row = ApiKey( + workspace_id="default", + name="budgeted", + key_hash=key_hash, + key_prefix=key_prefix, + budget_limit_cents=budget_limit_cents, + ) + s.add(row) + await s.commit() + await s.refresh(row) + return full_key, row.id + + +async def _add_billable_spend(factory, key_id: str, microcents: int) -> None: + from packages.db.models.request_log import RequestLog + + async with factory() as s: + s.add(RequestLog( + workspace_id="default", + api_key_id=key_id, + trace_id="budget-test-trace", + model_requested="gpt-4o-mini", + model_resolved="gpt-4o-mini", + provider="openai", + routing_strategy="balanced", + input_tokens=5, + output_tokens=2, + cost_microcents=microcents, + latency_ms=10, + status_code=200, + )) + await s.commit() + + +async def test_exhausted_budget_returns_429_without_upstream_call(budget_env): + make_client, fake, factory, _root = budget_env + key, key_id = await _make_budgeted_key(factory, budget_limit_cents=1) + # Pre-load spend past the 1-cent cap (10_000 microcents). + await _add_billable_spend(factory, key_id, microcents=20_000) + + async with await make_client(key) as c: + r = await c.post( + "/v1/chat/completions", + json={"model": "gpt-4o-mini", + "messages": [{"role": "user", "content": "hi"}]}, + ) + + assert r.status_code == 429, r.text + assert r.json()["error"]["type"] == "rate_limit_error" + fake.acompletion.assert_not_awaited() + + +async def test_blocked_request_writes_no_log_row(budget_env): + make_client, _fake, factory, _root = budget_env + key, key_id = await _make_budgeted_key(factory, budget_limit_cents=1) + await _add_billable_spend(factory, key_id, microcents=99_999) + + async with await make_client(key) as c: + await c.post( + "/v1/chat/completions", + json={"model": "gpt-4o-mini", + "messages": [{"role": "user", "content": "hi"}]}, + ) + + from sqlalchemy import func, select + + from packages.db.models.request_log import RequestLog + + async with factory() as s: + count = ( + await s.execute( + select(func.count()).select_from(RequestLog).where( + RequestLog.api_key_id == key_id + ) + ) + ).scalar_one() + assert count == 1 # only the pre-loaded history row + + +async def test_under_budget_key_serves_normally(budget_env): + make_client, fake, factory, _root = budget_env + key, _key_id = await _make_budgeted_key(factory, budget_limit_cents=100) + + async with await make_client(key) as c: + r = await c.post( + "/v1/chat/completions", + json={"model": "gpt-4o-mini", + "messages": [{"role": "user", "content": "hi"}]}, + ) + + assert r.status_code == 200, r.text + fake.acompletion.assert_awaited_once() + + +async def test_unbudgeted_root_key_unaffected(budget_env): + make_client, fake, _factory, root = budget_env + + async with await make_client(root) as c: + r = await c.post( + "/v1/chat/completions", + json={"model": "gpt-4o-mini", + "messages": [{"role": "user", "content": "hi"}]}, + ) + + assert r.status_code == 200, r.text + fake.acompletion.assert_awaited_once() + + +async def test_create_key_accepts_restrictions(budget_env): + make_client, _fake, factory, root = budget_env + + async with await make_client(root) as c: + r = await c.post("/v1/keys", json={ + "name": "team-a", + "model_allowlist": ["gpt-4o-mini"], + "budget_limit_cents": 500, + }) + + assert r.status_code == 201, r.text + body = r.json() + assert body["model_allowlist"] == ["gpt-4o-mini"] + assert body["budget_limit_cents"] == 500 + + from sqlalchemy import select + + from packages.db.models.api_key import ApiKey + + async with factory() as s: + row = ( + await s.execute(select(ApiKey).where(ApiKey.id == body["id"])) + ).scalar_one() + assert row.budget_limit_cents == 500 + assert row.model_allowlist == ["gpt-4o-mini"] diff --git a/tests/integration/test_keys_authz.py b/tests/integration/test_keys_authz.py new file mode 100644 index 0000000..8df84eb --- /dev/null +++ b/tests/integration/test_keys_authz.py @@ -0,0 +1,196 @@ +"""Key-management authorization tests. + +A restricted key (model_allowlist or budget_limit_cents set) must never be +able to mint, list, or revoke API keys — otherwise it could mint an +unrestricted sibling and bypass its own restrictions entirely. +See issue: restricted-key privilege escalation via POST /v1/keys. +""" + +import pytest + + +@pytest.fixture +async def seeded_keys(db_session): + """Seed the workspace root key plus one restricted and one budgeted key. + + Returns (root_full_key, restricted_full_key, budgeted_full_key). + """ + from app.seed import seed_initial_state + from packages.auth.hashing import generate_api_key + from packages.db.models.api_key import ApiKey + + seed = await seed_initial_state(db_session) + assert seed.api_key is not None + + def _make(**kwargs) -> str: + full_key, key_hash, key_prefix = generate_api_key() + row = ApiKey( + workspace_id="default", + name=kwargs.pop("name", "test"), + key_hash=key_hash, + key_prefix=key_prefix, + **kwargs, + ) + db_session.add(row) + return full_key + + # flush once so all rows land before any request reads them + restricted = _make(name="restricted", model_allowlist=["gpt-4o-mini"]) + budgeted = _make(name="budgeted", budget_limit_cents=500) + await db_session.commit() + return seed.api_key, restricted, budgeted + + +@pytest.fixture +async def keys_app(db_session, monkeypatch): + """FastAPI app with auth middleware and only the /v1/keys routes mounted.""" + monkeypatch.setenv("DATABASE_URL", str(db_session.bind.url)) + from fastapi import FastAPI + + from app.middleware.auth import AuthMiddleware + from packages.db import session as session_mod + + class _PassthroughFactory: + async def __aenter__(self): + return db_session + + async def __aexit__(self, *exc): + return False # propagate, don't close — fixture owns the session + + monkeypatch.setattr(session_mod, "_session_factory", lambda: _PassthroughFactory()) + + from app.routes.keys import router as keys_router + + app = FastAPI() + app.add_middleware(AuthMiddleware) + app.include_router(keys_router) + return app + + +async def _client(app): + from httpx import ASGITransport, AsyncClient + + return AsyncClient(transport=ASGITransport(app=app), base_url="http://t") + + +@pytest.mark.parametrize("which", [1, 2], ids=["allowlist-restricted", "budget-restricted"]) +async def test_restricted_key_cannot_create_keys(keys_app, seeded_keys, db_session, which): + keys, restricted, budgeted = seeded_keys + caller = (restricted, budgeted)[which - 1] + async with await _client(keys_app) as c: + r = await c.post( + "/v1/keys", + json={"name": "escalated"}, + headers={"Authorization": f"Bearer {caller}"}, + ) + assert r.status_code == 403 + # The escalation must not have persisted anything. + from sqlalchemy import func, select + + from packages.db.models.api_key import ApiKey + + count = ( + await db_session.execute(select(func.count()).select_from(ApiKey)) + ).scalar_one() + assert count == 3 # root + restricted + budgeted, nothing new + + +@pytest.mark.parametrize("which", [1, 2], ids=["allowlist-restricted", "budget-restricted"]) +async def test_restricted_key_cannot_list_keys(keys_app, seeded_keys, which): + keys, restricted, budgeted = seeded_keys + caller = (restricted, budgeted)[which - 1] + async with await _client(keys_app) as c: + r = await c.get("/v1/keys", headers={"Authorization": f"Bearer {caller}"}) + assert r.status_code == 403 + + +async def test_restricted_key_cannot_revoke_keys(keys_app, seeded_keys, db_session): + _, restricted, _budgeted = seeded_keys + from sqlalchemy import select + + from packages.db.models.api_key import ApiKey + + rows = (await db_session.execute(select(ApiKey))).scalars().all() + target_id = next(r.id for r in rows if r.name == "default") + async with await _client(keys_app) as c: + r = await c.delete( + f"/v1/keys/{target_id}", + headers={"Authorization": f"Bearer {restricted}"}, + ) + assert r.status_code == 403 + target = next(r for r in rows if r.name == "default") + assert target.is_active # untouched + + +async def test_unrestricted_key_retains_full_management(keys_app, seeded_keys): + root, _restricted, _budgeted = seeded_keys + h = {"Authorization": f"Bearer {root}"} + async with await _client(keys_app) as c: + listed = await c.get("/v1/keys", headers=h) + assert listed.status_code == 200 + + created = await c.post("/v1/keys", json={"name": "child"}, headers=h) + assert created.status_code == 201 + child_id = created.json()["id"] + + revoked = await c.delete(f"/v1/keys/{child_id}", headers=h) + assert revoked.status_code == 204 + + +# ── Workspace scoping (IDOR regression tests) ──────────────────────────── + + +async def _make_foreign_workspace_key(db_session) -> tuple[str, str]: + """A key belonging to a different workspace; returns (id, name).""" + from packages.auth.hashing import generate_api_key + from packages.db.models.api_key import ApiKey + from packages.db.models.workspace import Workspace + + db_session.add(Workspace(id="ws-other", name="Other", slug="other")) + await db_session.flush() + + full_key, key_hash, key_prefix = generate_api_key() + row = ApiKey( + workspace_id="ws-other", + name="foreign-key", + key_hash=key_hash, + key_prefix=key_prefix, + ) + db_session.add(row) + await db_session.commit() + return row.id, full_key + + +async def test_list_keys_hides_other_workspaces(keys_app, seeded_keys, db_session): + root, *_ = seeded_keys + foreign_id, _foreign_key = await _make_foreign_workspace_key(db_session) + + async with await _client(keys_app) as c: + r = await c.get( + "/v1/keys", headers={"Authorization": f"Bearer {root}"} + ) + + assert r.status_code == 200 + listed_ids = {k["id"] for k in r.json()["keys"]} + assert foreign_id not in listed_ids + + +async def test_revoke_rejects_other_workspaces_key(keys_app, seeded_keys, db_session): + from sqlalchemy import select + + from packages.db.models.api_key import ApiKey + + root, *_ = seeded_keys + foreign_id, _foreign_key = await _make_foreign_workspace_key(db_session) + + async with await _client(keys_app) as c: + r = await c.delete( + f"/v1/keys/{foreign_id}", + headers={"Authorization": f"Bearer {root}"}, + ) + + assert r.status_code == 404 + row = ( + await db_session.execute(select(ApiKey).where(ApiKey.id == foreign_id)) + ).scalar_one() + assert row.is_active # untouched diff --git a/tests/unit/test_budget_spend.py b/tests/unit/test_budget_spend.py new file mode 100644 index 0000000..21899d1 --- /dev/null +++ b/tests/unit/test_budget_spend.py @@ -0,0 +1,75 @@ +"""Unit tests for packages.auth.spend — lifetime spend aggregation.""" + +import pytest + +from packages.auth.spend import ( + MICROCENTS_PER_CENT, + budget_exceeded, + get_lifetime_spend_microcents, +) + + +@pytest.fixture +async def seeded_log(db_session): + """Two keys with a mix of billable / failed / soft-deleted rows.""" + from packages.db.models.api_key import ApiKey + from packages.db.models.request_log import RequestLog + + k1 = ApiKey(workspace_id="default", name="a", key_hash="h-a", key_prefix="p-a") + k2 = ApiKey(workspace_id="default", name="b", key_hash="h-b", key_prefix="p-b") + db_session.add_all([k1, k2]) + await db_session.flush() + + rows = [ + RequestLog( + workspace_id="default", api_key_id=k1.id, model_requested="m", + model_resolved="m", provider="openai", input_tokens=1, output_tokens=1, + cost_microcents=1000, status_code=200, routing_strategy="balanced", latency_ms=10, trace_id="t-1", + ), + RequestLog( + workspace_id="default", api_key_id=k1.id, model_requested="m", + model_resolved="m", provider="openai", input_tokens=1, output_tokens=1, + cost_microcents=500, status_code=200, routing_strategy="balanced", latency_ms=10, trace_id="t-1", + ), + # failed requests are not billable + RequestLog( + workspace_id="default", api_key_id=k1.id, model_requested="m", + model_resolved="m", provider="openai", input_tokens=9, output_tokens=9, + cost_microcents=999_999, status_code=503, routing_strategy="balanced", latency_ms=10, trace_id="t-3", + ), + # another key's spend must not leak in + RequestLog( + workspace_id="default", api_key_id=k2.id, model_requested="m", + model_resolved="m", provider="openai", input_tokens=2, output_tokens=2, + cost_microcents=777_777, status_code=200, routing_strategy="balanced", latency_ms=10, trace_id="t-4", + ), + ] + db_session.add_all(rows) + await db_session.commit() + return k1, k2 + + +async def test_spend_sums_only_billable_rows_for_the_key(db_session, seeded_log): + k1, _k2 = seeded_log + spend = await get_lifetime_spend_microcents(db_session, k1.id) + assert spend == 1500 + + +async def test_empty_history_is_zero(db_session, seeded_log): + _k1, k2 = seeded_log + from packages.db.models.api_key import ApiKey + + fresh = ApiKey(workspace_id="default", name="c", key_hash="h-c", key_prefix="p-c") + db_session.add(fresh) + await db_session.commit() + assert await get_lifetime_spend_microcents(db_session, fresh.id) == 0 + + +def test_budget_exceeded_boundary(): + assert budget_exceeded(10_000 - 1, 1) is False # just under 1 cent + assert budget_exceeded(10_000, 1) is True # exactly at the cap blocks + assert budget_exceeded(0, 1) is False + + +def test_microcent_conversion_constant(): + assert MICROCENTS_PER_CENT == 10_000 diff --git a/tests/unit/test_encryption_format.py b/tests/unit/test_encryption_format.py new file mode 100644 index 0000000..08f53ab --- /dev/null +++ b/tests/unit/test_encryption_format.py @@ -0,0 +1,84 @@ +"""Tests for the versioned ciphertext format and legacy-blob compatibility. + +The v1 format prefixes ``b"\\x01" + nonce + ct`` so key rotation becomes +possible later; legacy unversioned blobs (and the ~0.4% of them whose first +nonce byte collides with the version byte) must keep decrypting. +""" + +from __future__ import annotations + +import os + +import pytest +from cryptography.exceptions import InvalidTag + + +@pytest.fixture +def hex_key(monkeypatch): + """Pin a known 32-byte hex key via env (Settings-independent path).""" + monkeypatch.delenv("CREDENTIAL_ENCRYPTION_KEY", raising=False) + + from app import config as cfg + + cfg.get_settings.cache_clear() + s = cfg.Settings(_env_file=None, credential_encryption_key="ab" * 32) + monkeypatch.setattr(cfg, "get_settings", lambda: s) + return bytes.fromhex("ab" * 32) + + +def test_encrypt_produces_versioned_blob(hex_key): + from packages.auth.encryption import VERSION_BYTE, decrypt_credential, encrypt_credential + + blob = encrypt_credential("sk-secret") + assert blob[:1] == VERSION_BYTE + assert len(blob) == 1 + 12 + len(b"sk-secret") + 16 # ver+nonce+ct+tag + assert decrypt_credential(blob) == "sk-secret" + + +def test_decrypt_handles_legacy_unversioned_blob(hex_key): + from cryptography.hazmat.primitives.ciphers.aead import AESGCM + + from packages.auth.encryption import decrypt_credential + + aes = AESGCM(hex_key) + nonce = os.urandom(12) + assert nonce[:1] != b"\x01" + legacy = nonce + aes.encrypt(nonce, b"legacy-cred", None) + assert decrypt_credential(legacy) == "legacy-cred" + + +def test_decrypt_recovers_legacy_blob_whose_nonce_starts_with_version_byte(hex_key): + """A legacy blob with nonce[0] == 0x01 (~0.4% of old blobs) must not brick.""" + from cryptography.hazmat.primitives.ciphers.aead import AESGCM + + from packages.auth.encryption import VERSION_BYTE, decrypt_credential + + aes = AESGCM(hex_key) + nonce = VERSION_BYTE + os.urandom(11) + tricky_legacy = nonce + aes.encrypt(nonce, b"tricky", None) + assert tricky_legacy[:1] == b"\x01" + assert decrypt_credential(tricky_legacy) == "tricky" + + +def test_wrong_key_raises_invalid_tag(hex_key): + from app import config as cfg + from packages.auth.encryption import decrypt_credential, encrypt_credential + + blob = encrypt_credential("x") + + other = cfg.Settings(_env_file=None, credential_encryption_key="cd" * 32) + cfg_get = cfg.get_settings + cfg.get_settings = lambda: other # not monkeypatched; restored below + try: + with pytest.raises(InvalidTag): + decrypt_credential(blob) + finally: + cfg.get_settings = cfg_get + + +def test_truncated_blob_raises_rather_than_returning_garbage(hex_key): + from packages.auth.encryption import decrypt_credential, encrypt_credential + + blob = encrypt_credential("x") + with pytest.raises(InvalidTag): + decrypt_credential(blob[:8]) diff --git a/tests/unit/test_log_redaction.py b/tests/unit/test_log_redaction.py new file mode 100644 index 0000000..d3940fd --- /dev/null +++ b/tests/unit/test_log_redaction.py @@ -0,0 +1,26 @@ +"""redacted_url() must never leak credentials into loggable strings.""" + +from __future__ import annotations + + +def test_hides_password_on_postgres_style_url(): + from packages.db.engine import redacted_url + + out = redacted_url("postgresql+asyncpg://user:sup3rsecret@db.example.com:5432/orca") + assert "sup3rsecret" not in out + assert out.startswith("postgresql+asyncpg://user:") + assert "@db.example.com:5432/orca" in out + + +def test_password_free_url_survives_intact(): + from packages.db.engine import redacted_url + + url = "sqlite+aiosqlite:///./orca.db" + assert redacted_url(url) == url + + +def test_unparseable_input_is_replaced_not_echoed(): + from packages.db.engine import redacted_url + + garbage = "://not a url at all\x00" + assert redacted_url(garbage) == "" diff --git a/tests/unit/test_startup_guards.py b/tests/unit/test_startup_guards.py new file mode 100644 index 0000000..2f79165 --- /dev/null +++ b/tests/unit/test_startup_guards.py @@ -0,0 +1,108 @@ +"""Tests for packages.db.guards.assert_credential_encryption_ready. + +The guard must fail closed (RuntimeError) whenever the publicly-known dev +encryption key would protect real credentials — existing provider rows, or +any non-SQLite database — and allow fresh SQLite installs or explicit +opt-in. +""" + +from __future__ import annotations + +import pytest + + +def _allow_env(value: str) -> dict[str, str]: + return {"ORCA_ALLOW_INSECURE_DEV_KEY": value} + + +@pytest.fixture +async def guarded_db(tmp_sqlite_url): + """Engine + session factory over a fresh DB with tables created.""" + from packages.db.engine import build_engine + from packages.db.models.base import Base + + engine = build_engine(tmp_sqlite_url) + async with engine.begin() as conn: + await conn.run_sync(Base.metadata.create_all) + + from sqlalchemy.ext.asyncio import async_sessionmaker + + yield tmp_sqlite_url, async_sessionmaker(engine, expire_on_commit=False) + + await engine.dispose() + + +async def _add_provider_key(make_session) -> None: + from packages.db.models.provider_key import ProviderKey + + async with make_session() as s: + s.add(ProviderKey(provider="openai", encrypted_key=b"x" * 40, key_prefix="sk-...abcd")) + await s.commit() + + +async def test_fresh_sqlite_without_keys_is_allowed(guarded_db, monkeypatch): + from packages.db.guards import assert_credential_encryption_ready + + url, factory = guarded_db + await assert_credential_encryption_ready( + make_session=factory, database_url=url, + ) + + +async def test_existing_provider_keys_fail_closed(guarded_db): + from packages.db.guards import assert_credential_encryption_ready + + url, factory = guarded_db + await _add_provider_key(factory) + + with pytest.raises(RuntimeError, match="CREDENTIAL_ENCRYPTION_KEY"): + await assert_credential_encryption_ready( + make_session=factory, database_url=url, + ) + + +async def test_non_sqlite_requires_explicit_key_even_when_empty(guarded_db): + from packages.db.guards import assert_credential_encryption_ready + + url, factory = guarded_db + with pytest.raises(RuntimeError, match="non-SQLite"): + await assert_credential_encryption_ready( + make_session=factory, + # storage stays on sqlite; only the *claimed* URL is postgres + database_url="postgresql+asyncpg://user:pw@db.example/orca", + ) + + +async def test_opt_in_flag_bypasses_the_guard(guarded_db): + from packages.db.guards import assert_credential_encryption_ready + + url, factory = guarded_db + await _add_provider_key(factory) + + await assert_credential_encryption_ready( + make_session=factory, database_url=url, + os_environ=_allow_env("1"), + ) + await assert_credential_encryption_ready( + make_session=factory, database_url=url, + allow_insecure_dev_key=True, + ) + + +async def test_guard_no_ops_when_real_key_configured(guarded_db, monkeypatch): + """When Settings carries a real key, is_using_insecure_dev_key() is False + and the guard returns immediately regardless of stored rows.""" + from app import config as cfg + from packages.db.guards import assert_credential_encryption_ready + + cfg.get_settings.cache_clear() + monkeypatch.delenv("CREDENTIAL_ENCRYPTION_KEY", raising=False) + real = cfg.Settings(_env_file=None, credential_encryption_key="11" * 32) + monkeypatch.setattr(cfg, "get_settings", lambda: real) + + url, factory = guarded_db + await _add_provider_key(factory) + + await assert_credential_encryption_ready( + make_session=factory, database_url=url, + ) diff --git a/tests/unit/test_unhandled_logging.py b/tests/unit/test_unhandled_logging.py new file mode 100644 index 0000000..aa869c4 --- /dev/null +++ b/tests/unit/test_unhandled_logging.py @@ -0,0 +1,77 @@ +"""Both last-resort error paths must emit a structured-log event. + +Before the fix, the catch-all 500 handler and the auth middleware's +generic-exception 503 branch discarded the exception silently, making +production incidents undebuggable. +""" + +from __future__ import annotations + +import structlog +import structlog.testing + + +async def test_unhandled_exception_handler_returns_envelope_and_logs(): + from fastapi import FastAPI + from fastapi.testclient import TestClient + + from app.main import unhandled_exception_handler + + app = FastAPI() + app.add_exception_handler(Exception, unhandled_exception_handler) + + @app.get("/boom") + async def boom(): + raise RuntimeError("kaboom") + + # ServerErrorMiddleware sends the handler's response and then re-raises + # to the transport by design, so the client must not re-raise either. + with structlog.testing.capture_logs() as cap: + with TestClient(app, raise_server_exceptions=False) as c: + r = c.get("/boom") + + assert r.status_code == 500 + assert r.json()["error"] == {"message": "Internal server error", "type": "server_error"} + + events = [e for e in cap if e.get("event") == "unhandled_exception"] + assert events, "expected an unhandled_exception log event" + assert events[0]["path"].endswith("/boom") + assert events[0].get("exc_info") is not None + + +async def test_auth_middleware_db_failure_returns_503_and_logs(monkeypatch): + from fastapi import FastAPI + from httpx import ASGITransport, AsyncClient + + from app.middleware.auth import AuthMiddleware + from packages.db import session as session_mod + + class _ExplodingFactory: + async def __aenter__(self): + raise ConnectionError("db gone") + + async def __aexit__(self, *exc): + return False + + monkeypatch.setattr(session_mod, "_session_factory", lambda: _ExplodingFactory()) + + app = FastAPI() + app.add_middleware(AuthMiddleware) + + @app.get("/v1/anything") + async def anything(): + return {"ok": True} + + with structlog.testing.capture_logs() as cap: + async with AsyncClient( + transport=ASGITransport(app=app), + base_url="http://t", + headers={"Authorization": "Bearer sk-orca-somekey"}, + ) as c: + r = await c.get("/v1/anything") + + assert r.status_code == 503 + assert r.json()["error"]["type"] == "server_error" + + events = [e for e in cap if e.get("event") == "auth_middleware_error"] + assert events, "expected auth_middleware_error to be logged"