mirror of
https://github.com/langgenius/dify.git
synced 2026-07-15 17:37:38 +08:00
220 lines
7.7 KiB
Python
220 lines
7.7 KiB
Python
from __future__ import annotations
|
|
|
|
import uuid
|
|
from datetime import UTC, datetime, timedelta
|
|
from unittest.mock import MagicMock
|
|
|
|
import pytest
|
|
from sqlalchemy.orm import Session
|
|
|
|
from libs.oauth_bearer import TOKEN_CACHE_KEY_FMT, AuthContext, SubjectType, TokenType
|
|
from models.oauth import OAuthAccessToken
|
|
from services.oauth_device_flow import (
|
|
list_active_sessions,
|
|
revoke_oauth_token,
|
|
subject_match_clauses,
|
|
token_belongs_to_subject,
|
|
)
|
|
|
|
ACCOUNT_ID = uuid.UUID("11111111-1111-1111-1111-111111111111")
|
|
OTHER_ACCOUNT_ID = uuid.UUID("22222222-2222-2222-2222-222222222222")
|
|
TOKEN_ID = uuid.UUID("33333333-3333-3333-3333-333333333333")
|
|
OTHER_TOKEN_ID = uuid.UUID("44444444-4444-4444-4444-444444444444")
|
|
|
|
|
|
def _token(
|
|
*,
|
|
token_id: uuid.UUID = TOKEN_ID,
|
|
account_id: uuid.UUID | None = ACCOUNT_ID,
|
|
subject_email: str = "user@example.com",
|
|
subject_issuer: str = "dify:account",
|
|
token_hash: str | None = "live-hash",
|
|
expires_at: datetime | None = None,
|
|
revoked_at: datetime | None = None,
|
|
created_at: datetime | None = None,
|
|
) -> OAuthAccessToken:
|
|
token = OAuthAccessToken(
|
|
subject_email=subject_email,
|
|
subject_issuer=subject_issuer,
|
|
account_id=str(account_id) if account_id is not None else None,
|
|
client_id="difyctl",
|
|
device_label="test-device",
|
|
prefix="dfoa_" if account_id is not None else "dfoe_",
|
|
token_hash=token_hash,
|
|
expires_at=expires_at or datetime.now(UTC) + timedelta(days=1),
|
|
revoked_at=revoked_at,
|
|
)
|
|
token.id = str(token_id)
|
|
if created_at is not None:
|
|
token.created_at = created_at
|
|
return token
|
|
|
|
|
|
def _account_ctx(*, account_id: uuid.UUID = ACCOUNT_ID) -> AuthContext:
|
|
return AuthContext(
|
|
subject_type=SubjectType.ACCOUNT,
|
|
subject_email="user@example.com",
|
|
subject_issuer="dify:account",
|
|
account_id=account_id,
|
|
client_id="difyctl",
|
|
scopes=frozenset({"full"}),
|
|
token_id=uuid.uuid4(),
|
|
token_type=TokenType.OAUTH_ACCOUNT,
|
|
expires_at=None,
|
|
token_hash="h1",
|
|
verified_tenants={},
|
|
)
|
|
|
|
|
|
def _sso_ctx() -> AuthContext:
|
|
return AuthContext(
|
|
subject_type=SubjectType.EXTERNAL_SSO,
|
|
subject_email="sso@partner.com",
|
|
subject_issuer="https://idp.partner.com",
|
|
account_id=None,
|
|
client_id="difyctl",
|
|
scopes=frozenset({"apps:run"}),
|
|
token_id=uuid.uuid4(),
|
|
token_type=TokenType.OAUTH_EXTERNAL_SSO,
|
|
expires_at=None,
|
|
token_hash="h1",
|
|
verified_tenants={},
|
|
)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# subject_match_clauses
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_subject_match_clauses_account_matches_only_account_id():
|
|
clauses = subject_match_clauses(_account_ctx())
|
|
assert len(clauses) == 1
|
|
assert "account_id" in str(clauses[0])
|
|
|
|
|
|
def test_subject_match_clauses_external_sso_requires_null_account_id():
|
|
"""External SSO must additionally require ``account_id IS NULL`` so a
|
|
same-email account-flow row from a federated tenant cannot be
|
|
enumerated/revoked through an SSO bearer.
|
|
"""
|
|
clauses = subject_match_clauses(_sso_ctx())
|
|
assert len(clauses) == 3
|
|
rendered = " ".join(str(c) for c in clauses)
|
|
assert "subject_email" in rendered
|
|
assert "subject_issuer" in rendered
|
|
assert "account_id IS NULL" in rendered
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# revoke_oauth_token
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@pytest.mark.parametrize("sqlite_session", [(OAuthAccessToken,)], indirect=True)
|
|
def test_revoke_oauth_token_invalidates_redis_cache_when_live_hash_seen(sqlite_session: Session):
|
|
"""Happy path: snapshot finds a live ``token_hash`` → UPDATE runs +
|
|
Redis cache entry is DEL'd so the next bearer probe re-reads the now
|
|
revoked row from DB.
|
|
"""
|
|
sqlite_session.add(_token())
|
|
sqlite_session.commit()
|
|
|
|
redis = MagicMock()
|
|
|
|
revoke_oauth_token(redis, str(TOKEN_ID), session=sqlite_session)
|
|
|
|
assert not sqlite_session.in_transaction()
|
|
persisted = sqlite_session.get(OAuthAccessToken, str(TOKEN_ID))
|
|
assert persisted is not None
|
|
assert persisted.token_hash is None
|
|
assert persisted.revoked_at is not None
|
|
redis.delete.assert_called_once_with(TOKEN_CACHE_KEY_FMT.format(hash="live-hash"))
|
|
|
|
|
|
@pytest.mark.parametrize("sqlite_session", [(OAuthAccessToken,)], indirect=True)
|
|
def test_revoke_oauth_token_is_idempotent_when_already_revoked(sqlite_session: Session):
|
|
"""Second call (or race-loser): no live hash → UPDATE still runs (it
|
|
is itself idempotent thanks to ``WHERE revoked_at IS NULL``) but the
|
|
Redis invalidation is skipped because there's no cache entry to
|
|
drop.
|
|
"""
|
|
revoked_at = datetime.now(UTC) - timedelta(minutes=1)
|
|
sqlite_session.add(_token(token_hash=None, revoked_at=revoked_at))
|
|
sqlite_session.commit()
|
|
|
|
redis = MagicMock()
|
|
|
|
revoke_oauth_token(redis, str(TOKEN_ID), session=sqlite_session)
|
|
|
|
assert not sqlite_session.in_transaction()
|
|
persisted = sqlite_session.get(OAuthAccessToken, str(TOKEN_ID))
|
|
assert persisted is not None
|
|
assert persisted.token_hash is None
|
|
assert persisted.revoked_at is not None
|
|
redis.delete.assert_not_called()
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# list_active_sessions / token_belongs_to_subject
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@pytest.mark.parametrize("sqlite_session", [(OAuthAccessToken,)], indirect=True)
|
|
def test_list_active_sessions_returns_only_live_subject_tokens(sqlite_session: Session):
|
|
"""Only live, hashed rows for the authenticated subject are returned newest-first."""
|
|
|
|
now = datetime.now(UTC)
|
|
active_new = _token(token_id=TOKEN_ID, created_at=now - timedelta(minutes=1))
|
|
active_old = _token(token_id=OTHER_TOKEN_ID, created_at=now - timedelta(minutes=2))
|
|
expired = _token(
|
|
token_id=uuid.UUID(int=5),
|
|
expires_at=now - timedelta(seconds=1),
|
|
created_at=now - timedelta(minutes=3),
|
|
)
|
|
revoked = _token(
|
|
token_id=uuid.UUID(int=6),
|
|
token_hash=None,
|
|
revoked_at=now - timedelta(seconds=1),
|
|
created_at=now - timedelta(minutes=4),
|
|
)
|
|
hashless = _token(
|
|
token_id=uuid.UUID(int=7),
|
|
token_hash=None,
|
|
created_at=now - timedelta(minutes=5),
|
|
)
|
|
other_account = _token(
|
|
token_id=uuid.UUID(int=8),
|
|
account_id=OTHER_ACCOUNT_ID,
|
|
created_at=now - timedelta(minutes=6),
|
|
)
|
|
external_sso = _token(
|
|
token_id=uuid.UUID(int=9),
|
|
account_id=None,
|
|
subject_email="user@example.com",
|
|
subject_issuer="https://idp.example.com",
|
|
created_at=now - timedelta(minutes=7),
|
|
)
|
|
sqlite_session.add_all([active_new, active_old, expired, revoked, hashless, other_account, external_sso])
|
|
sqlite_session.commit()
|
|
|
|
out = list_active_sessions(_account_ctx(), now, session=sqlite_session)
|
|
|
|
assert [token.id for token in out] == [str(TOKEN_ID), str(OTHER_TOKEN_ID)]
|
|
|
|
|
|
@pytest.mark.parametrize("sqlite_session", [(OAuthAccessToken,)], indirect=True)
|
|
def test_token_belongs_to_subject_true_when_row_present(sqlite_session: Session):
|
|
sqlite_session.add(_token())
|
|
sqlite_session.commit()
|
|
|
|
assert token_belongs_to_subject(str(TOKEN_ID), _account_ctx(), session=sqlite_session) is True
|
|
|
|
|
|
@pytest.mark.parametrize("sqlite_session", [(OAuthAccessToken,)], indirect=True)
|
|
def test_token_belongs_to_subject_false_for_other_account(sqlite_session: Session):
|
|
sqlite_session.add(_token(account_id=OTHER_ACCOUNT_ID))
|
|
sqlite_session.commit()
|
|
|
|
assert token_belongs_to_subject(str(TOKEN_ID), _account_ctx(), session=sqlite_session) is False
|