-
Notifications
You must be signed in to change notification settings - Fork 173
fix(keys): scope key list/revoke and block restricted keys #89
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
Changes from all commits
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -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, | ||
|
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. 🟠 P1 Enforce budget_limit_cents in the request path before exposing it as a settable restriction This change makes |
||
| ) | ||
| 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: | ||
|
|
||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,221 @@ | ||
| """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 | ||
|
|
||
|
|
||
| async def test_create_key_accepts_restrictions(keys_app, seeded_keys, db_session): | ||
| root, *_ = seeded_keys | ||
| h = {"Authorization": f"Bearer {root}"} | ||
| async with await _client(keys_app) as c: | ||
| r = await c.post( | ||
| "/v1/keys", | ||
| json={"name": "team-a", "model_allowlist": ["gpt-4o-mini"], "budget_limit_cents": 500}, | ||
| headers=h, | ||
| ) | ||
| 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 | ||
|
|
||
| row = ( | ||
| await db_session.execute(select(ApiKey).where(ApiKey.id == body["id"])) | ||
| ).scalar_one() | ||
| assert row.budget_limit_cents == 500 | ||
| assert row.model_allowlist == ["gpt-4o-mini"] | ||
|
|
||
|
|
||
| # ── 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 |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
🟠 P1 Apply the unrestricted-key gate to the equivalent management endpoints (PUT/DELETE /v1/providers/*, PUT /v1/routing) that restricted keys can still call
The change introduces restricted ("child") keys via POST /v1/keys and claims (docstring at keys.py:33-38) that denying them /v1/keys "closes the escalation path", because an unrestricted key holds "the maximum privilege ... (same trust level as PUT /v1/providers/*)". But the gate is applied only to the three /v1/keys routes. A key minted with model_allowlist/budget_limit_cents is fully authenticated by the middleware and can still call PUT /v1/providers/{provider} (app/routes/providers.py:140, takes
_kcand ignores it) with an attacker-controlled credential, or DELETE /v1/providers/{provider} (providers.py:198), PUT /v1/routing (routing.py:55), and the quality override routes — none of which call require_unrestricted. Replacing the workspace's upstream credential with one the restricted-key holder owns means every subsequent request from every key (including the operator's unrestricted root key) authenticates upstream with the attacker's account, so the attacker receives all prompts/responses and can also delete providers for a full outage. The allowlist/budget restriction on the minted key is thereby trivially bypassed, and the escalation the commit claims to close remains open via the equivalent management endpoints. Fix: import and call require_unrestricted in set_provider_key, delete_provider_key, update_routing (and the quality override/refresh routes).