From 190ad77e91503a7277dbc89e42d0fda66d413373 Mon Sep 17 00:00:00 2001 From: codada Date: Fri, 9 Oct 2026 00:59:34 +0000 Subject: [PATCH 1/3] refactor(api): validate OAuth bearer token cache with Pydantic --- api/libs/oauth_bearer.py | 26 ++++- .../unit_tests/libs/test_oauth_bearer.py | 108 +++++++++++++++++- 2 files changed, 126 insertions(+), 8 deletions(-) diff --git a/api/libs/oauth_bearer.py b/api/libs/oauth_bearer.py index 0832cd8580f151..2cf644fd72b90a 100644 --- a/api/libs/oauth_bearer.py +++ b/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 diff --git a/api/tests/unit_tests/libs/test_oauth_bearer.py b/api/tests/unit_tests/libs/test_oauth_bearer.py index 1e71354dba39e3..0af722437531dc 100644 --- a/api/tests/unit_tests/libs/test_oauth_bearer.py +++ b/api/tests/unit_tests/libs/test_oauth_bearer.py @@ -2,14 +2,17 @@ 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 import pytest from constants.oauth_bearer import Scope, SubjectType, TokenType -from libs.oauth_bearer import ResolvedRow, _TokenTypeResolver +from libs.oauth_bearer import OAuthAccessTokenResolver, ResolvedRow, _TokenTypeResolver +from models import OAuthAccessToken def _row(account_id: uuid.UUID | None) -> ResolvedRow: @@ -49,3 +52,104 @@ def test_scope_sets_are_pinned_exactly() -> None: ) def test_for_token_classifies_by_prefix_and_refuses_the_rest(token: str, expected: TokenType | None) -> None: 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(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 = MagicMock() + redis.get.return_value = raw + resolver = OAuthAccessTokenResolver(MagicMock(), redis) + + assert resolver.cache_get("token-hash") == row + + +def test_token_cache_accepts_legacy_entries_without_client_id() -> None: + row = replace(_row(None), client_id=None, expires_at=None) + payload = row.to_cache() + del payload["client_id"] + redis = MagicMock() + redis.get.return_value = json.dumps(payload) + resolver = OAuthAccessTokenResolver(MagicMock(), 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(raw: str | bytes | bytearray | None) -> None: + redis = MagicMock() + redis.get.return_value = raw + resolver = OAuthAccessTokenResolver(MagicMock(), 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(raw: str | bytes) -> None: + redis = MagicMock() + redis.get.return_value = raw + resolver = OAuthAccessTokenResolver(MagicMock(), 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(field: str, value: object) -> None: + payload: dict[str, object] = dict(_row(None).to_cache()) + payload[field] = value + redis = MagicMock() + redis.get.return_value = json.dumps(payload) + resolver = OAuthAccessTokenResolver(MagicMock(), redis) + + assert resolver.cache_get("token-hash") is None + + +def test_token_resolver_recovers_from_malformed_cache() -> 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), + ) + session = MagicMock() + session.query.return_value.filter.return_value.one_or_none.return_value = token + redis = MagicMock() + redis.get.return_value = "[]" + 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 + redis.setex.assert_called_once() + assert json.loads(redis.setex.call_args.args[2]) == resolved.to_cache() From 3a62601ba4d710bbe991d19dad1050fb353e8ba5 Mon Sep 17 00:00:00 2001 From: Codada Date: Sat, 10 Oct 2026 00:18:48 +0000 Subject: [PATCH 2/3] test(api): use SQLite for OAuth cache recovery regression --- .../unit_tests/libs/test_oauth_bearer.py | 32 ++++++++++--------- 1 file changed, 17 insertions(+), 15 deletions(-) diff --git a/api/tests/unit_tests/libs/test_oauth_bearer.py b/api/tests/unit_tests/libs/test_oauth_bearer.py index 9ac1f47d8c7a32..87f3c7dc99e091 100644 --- a/api/tests/unit_tests/libs/test_oauth_bearer.py +++ b/api/tests/unit_tests/libs/test_oauth_bearer.py @@ -10,7 +10,7 @@ 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 libs.oauth_bearer import OAuthAccessTokenResolver, ResolvedRow, _TokenTypeResolver @@ -68,7 +68,7 @@ def test_token_cache_round_trip(account_id: uuid.UUID | None, encoding: str) -> raw = bytearray(payload.encode()) redis = MagicMock() redis.get.return_value = raw - resolver = OAuthAccessTokenResolver(MagicMock(), redis) + resolver = OAuthAccessTokenResolver(sessionmaker(), redis) assert resolver.cache_get("token-hash") == row @@ -79,7 +79,7 @@ def test_token_cache_accepts_legacy_entries_without_client_id() -> None: del payload["client_id"] redis = MagicMock() redis.get.return_value = json.dumps(payload) - resolver = OAuthAccessTokenResolver(MagicMock(), redis) + resolver = OAuthAccessTokenResolver(sessionmaker(), redis) assert resolver.cache_get("token-hash") == row @@ -88,7 +88,7 @@ def test_token_cache_accepts_legacy_entries_without_client_id() -> None: def test_token_cache_preserves_miss_and_negative_entries(raw: str | bytes | bytearray | None) -> None: redis = MagicMock() redis.get.return_value = raw - resolver = OAuthAccessTokenResolver(MagicMock(), redis) + resolver = OAuthAccessTokenResolver(sessionmaker(), redis) assert resolver.cache_get("token-hash") == (None if raw is None else "invalid") @@ -100,7 +100,7 @@ def test_token_cache_preserves_miss_and_negative_entries(raw: str | bytes | byte def test_token_cache_treats_malformed_payload_as_miss(raw: str | bytes) -> None: redis = MagicMock() redis.get.return_value = raw - resolver = OAuthAccessTokenResolver(MagicMock(), redis) + resolver = OAuthAccessTokenResolver(sessionmaker(), redis) assert resolver.cache_get("token-hash") is None @@ -124,12 +124,12 @@ def test_token_cache_treats_invalid_fields_as_miss(field: str, value: object) -> payload[field] = value redis = MagicMock() redis.get.return_value = json.dumps(payload) - resolver = OAuthAccessTokenResolver(MagicMock(), redis) + resolver = OAuthAccessTokenResolver(sessionmaker(), redis) assert resolver.cache_get("token-hash") is None -def test_token_resolver_recovers_from_malformed_cache() -> None: +def test_token_resolver_recovers_from_malformed_cache(sqlite_session_factory: sessionmaker[Session]) -> None: account_id = uuid.uuid4() token = OAuthAccessToken( subject_email="who@example.com", @@ -141,17 +141,19 @@ def test_token_resolver_recovers_from_malformed_cache() -> None: token_hash="token-hash", expires_at=datetime.now(UTC) + timedelta(hours=1), ) - session = MagicMock() - session.query.return_value.filter.return_value.one_or_none.return_value = token redis = MagicMock() redis.get.return_value = "[]" - resolver = OAuthAccessTokenResolver(lambda: session, redis).for_token_type(TokenType.OAUTH_ACCOUNT) - resolved = resolver.resolve("token-hash") + with sqlite_session_factory() as session: + session.add(token) + session.flush() + resolver = OAuthAccessTokenResolver(lambda: session, redis).for_token_type(TokenType.OAUTH_ACCOUNT) - 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 + 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 redis.setex.assert_called_once() assert json.loads(redis.setex.call_args.args[2]) == resolved.to_cache() From 719ae74f9fc793bf50c58ee44157b314b090b21d Mon Sep 17 00:00:00 2001 From: Codada <108800975+Redestiny@users.noreply.github.com> Date: Tue, 20 Oct 2026 14:53:22 +0800 Subject: [PATCH 3/3] test(api): exercise OAuth cache through Redis transport fixture --- .../unit_tests/libs/test_oauth_bearer.py | 58 ++++++++++++------- 1 file changed, 36 insertions(+), 22 deletions(-) diff --git a/api/tests/unit_tests/libs/test_oauth_bearer.py b/api/tests/unit_tests/libs/test_oauth_bearer.py index 87f3c7dc99e091..28f8e78b4c4013 100644 --- a/api/tests/unit_tests/libs/test_oauth_bearer.py +++ b/api/tests/unit_tests/libs/test_oauth_bearer.py @@ -12,8 +12,9 @@ from redis import Redis from sqlalchemy.orm import Session, sessionmaker -from constants.oauth_bearer import Scope, SubjectType, TokenType -from libs.oauth_bearer import OAuthAccessTokenResolver, ResolvedRow, _TokenTypeResolver +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, OAuthAccessTokenResolver, ResolvedRow, _TokenTypeResolver from models import OAuthAccessToken @@ -58,7 +59,9 @@ def test_for_token_classifies_by_prefix_and_refuses_the_rest(token: str, expecte @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(account_id: uuid.UUID | None, encoding: str) -> None: +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 @@ -66,28 +69,32 @@ def test_token_cache_round_trip(account_id: uuid.UUID | None, encoding: str) -> raw = payload.encode() elif encoding == "bytearray": raw = bytearray(payload.encode()) - redis = MagicMock() - redis.get.return_value = raw + 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() -> None: +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 = MagicMock() - redis.get.return_value = json.dumps(payload) + 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(raw: str | bytes | bytearray | None) -> None: - redis = MagicMock() - redis.get.return_value = raw +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") @@ -97,9 +104,11 @@ def test_token_cache_preserves_miss_and_negative_entries(raw: str | bytes | byte "raw", ["not json", "null", "[]", "42", '"text"', "{}", b"\xff"], ) -def test_token_cache_treats_malformed_payload_as_miss(raw: str | bytes) -> None: - redis = MagicMock() - redis.get.return_value = raw +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 @@ -119,17 +128,21 @@ def test_token_cache_treats_malformed_payload_as_miss(raw: str | bytes) -> None: ("expires_at", "not a timestamp"), ], ) -def test_token_cache_treats_invalid_fields_as_miss(field: str, value: object) -> None: +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 = MagicMock() - redis.get.return_value = json.dumps(payload) + 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]) -> 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", @@ -141,8 +154,8 @@ def test_token_resolver_recovers_from_malformed_cache(sqlite_session_factory: se token_hash="token-hash", expires_at=datetime.now(UTC) + timedelta(hours=1), ) - redis = MagicMock() - redis.get.return_value = "[]" + redis, commands = redis_transport + commands.return_value = "[]" with sqlite_session_factory() as session: session.add(token) @@ -155,5 +168,6 @@ def test_token_resolver_recovers_from_malformed_cache(sqlite_session_factory: se assert resolved.token_id == uuid.UUID(token.id) assert resolved.subject_email == token.subject_email assert resolved.account_id == account_id - redis.setex.assert_called_once() - assert json.loads(redis.setex.call_args.args[2]) == resolved.to_cache() + commands.assert_called_with( + "SETEX", TOKEN_CACHE_KEY_FMT.format(hash="token-hash"), POSITIVE_TTL_SECONDS, json.dumps(resolved.to_cache()) + )