refactor(api): validate OAuth bearer token cache with Pydantic

api/libs/oauth_bearer.py

@@ -11,10 +11,12 @@
 from dataclasses import dataclass
 from datetime import UTC, datetime
 from functools import wraps
-from typing import Literal, Protocol
+from typing import Literal, NotRequired, Protocol
 
+from pydantic import TypeAdapter
 from sqlalchemy import update
 from sqlalchemy.orm import Session
+from typing_extensions import TypedDict
 from werkzeug.exceptions import ServiceUnavailable
 
 from configs import dify_config
@@ -73,6 +75,18 @@ def try_get_auth_ctx() -> AuthContext | None:
     return _auth_ctx_var.get(None)
 
 
+class _CachedResolvedRow(TypedDict):
+    subject_email: str | None
+    subject_issuer: str | None
+    account_id: str | None
+    client_id: NotRequired[str | None]
+    token_id: str
+    expires_at: str | None
+
+
+_RESOLVED_ROW_CACHE_ADAPTER = TypeAdapter(_CachedResolvedRow)
+
+
 @dataclass(frozen=True, slots=True)
 class ResolvedRow:
     subject_email: str | None
@@ -82,7 +96,7 @@ class ResolvedRow:
     token_id: uuid.UUID
     expires_at: datetime | None
 
-    def to_cache(self) -> dict:
+    def to_cache(self) -> _CachedResolvedRow:
         return {
             "subject_email": self.subject_email,
             "subject_issuer": self.subject_issuer,
@@ -93,7 +107,7 @@ def to_cache(self) -> dict:
         }
 
     @classmethod
-    def from_cache(cls, data: dict) -> ResolvedRow:
+    def from_cache(cls, data: _CachedResolvedRow) -> ResolvedRow:
         return cls(
             subject_email=data["subject_email"],
             subject_issuer=data["subject_issuer"],
@@ -193,11 +207,11 @@ def cache_get(self, token_hash: str) -> ResolvedRow | None | Literal["invalid"]:
         raw = self._redis.get(self._cache_key(token_hash))
         if raw is None:
             return None
-        text = raw.decode() if isinstance(raw, (bytes, bytearray)) else raw
-        if text == "invalid":
+        if raw in ("invalid", b"invalid"):
             return "invalid"
         try:
-            return ResolvedRow.from_cache(json.loads(text))
+            data = _RESOLVED_ROW_CACHE_ADAPTER.validate_json(raw, strict=True)
+            return ResolvedRow.from_cache(data)
         except (ValueError, KeyError):
             logger.warning("auth:token cache entry malformed; treating as miss")
             return None

api/tests/unit_tests/libs/test_oauth_bearer.py

@@ -1,22 +1,27 @@
 from __future__ import annotations
 
+import json
 import uuid
-from datetime import UTC, datetime
+from dataclasses import replace
+from datetime import UTC, datetime, timedelta
 from unittest.mock import MagicMock, patch
 
 import pytest
 from redis import Redis
-from sqlalchemy.orm import sessionmaker
+from sqlalchemy.orm import Session, sessionmaker
 
-from constants.oauth_bearer import Scope, SubjectType, TokenType
+from constants.oauth_bearer import TOKEN_CACHE_KEY_FMT, Scope, SubjectType, TokenType
+from extensions.ext_redis import RedisClientWrapper
 from libs.oauth_bearer import (
+    POSITIVE_TTL_SECONDS,
     BearerAuthenticator,
     InvalidBearerError,
     OAuthAccessTokenResolver,
     ResolvedRow,
     Resolver,
     _TokenTypeResolver,
 )
+from models import OAuthAccessToken
 
 """The bearer catalog: which subject each token type serves and what it may do."""
 
@@ -60,6 +65,122 @@ def test_for_token_classifies_by_prefix_and_refuses_the_rest(token: str, expecte
     assert TokenType.for_token(token) is expected
 
 
+@pytest.mark.parametrize("account_id", [None, uuid.UUID("12345678-1234-5678-1234-567812345678")])
+@pytest.mark.parametrize("encoding", ["str", "bytes", "bytearray"])
+def test_token_cache_round_trip(
+    redis_transport: tuple[RedisClientWrapper, MagicMock], account_id: uuid.UUID | None, encoding: str
+) -> None:
+    row = _row(account_id)
+    payload = json.dumps(row.to_cache())
+    raw: str | bytes | bytearray = payload
+    if encoding == "bytes":
+        raw = payload.encode()
+    elif encoding == "bytearray":
+        raw = bytearray(payload.encode())
+    redis, commands = redis_transport
+    commands.return_value = raw
+    resolver = OAuthAccessTokenResolver(sessionmaker(), redis)
+
+    assert resolver.cache_get("token-hash") == row
+
+
+def test_token_cache_accepts_legacy_entries_without_client_id(
+    redis_transport: tuple[RedisClientWrapper, MagicMock],
+) -> None:
+    row = replace(_row(None), client_id=None, expires_at=None)
+    payload = row.to_cache()
+    del payload["client_id"]
+    redis, commands = redis_transport
+    commands.return_value = json.dumps(payload)
+    resolver = OAuthAccessTokenResolver(sessionmaker(), redis)
+
+    assert resolver.cache_get("token-hash") == row
+
+
+@pytest.mark.parametrize("raw", [None, "invalid", b"invalid", bytearray(b"invalid")])
+def test_token_cache_preserves_miss_and_negative_entries(
+    redis_transport: tuple[RedisClientWrapper, MagicMock], raw: str | bytes | bytearray | None
+) -> None:
+    redis, commands = redis_transport
+    commands.return_value = raw
+    resolver = OAuthAccessTokenResolver(sessionmaker(), redis)
+
+    assert resolver.cache_get("token-hash") == (None if raw is None else "invalid")
+
+
+@pytest.mark.parametrize(
+    "raw",
+    ["not json", "null", "[]", "42", '"text"', "{}", b"\xff"],
+)
+def test_token_cache_treats_malformed_payload_as_miss(
+    redis_transport: tuple[RedisClientWrapper, MagicMock], raw: str | bytes
+) -> None:
+    redis, commands = redis_transport
+    commands.return_value = raw
+    resolver = OAuthAccessTokenResolver(sessionmaker(), redis)
+
+    assert resolver.cache_get("token-hash") is None
+
+
+@pytest.mark.parametrize(
+    ("field", "value"),
+    [
+        ("subject_email", 42),
+        ("subject_issuer", []),
+        ("account_id", 42),
+        ("account_id", "not a uuid"),
+        ("client_id", {}),
+        ("token_id", None),
+        ("token_id", "not a uuid"),
+        ("expires_at", 42),
+        ("expires_at", "not a timestamp"),
+    ],
+)
+def test_token_cache_treats_invalid_fields_as_miss(
+    redis_transport: tuple[RedisClientWrapper, MagicMock], field: str, value: object
+) -> None:
+    payload: dict[str, object] = dict(_row(None).to_cache())
+    payload[field] = value
+    redis, commands = redis_transport
+    commands.return_value = json.dumps(payload)
+    resolver = OAuthAccessTokenResolver(sessionmaker(), redis)
+
+    assert resolver.cache_get("token-hash") is None
+
+
+def test_token_resolver_recovers_from_malformed_cache(
+    sqlite_session_factory: sessionmaker[Session], redis_transport: tuple[RedisClientWrapper, MagicMock]
+) -> None:
+    account_id = uuid.uuid4()
+    token = OAuthAccessToken(
+        subject_email="who@example.com",
+        subject_issuer="dify:account",
+        account_id=str(account_id),
+        client_id="client",
+        device_label="test-device",
+        prefix=TokenType.OAUTH_ACCOUNT.prefix,
+        token_hash="token-hash",
+        expires_at=datetime.now(UTC) + timedelta(hours=1),
+    )
+    redis, commands = redis_transport
+    commands.return_value = "[]"
+
+    with sqlite_session_factory() as session:
+        session.add(token)
+        session.flush()
+        resolver = OAuthAccessTokenResolver(lambda: session, redis).for_token_type(TokenType.OAUTH_ACCOUNT)
+
+        resolved = resolver.resolve("token-hash")
+
+        assert resolved is not None
+        assert resolved.token_id == uuid.UUID(token.id)
+        assert resolved.subject_email == token.subject_email
+        assert resolved.account_id == account_id
+    commands.assert_called_with(
+        "SETEX", TOKEN_CACHE_KEY_FMT.format(hash="token-hash"), POSITIVE_TTL_SECONDS, json.dumps(resolved.to_cache())
+    )
+
+
 def _authenticator_with(resolver: Resolver) -> BearerAuthenticator:
     return BearerAuthenticator({TokenType.OAUTH_ACCOUNT: resolver})