fix: keep legacy api keys active

suguanYang · suguanYang · commit 781ee644aade · 2026-05-06T01:53:49.000Z
diff --git a/apps/api/alembic/versions/f6a7b8c9d0e1_key_api_key_hashes.py b/apps/api/alembic/versions/f6a7b8c9d0e1_key_api_key_hashes.py
@@ -8,18 +8,19 @@
 
 from typing import Sequence, Union
 
-from alembic import op
-
-
 revision: str = "f6a7b8c9d0e1"
 down_revision: Union[str, Sequence[str], None] = "e5f6a7b8c9d0"
 branch_labels: Union[str, Sequence[str], None] = None
 depends_on: Union[str, Sequence[str], None] = None
 
 
 def upgrade() -> None:
-    op.execute("UPDATE api_keys SET is_active = false")
+    # Existing API key hashes are legacy unsalted SHA-256 values. They remain
+    # active for backwards compatibility; runtime lookup accepts both legacy
+    # SHA-256 and new keyed HMAC-SHA-256 digests.
+    pass
 
 
 def downgrade() -> None:
+    # No schema or data migration is needed.
     pass
diff --git a/apps/api/app/core/dependencies.py b/apps/api/app/core/dependencies.py
@@ -17,7 +17,7 @@
     AuthException,
     PermissionDeniedException,
 )
-from shared.utils.api_key_hashing import hash_api_key
+from shared.utils.api_key_hashing import get_api_key_lookup_hashes
 
 # Standard JWKS endpoint path (fixed, following OpenID Connect convention)
 JWKS_ENDPOINT_PATH = "/api/auth/jwks"
@@ -186,19 +186,27 @@ async def get_current_user_id(
     # Mode 1: API Key verification (for external clients)
     if token.startswith("sk_"):
         # Check identity cache first — skip DB on cache hit
-        api_key_hash = hash_api_key(token)
+        api_key_hashes = get_api_key_lookup_hashes(token)
         try:
-            cached = await identity_cache.get_cached_identity(
-                redis_pool_manager.get_redis_service(),
-                identity_cache._apikey_key(api_key_hash),
-            )
+            cached = None
+            matched_api_key_hash = None
+            redis_service = redis_pool_manager.get_redis_service()
+            for api_key_hash in api_key_hashes:
+                cached = await identity_cache.get_cached_identity(
+                    redis_service,
+                    identity_cache._apikey_key(api_key_hash),
+                )
+                if cached is not None:
+                    matched_api_key_hash = api_key_hash
+                    break
             if cached is not None:
                 cached_user_id = cached.get("user_id")
                 cached_user_tier = cached.get("user_tier")
                 if cached_user_id and isinstance(cached_user_tier, str):
                     request.state.cached_user_tier = cached_user_tier
                     request.state.cached_identity_hit = True
                     request.state.user_id = cached_user_id
+                    request.state.api_key_hash = matched_api_key_hash
                     _enforce_guest_api_key_scope(route_path, cached_user_tier)
                     return cached_user_id
         except PermissionDeniedException:
@@ -213,6 +221,7 @@ async def get_current_user_id(
             request.state.cached_user_tier = identity.user_tier
             request.state.cached_identity_hit = False
             request.state.user_id = identity.user_id
+            request.state.api_key_hash = identity.key_hash
             _enforce_guest_api_key_scope(route_path, identity.user_tier)
             return identity.user_id
         else:
diff --git a/apps/api/app/repositories/api_key_repository.py b/apps/api/app/repositories/api_key_repository.py
@@ -32,6 +32,16 @@ async def get_by_key_hash(
         )
         return result.scalar_one_or_none()
 
+    async def get_by_key_hashes(
+        self, session: AsyncSession, key_hashes: Sequence[str]
+    ) -> Optional[APIKey]:
+        """Get an API key by candidate key hash, preserving candidate priority."""
+        for key_hash in key_hashes:
+            api_key = await self.get_by_key_hash(session, key_hash)
+            if api_key is not None:
+                return api_key
+        return None
+
     async def get_by_user_id(
         self, session: AsyncSession, user_id: str
     ) -> Sequence[APIKey]:
diff --git a/apps/api/app/services/auth/api_key_service.py b/apps/api/app/services/auth/api_key_service.py
@@ -22,7 +22,7 @@
 )
 from shared.models.database.api_key import APIKey
 from shared.models.database.user_balance import UserBalance
-from shared.utils.api_key_hashing import hash_api_key
+from shared.utils.api_key_hashing import get_api_key_lookup_hashes, hash_api_key
 
 _DEFAULT_USER_TIER: str = "free"
 
@@ -33,6 +33,7 @@ class APIKeyIdentity:
 
     user_id: str
     user_tier: str
+    key_hash: str
 
 
 class APIKeyService:
@@ -116,8 +117,8 @@ async def validate_api_key_identity(
         api_key: str,
     ) -> Optional[APIKeyIdentity]:
         """Validate API key and return the authenticated identity."""
-        key_hash = hash_api_key(api_key)
-        api_key_record = await self.repository.get_by_key_hash(session, key_hash)
+        key_hashes = get_api_key_lookup_hashes(api_key)
+        api_key_record = await self.repository.get_by_key_hashes(session, key_hashes)
 
         if not api_key_record or not api_key_record.is_valid():
             return None
@@ -129,6 +130,7 @@ async def validate_api_key_identity(
         return APIKeyIdentity(
             user_id=user_id,
             user_tier=user_tier,
+            key_hash=str(api_key_record.key_hash),
         )
 
     async def _resolve_user_tier(
@@ -267,8 +269,8 @@ async def check_module_permission(
         self, session: AsyncSession, api_key: str, module: str
     ) -> bool:
         """Check whether an API key can access the requested module."""
-        key_hash = hash_api_key(api_key)
-        api_key_record = await self.repository.get_by_key_hash(session, key_hash)
+        key_hashes = get_api_key_lookup_hashes(api_key)
+        api_key_record = await self.repository.get_by_key_hashes(session, key_hashes)
 
         if not api_key_record or not api_key_record.is_valid():
             return False
diff --git a/apps/api/app/services/rate_limit/dependencies.py b/apps/api/app/services/rate_limit/dependencies.py
@@ -16,7 +16,7 @@
 
 import math
 from datetime import datetime, timezone
-from typing import AsyncGenerator
+from typing import Any, AsyncGenerator
 
 from app.core.dependencies import get_current_user_id
 from app.services.rate_limit.config import (
@@ -41,9 +41,9 @@
 from shared.core.logging import log_context
 from shared.core.state_machine.states import JobStatus
 from shared.models.database.api_key import APIKey
-from shared.utils.api_key_hashing import hash_api_key
 from shared.models.database.job import Job
 from shared.models.database.user_balance import UserBalance
+from shared.utils.api_key_hashing import get_api_key_lookup_hashes
 
 _DEFAULT_TIER: str = "free"
 _ACTIVE_JOB_STATES: tuple[str, ...] = (
@@ -140,6 +140,21 @@ async def _resolve_apikey_cache_ttl_seconds(api_key_hash: str) -> int:
         return max_ttl_seconds
 
 
+async def _get_cached_apikey_identity(
+    redis_service: Any,
+    api_key_hashes: list[str],
+) -> tuple[dict | None, str | None]:
+    """Return cached API key identity and the hash key that matched it."""
+    for api_key_hash in api_key_hashes:
+        cached = await identity_cache.get_cached_identity(
+            redis_service,
+            identity_cache._apikey_key(api_key_hash),
+        )
+        if cached is not None:
+            return cached, api_key_hash
+    return None, None
+
+
 # ---------------------------------------------------------------------------
 # with_current_user -- Layer 0 (matched system limit)
 # ---------------------------------------------------------------------------
@@ -172,11 +187,16 @@ async def with_current_user(
     if isinstance(user_tier, str) and stashed_user_id == user_id:
         if cached_identity_hit is False:
             token = _extract_bearer_token(request.headers.get("authorization"))
-            api_key_hash = None
+            api_key_hash = getattr(request.state, "api_key_hash", None)
             is_api_key_auth = isinstance(token, str) and token.startswith("sk_")
-            if token is not None and is_api_key_auth:
-                api_key_hash = hash_api_key(token)
-            if is_api_key_auth and api_key_hash:
+            if (
+                is_api_key_auth
+                and not api_key_hash
+                and isinstance(token, str)
+            ):
+                api_key_hashes = get_api_key_lookup_hashes(token)
+                api_key_hash = api_key_hashes[0] if api_key_hashes else None
+            if is_api_key_auth and isinstance(api_key_hash, str):
                 try:
                     ttl_seconds = await _resolve_apikey_cache_ttl_seconds(api_key_hash)
                     await identity_cache.set_apikey_identity(
@@ -194,24 +214,35 @@ async def with_current_user(
                     )
     else:
         token = _extract_bearer_token(request.headers.get("authorization"))
-        api_key_hash = None
+        api_key_hashes: list[str] = []
         is_api_key_auth = isinstance(token, str) and token.startswith("sk_")
         if token is not None and is_api_key_auth:
-            api_key_hash = hash_api_key(token)
+            api_key_hashes = get_api_key_lookup_hashes(token)
         cache_key: str = (
-            identity_cache._apikey_key(api_key_hash)
-            if is_api_key_auth and api_key_hash
+            identity_cache._apikey_key(api_key_hashes[0])
+            if is_api_key_auth and api_key_hashes
             else identity_cache._jwt_key(user_id)
         )
         try:
-            cached: dict | None = await identity_cache.get_cached_identity(
-                redis_service, cache_key
-            )
+            if is_api_key_auth and api_key_hashes:
+                cached, matched_api_key_hash = await _get_cached_apikey_identity(
+                    redis_service,
+                    api_key_hashes,
+                )
+            else:
+                cached = await identity_cache.get_cached_identity(
+                    redis_service,
+                    cache_key,
+                )
+                matched_api_key_hash = None
             if cached is not None:
                 user_tier = cached.get("user_tier", _DEFAULT_TIER)
+                if matched_api_key_hash:
+                    request.state.api_key_hash = matched_api_key_hash
             else:
                 user_tier = await _resolve_user_tier_from_db(user_id)
-                if is_api_key_auth and api_key_hash:
+                api_key_hash = getattr(request.state, "api_key_hash", None)
+                if is_api_key_auth and isinstance(api_key_hash, str):
                     ttl_seconds = await _resolve_apikey_cache_ttl_seconds(api_key_hash)
                     await identity_cache.set_apikey_identity(
                         redis_service,
diff --git a/apps/api/tests/contract/test_api_key_contract.py b/apps/api/tests/contract/test_api_key_contract.py
@@ -6,6 +6,9 @@
 import pytest
 from httpx import AsyncClient
 
+from tests.support.contract_database import ContractDatabase
+from shared.utils.api_key_hashing import hash_legacy_api_key
+
 
 @pytest.mark.asyncio
 async def test_should_revoke_a_created_api_key_through_http_only(
@@ -78,6 +81,37 @@ async def test_should_revoke_a_created_api_key_through_http_only(
     assert "details" not in error
 
 
+@pytest.mark.asyncio
+async def test_should_accept_an_active_legacy_sha256_api_key_hash(
+    api_client_factory: Callable[[], AbstractAsyncContextManager[AsyncClient]],
+) -> None:
+    user_id = f"legacy-user-{uuid4().hex[:12]}"
+    raw_api_key = f"sk_legacy_{uuid4().hex}"
+    legacy_key_hash = hash_legacy_api_key(raw_api_key)
+
+    async with api_client_factory() as api_client:
+        await ContractDatabase.insert_authenticated_user(
+            user_id=user_id,
+            api_key=raw_api_key,
+            user_tier="tier_5",
+        )
+        await ContractDatabase.execute(
+            """
+            UPDATE api_keys
+            SET key_hash = :legacy_key_hash
+            WHERE user_id = :user_id
+            """,
+            {
+                "legacy_key_hash": legacy_key_hash,
+                "user_id": user_id,
+            },
+        )
+        api_client.headers.update({"Authorization": f"Bearer {raw_api_key}"})
+        response = await api_client.get("/api/v1/jobs")
+
+    assert response.status_code == 200
+
+
 @pytest.mark.asyncio
 async def test_should_regenerate_an_api_key_and_invalidate_the_previous_raw_key(
     developer_api_client_factory: Callable[
diff --git a/packages/shared-python/shared/utils/api_key_hashing.py b/packages/shared-python/shared/utils/api_key_hashing.py
@@ -3,6 +3,9 @@
 import hmac
 from hashlib import sha256
 
+LEGACY_API_KEY_HASH_VERSION = "sha256"
+CURRENT_API_KEY_HASH_VERSION = "hmac-sha256"
+
 
 def hash_api_key(api_key: str) -> str:
     """Return a deterministic keyed digest for API key lookup."""
@@ -11,3 +14,17 @@ def hash_api_key(api_key: str) -> str:
     secret_key = settings.SECRET_KEY.encode("utf-8")
     api_key_bytes = api_key.encode("utf-8")
     return hmac.new(secret_key, api_key_bytes, sha256).hexdigest()
+
+
+def hash_legacy_api_key(api_key: str) -> str:
+    """Return the legacy unkeyed digest used before keyed API key hashing."""
+    return sha256(api_key.encode("utf-8")).hexdigest()
+
+
+def get_api_key_lookup_hashes(api_key: str) -> list[str]:
+    """Return hashes that may identify this API key, strongest first."""
+    current_hash = hash_api_key(api_key)
+    legacy_hash = hash_legacy_api_key(api_key)
+    if legacy_hash == current_hash:
+        return [current_hash]
+    return [current_hash, legacy_hash]
