refactor(persistence): ApiKeyRepository hot path via indexed prefix + constant-time hash verify (Stage 1 PR1)

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

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:

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()