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>
This commit is contained in:
@@ -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:
|
||||
|
||||
@@ -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()
|
||||
|
||||
Reference in New Issue
Block a user