diff --git a/backend/packages/harness/deerflow/persistence/api_key/sql.py b/backend/packages/harness/deerflow/persistence/api_key/sql.py index 09613de3..e35f4df8 100644 --- a/backend/packages/harness/deerflow/persistence/api_key/sql.py +++ b/backend/packages/harness/deerflow/persistence/api_key/sql.py @@ -1,7 +1,9 @@ """SQLAlchemy-backed API key repository (Stage 1 PR1). -``get_active_by_hash`` is the auth hot path. ``revoked_at IS NULL`` -rides the partial index ``idx_api_keys_active``; expiry is filtered in +``get_active_by_hash`` is the auth hot path: it looks the key up by its +public ``key_prefix`` — UNIQUE and covered by the partial index +``idx_api_keys_active`` (WHERE revoked_at IS NULL) — then verifies the +full ``key_hash`` with a constant-time compare. Expiry is filtered in Python so the behaviour is identical across sqlite/postgres drivers. ``_row_to_dict`` deliberately omits ``key_hash`` — no dict this @@ -10,6 +12,7 @@ repository returns ever carries the secret material. from __future__ import annotations +import secrets import uuid from datetime import UTC, datetime from typing import Any @@ -70,17 +73,25 @@ class ApiKeyRepository: row = await session.get(ApiKeyRow, key_id) return self._row_to_dict(row) if row else None - async def get_active_by_hash(self, key_hash: str) -> dict[str, Any] | None: - """Auth hot path: return the key iff not revoked and not expired.""" + async def get_active_by_hash(self, key_hash: str, *, key_prefix: str) -> dict[str, Any] | None: + """Auth hot path: resolve an active, unexpired key. + + Looks the key up by its public ``key_prefix`` — UNIQUE and covered + by the partial index ``idx_api_keys_active`` (WHERE revoked_at IS + NULL) — then verifies the full ``key_hash`` with a constant-time + compare. Returns None on miss / hash mismatch / revoked / expired. + Expiry is filtered in Python so behaviour is driver-agnostic + (sqlite returns naive datetimes; postgres returns aware). + """ async with self._sf() as session: - result = await session.execute(select(ApiKeyRow).where(ApiKeyRow.key_hash == key_hash, ApiKeyRow.revoked_at.is_(None))) + result = await session.execute(select(ApiKeyRow).where(ApiKeyRow.key_prefix == key_prefix, ApiKeyRow.revoked_at.is_(None))) row = result.scalar_one_or_none() if row is None: return None - if row.expires_at is not None: - expires_at = row.expires_at - # SQLite (aiosqlite) returns naive datetimes even for DateTime(timezone=True); - # treat them as UTC so the comparison works driver-agnostically. + if not secrets.compare_digest(row.key_hash, key_hash): + return None + expires_at = row.expires_at + if expires_at is not None: if expires_at.tzinfo is None: expires_at = expires_at.replace(tzinfo=UTC) if expires_at <= datetime.now(UTC): @@ -95,7 +106,7 @@ class ApiKeyRepository: async def revoke(self, key_id: str) -> None: """Soft-revoke: set ``revoked_at`` (row is kept for audit).""" async with self._sf() as session: - await session.execute(update(ApiKeyRow).where(ApiKeyRow.id == key_id).values(revoked_at=datetime.now(UTC))) + await session.execute(update(ApiKeyRow).where(ApiKeyRow.id == key_id, ApiKeyRow.revoked_at.is_(None)).values(revoked_at=datetime.now(UTC))) await session.commit() async def touch_last_used(self, key_id: str) -> None: diff --git a/backend/tests/test_api_key_repo.py b/backend/tests/test_api_key_repo.py index e1592ab4..fbdd7fa6 100644 --- a/backend/tests/test_api_key_repo.py +++ b/backend/tests/test_api_key_repo.py @@ -64,7 +64,7 @@ async def test_create_then_get_active_by_hash(tmp_path): gen, created = await _mint(repo) assert created["key_prefix"] == gen.prefix assert "key_hash" not in created # never expose the hash in dicts - found = await repo.get_active_by_hash(gen.key_hash) + found = await repo.get_active_by_hash(gen.key_hash, key_prefix=gen.prefix) assert found is not None assert found["id"] == created["id"] assert found["scopes"] == "threads:read" @@ -76,7 +76,17 @@ async def test_get_active_by_hash_miss_returns_none(tmp_path): repo = await _make_repo(tmp_path) try: await _seed_sa(repo) - assert await repo.get_active_by_hash("deadbeef") is None + assert await repo.get_active_by_hash("deadbeef", key_prefix="dfk_live_nomatch0") is None + finally: + await _cleanup() + + +async def test_wrong_hash_for_valid_prefix_returns_none(tmp_path): + repo = await _make_repo(tmp_path) + try: + await _seed_sa(repo) + gen, _ = await _mint(repo) + assert await repo.get_active_by_hash("0" * 64, key_prefix=gen.prefix) is None finally: await _cleanup() @@ -87,7 +97,7 @@ async def test_revoked_key_not_active(tmp_path): await _seed_sa(repo) gen, created = await _mint(repo) await repo.revoke(created["id"]) - assert await repo.get_active_by_hash(gen.key_hash) is None + assert await repo.get_active_by_hash(gen.key_hash, key_prefix=gen.prefix) is None finally: await _cleanup() @@ -98,7 +108,7 @@ async def test_expired_key_not_active(tmp_path): await _seed_sa(repo) past = datetime.now(UTC) - timedelta(hours=1) gen, _ = await _mint(repo, expires_at=past) - assert await repo.get_active_by_hash(gen.key_hash) is None + assert await repo.get_active_by_hash(gen.key_hash, key_prefix=gen.prefix) is None finally: await _cleanup() @@ -109,7 +119,7 @@ async def test_future_expiry_still_active(tmp_path): await _seed_sa(repo) future = datetime.now(UTC) + timedelta(hours=1) gen, _ = await _mint(repo, expires_at=future) - assert await repo.get_active_by_hash(gen.key_hash) is not None + assert await repo.get_active_by_hash(gen.key_hash, key_prefix=gen.prefix) is not None finally: await _cleanup() @@ -138,3 +148,21 @@ async def test_list_by_service_account(tmp_path): assert all("key_hash" not in r for r in rows) finally: await _cleanup() + + +async def test_list_by_service_account_excludes_other_sa(tmp_path): + repo = await _make_repo(tmp_path) + try: + await _seed_sa(repo) + async with repo._sf() as session: + from deerflow.persistence.service_account.model import ServiceAccountRow + + session.add(ServiceAccountRow(id="sa-2", workspace_id="w-1", name="bot2", role="member", identity_mode="collapsed", status="active", created_by="u-alice")) + await session.commit() + await _mint(repo) # belongs to sa-1 + g2 = generate_api_key("live") + await repo.create(service_account_id="sa-2", key_prefix=g2.prefix, key_hash=g2.key_hash, name="k2", scopes="") + rows = await repo.list_by_service_account("sa-1") + assert len(rows) == 1 + finally: + await _cleanup()