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 634227fa4749f9..21eda004e282c7 100644 --- a/api/tests/unit_tests/libs/test_oauth_bearer.py +++ b/api/tests/unit_tests/libs/test_oauth_bearer.py @@ -1,15 +1,19 @@ 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, @@ -17,6 +21,7 @@ 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})