diff --git a/api/.importlinter b/api/.importlinter index df3929e6f89..3933ee443ed 100644 --- a/api/.importlinter +++ b/api/.importlinter @@ -203,6 +203,7 @@ forbidden_modules = name = Account application services and contracts are framework and persistence neutral type = forbidden source_modules = + services.account_access_service services.account_avatar_service services.account_change_email_ports services.account_change_email_service @@ -219,6 +220,7 @@ source_modules = services.account_password_service services.account_ports services.account_profile_service + services.entities.account_access_entities services.entities.account_entities services.entities.account_login_entities services.entities.account_oauth_entities diff --git a/api/controllers/openapi/_contract.py b/api/controllers/openapi/_contract.py index a7dcf9093da..aa92cdf411d 100644 --- a/api/controllers/openapi/_contract.py +++ b/api/controllers/openapi/_contract.py @@ -5,8 +5,9 @@ reference — emitting the Swagger schema AND doing the runtime validation/ serialisation — so the advertised and enforced contracts can't drift. Validation failures map to a single shape: 422. -They must sit BELOW ``@auth_router.guard`` so auth runs before validation and the -``view.__wrapped__`` unit-test seam unwraps exactly the guard layer. +They must sit below route admission (or a direct ``@auth_router.guard``) so +authentication runs before validation and the ``view.__wrapped__`` unit-test +seam can bypass the outer admission layer. """ from __future__ import annotations diff --git a/api/controllers/openapi/account.py b/api/controllers/openapi/account.py index f1b02ef115e..6ddd9a14ef2 100644 --- a/api/controllers/openapi/account.py +++ b/api/controllers/openapi/account.py @@ -1,12 +1,11 @@ from __future__ import annotations from datetime import UTC, datetime +from uuid import UUID from flask_restx import Resource -from sqlalchemy.orm import Session -from werkzeug.exceptions import NotFound +from werkzeug.exceptions import NotFound, Unauthorized -from controllers.common.session import with_session from controllers.openapi import openapi_ns from controllers.openapi._contract import accepts, returns from controllers.openapi._models import ( @@ -18,114 +17,77 @@ from controllers.openapi._models import ( SessionRow, WorkspacePayload, ) -from controllers.openapi.auth.composition import auth_router -from controllers.openapi.auth.data import AuthData -from extensions.ext_redis import redis_client -from libs.oauth_bearer import ( - Scope, - TokenType, - get_auth_ctx, -) -from libs.rate_limit import ( - LIMIT_ME_PER_ACCOUNT, - enforce, -) -from services.account_service import AccountService, TenantService -from services.oauth_device_flow import ( - list_active_sessions, - revoke_oauth_token, - token_belongs_to_subject, -) +from controllers.openapi.flask_admission import openapi_account_admission +from extensions.ext_application_services import application_services +from libs.oauth_bearer import Scope +from libs.rate_limit import LIMIT_ME_PER_ACCOUNT +from machinery.context import AccountRequestContext +from services.account_errors import AccountNotFoundError, AccountSessionNotFoundError +from services.entities.account_access_entities import AccountSessionSnapshot, AccountWorkspaceSnapshot +from services.entities.account_entities import AccountSnapshot @openapi_ns.route("/account") class AccountApi(Resource): - @auth_router.guard(scope=Scope.FULL, allowed_token_types=frozenset({TokenType.OAUTH_ACCOUNT})) + @openapi_account_admission(scope=Scope.FULL, rate_limit=LIMIT_ME_PER_ACCOUNT) @returns(200, AccountResponse, description="Account info") - @with_session(write=False) - def get(self, session: Session, *, auth_data: AuthData): - enforce(LIMIT_ME_PER_ACCOUNT, key=f"account:{auth_data.account_id}") - - account_id_str = str(auth_data.account_id) if auth_data.account_id else None - account = AccountService.get_account_by_id(account_id_str, session=session) if account_id_str else None - memberships = TenantService.get_account_memberships(account_id_str, session=session) if account_id_str else [] - default_ws_id = _pick_default_workspace(memberships) - + def get(self, request_context: AccountRequestContext): + try: + snapshot = application_services().accounts.access.get(request_context) + except AccountNotFoundError: + raise Unauthorized("account not found") from None return AccountResponse( subject_type="account", - subject_email=account.email if account else None, - account=_account_payload(account) if account else None, - workspaces=[_workspace_payload(m) for m in memberships], - default_workspace_id=default_ws_id, + subject_email=snapshot.account.email, + account=_account_payload(snapshot.account), + workspaces=[_workspace_payload(workspace) for workspace in snapshot.workspaces], + default_workspace_id=snapshot.default_workspace_id, ) @openapi_ns.route("/account/sessions/self") class AccountSessionsSelfApi(Resource): - @auth_router.guard(scope=Scope.FULL, allowed_token_types=frozenset({TokenType.OAUTH_ACCOUNT})) + @openapi_account_admission(scope=Scope.FULL) @returns(200, RevokeResponse, description="Session revoked") - @with_session - def delete(self, session: Session, *, auth_data: AuthData): - revoke_oauth_token(redis_client, str(auth_data.token_id), session=session) + def delete(self, request_context: AccountRequestContext): + application_services().accounts.access.revoke_current_session(request_context) return RevokeResponse(status="revoked") @openapi_ns.route("/account/sessions") class AccountSessionsApi(Resource): - @auth_router.guard(scope=Scope.FULL, allowed_token_types=frozenset({TokenType.OAUTH_ACCOUNT})) + @openapi_account_admission(scope=Scope.FULL) @returns(200, SessionListResponse, description="Session list") @accepts(query=SessionListQuery) - @with_session(write=False) - def get(self, session: Session, *, auth_data: AuthData, query: SessionListQuery): - # SessionListQuery enforces the advertised bounds (extra='forbid', page>=1, - # 1<=limit<=MAX_PAGE_LIMIT) so the server rejects out-of-range paging rather - # than silently coercing (e.g. page=0 -> empty slice). - ctx = get_auth_ctx() - now = datetime.now(UTC) - page = query.page - limit = query.limit - - all_rows = list_active_sessions(ctx, now, session=session) - - total = len(all_rows) - sliced = all_rows[(page - 1) * limit : page * limit] - - items = [ - SessionRow( - id=str(r.id), - prefix=r.prefix, - client_id=r.client_id, - device_label=r.device_label, - created_at=_iso(r.created_at), - last_used_at=_iso(r.last_used_at), - expires_at=_iso(r.expires_at), - ) - for r in sliced - ] - + def get(self, request_context: AccountRequestContext, *, query: SessionListQuery): + page = application_services().accounts.access.list_sessions( + request_context, + page=query.page, + limit=query.limit, + ) return SessionListResponse( - page=page, - limit=limit, - total=total, - has_more=page * limit < total, - data=items, + page=page.page, + limit=page.limit, + total=page.total, + has_more=page.has_more, + data=[_session_row(session) for session in page.items], ) @openapi_ns.route("/account/sessions/") class AccountSessionByIdApi(Resource): - @auth_router.guard(scope=Scope.FULL, allowed_token_types=frozenset({TokenType.OAUTH_ACCOUNT})) + @openapi_account_admission(scope=Scope.FULL) @returns(200, RevokeResponse, description="Session revoked") - @with_session - def delete(self, session: Session, session_id: str, *, auth_data: AuthData): - ctx = get_auth_ctx() - - # 404 (not 403) on cross-subject so the endpoint doesn't leak - # token IDs that belong to other subjects. - if not token_belongs_to_subject(session_id, ctx, session=session): - raise NotFound("session not found") - - revoke_oauth_token(redis_client, session_id, session=session) + def delete(self, request_context: AccountRequestContext, session_id: str): + try: + token_id = str(UUID(session_id)) + except ValueError: + raise NotFound("session not found") from None + try: + application_services().accounts.access.revoke_session(request_context, token_id=token_id) + except AccountSessionNotFoundError: + # Do not reveal whether a token ID belongs to another account. + raise NotFound("session not found") from None return RevokeResponse(status="revoked") @@ -137,19 +99,21 @@ def _iso(dt: datetime | None) -> str | None: return dt.isoformat().replace("+00:00", "Z") -def _pick_default_workspace(memberships) -> str | None: - if not memberships: - return None - for join, tenant in memberships: - if getattr(join, "current", False): - return str(tenant.id) - return str(memberships[0][1].id) +def _session_row(session: AccountSessionSnapshot) -> SessionRow: + return SessionRow( + id=session.id, + prefix=session.prefix, + client_id=session.client_id, + device_label=session.device_label, + created_at=_iso(session.created_at), + last_used_at=_iso(session.last_used_at), + expires_at=_iso(session.expires_at), + ) -def _workspace_payload(row) -> WorkspacePayload: - join, tenant = row - return WorkspacePayload(id=str(tenant.id), name=tenant.name, role=getattr(join, "role", "")) +def _workspace_payload(workspace: AccountWorkspaceSnapshot) -> WorkspacePayload: + return WorkspacePayload(id=workspace.id, name=workspace.name, role=workspace.role) -def _account_payload(account) -> AccountPayload: - return AccountPayload(id=str(account.id), email=account.email, name=account.name) +def _account_payload(account: AccountSnapshot) -> AccountPayload: + return AccountPayload(id=account.id, email=account.email, name=account.name) diff --git a/api/controllers/openapi/auth/pipeline.py b/api/controllers/openapi/auth/pipeline.py index 6a68e5c91fe..9f14105a942 100644 --- a/api/controllers/openapi/auth/pipeline.py +++ b/api/controllers/openapi/auth/pipeline.py @@ -119,8 +119,8 @@ class PipelineRouter: """Entry point for openapi auth. `guard()` is the decorator that endpoints attach to. It applies - global gates (edition, token type) then dispatches to the matching - `PipelineRoute` for the token type. + global gates (edition, license, token type) then dispatches to the + matching `PipelineRoute` for the token type. """ def __init__(self, routes: dict[TokenType, PipelineRoute]) -> None: @@ -132,6 +132,7 @@ class PipelineRouter: scope: Scope | None = None, allowed_token_types: frozenset[TokenType] | None = None, edition: frozenset[DeploymentEdition] | None = None, + require_valid_enterprise_license: bool = False, workspace_membership: bool = False, allowed_roles: frozenset[TenantAccountRole] | None = None, rbac: RBACRequirement | None = None, @@ -140,6 +141,7 @@ class PipelineRouter: scope=scope, allowed_token_types=allowed_token_types, edition=edition, + require_valid_enterprise_license=require_valid_enterprise_license, workspace_membership=workspace_membership, allowed_roles=allowed_roles, rbac=rbac, @@ -151,6 +153,7 @@ class PipelineRouter: scope: Scope | None = None, allowed_token_types: frozenset[TokenType] | None = None, edition: frozenset[DeploymentEdition] | None = None, + require_valid_enterprise_license: bool = False, allowed_roles: frozenset[TenantAccountRole] | None = None, rbac: RBACRequirement | None = None, ) -> Callable: @@ -158,6 +161,7 @@ class PipelineRouter: scope=scope, allowed_token_types=allowed_token_types, edition=edition, + require_valid_enterprise_license=require_valid_enterprise_license, workspace_membership=True, allowed_roles=allowed_roles, rbac=rbac, @@ -169,6 +173,7 @@ class PipelineRouter: scope: Scope | None, allowed_token_types: frozenset[TokenType] | None, edition: frozenset[DeploymentEdition] | None, + require_valid_enterprise_license: bool, workspace_membership: bool, allowed_roles: frozenset[TenantAccountRole] | None, rbac: RBACRequirement | None, @@ -183,6 +188,7 @@ class PipelineRouter: scope=scope, allowed_token_types=allowed_token_types, edition=edition, + require_valid_enterprise_license=require_valid_enterprise_license, workspace_membership=workspace_membership, allowed_roles=allowed_roles, rbac=rbac, @@ -201,6 +207,7 @@ class PipelineRouter: scope: Scope | None, allowed_token_types: frozenset[TokenType] | None, edition: frozenset[DeploymentEdition] | None, + require_valid_enterprise_license: bool, workspace_membership: bool = False, allowed_roles: frozenset[TenantAccountRole] | None = None, rbac: RBACRequirement | None = None, @@ -210,7 +217,9 @@ class PipelineRouter: raise NotFound() license_checked = False - if edition is not None and DeploymentEdition.ENTERPRISE in edition: + if dify_config.DEPLOYMENT_EDITION == DeploymentEdition.ENTERPRISE and ( + require_valid_enterprise_license or (edition is not None and DeploymentEdition.ENTERPRISE in edition) + ): _check_license() license_checked = True diff --git a/api/controllers/openapi/flask_admission.py b/api/controllers/openapi/flask_admission.py new file mode 100644 index 00000000000..7e485c70db6 --- /dev/null +++ b/api/controllers/openapi/flask_admission.py @@ -0,0 +1,88 @@ +"""Flask adapter for account-authenticated OpenAPI admission.""" + +from __future__ import annotations + +from collections.abc import Callable +from functools import wraps +from typing import Concatenate + +from flask import Response, request +from werkzeug.exceptions import Unauthorized + +from controllers.openapi.auth.composition import auth_router +from controllers.openapi.auth.data import AuthData +from core.logging.context import get_request_id, get_trace_id +from enums import DeploymentEdition +from libs.oauth_bearer import Scope, TokenType +from libs.rate_limit import RateLimit, enforce +from machinery.context import AccountRequestContext +from models.account import Account, AccountStatus + + +def openapi_account_admission[T, **P, R]( + *, + scope: Scope, + editions: frozenset[DeploymentEdition] | None = None, + require_initialized: bool = True, + require_valid_enterprise_license: bool = True, + rate_limit: RateLimit | None = None, +) -> Callable[ + [Callable[Concatenate[T, AccountRequestContext, P], R]], + Callable[Concatenate[T, P], R | Response], +]: + """Authenticate an account bearer and inject framework-neutral identity. + + Client-version admission remains attached to the OpenAPI blueprint so it + can also reject requests for removed routes. Edition and Enterprise + license checks are delegated to the shared auth router before the stable + context is constructed. + """ + + def decorator( + view: Callable[Concatenate[T, AccountRequestContext, P], R], + ) -> Callable[Concatenate[T, P], R | Response]: + @wraps(view) + def inject_request_context( + self: T, + /, + *args: P.args, + **kwargs: P.kwargs, + ) -> R: + auth_data = kwargs.pop("auth_data", None) + if not isinstance(auth_data, AuthData): + raise RuntimeError("OpenAPI auth pipeline did not provide valid AuthData") + account = auth_data.caller + if not isinstance(account, Account) or auth_data.account_id is None: + raise Unauthorized("account not found") + if require_initialized and account.status == AccountStatus.UNINITIALIZED: + raise Unauthorized("account not initialized") + + account_id = str(auth_data.account_id) + if rate_limit is not None: + enforce(rate_limit, key=f"account:{account_id}") + + context = AccountRequestContext( + request_id=get_request_id(), + trace_id=get_trace_id() or request.headers.get("X-Trace-Id"), + account_id=account_id, + access_token_id=str(auth_data.token_id) if auth_data.token_id is not None else None, + ) + return view(self, context, *args, **kwargs) + + authenticated = auth_router.guard( + scope=scope, + allowed_token_types=frozenset({TokenType.OAUTH_ACCOUNT}), + edition=editions, + require_valid_enterprise_license=require_valid_enterprise_license, + )(inject_request_context) + + # Keep one stable test seam: one ``__wrapped__`` skips route admission + # and reaches input parsing/response handling. Client-version admission + # stays blueprint-wide by design. + @wraps(view) + def admitted(self: T, /, *args: P.args, **kwargs: P.kwargs) -> R | Response: + return authenticated(self, *args, **kwargs) + + return admitted + + return decorator diff --git a/api/extensions/ext_application_services.py b/api/extensions/ext_application_services.py index 7863e5a5b72..f2328234f19 100644 --- a/api/extensions/ext_application_services.py +++ b/api/extensions/ext_application_services.py @@ -5,6 +5,7 @@ import time from collections.abc import Mapping from dataclasses import dataclass from datetime import UTC, datetime +from functools import partial from typing import cast from uuid import uuid4 @@ -23,9 +24,10 @@ from core.tools.tool_file_manager import ToolFileManager from enums import DeploymentEdition, WebAppAccessMode from extensions.ext_redis import RedisClientWrapper, redis_client from extensions.ext_storage import storage -from libs.datetime_utils import naive_utc_now +from libs.datetime_utils import naive_utc_now, utc_now from libs.helper import RateLimiter from libs.oauth import GitHubOAuth, GoogleOAuth +from libs.oauth_bearer import invalidate_oauth_token_cache from libs.passport import PassportService from repositories.account_activation_repository import SQLAlchemyAccountActivationRepository from repositories.account_integration_repository import SQLAlchemyAccountIntegrationRepository @@ -44,6 +46,7 @@ from repositories.explore_banner_query_repository import ExploreBannerQueryRepos from repositories.factory import DifyAPIRepositoryFactory from repositories.file_grant_repository import FileGrantRepository from repositories.installation_state_repository import InstallationStateRepository +from repositories.oauth_access_token_repository import SQLAlchemyOAuthAccessTokenRepository from repositories.oauth_server_repository import RedisOAuthServerTokenRepository, SQLAlchemyOAuthServerRepository from repositories.recommended_app_catalog_repository import DatabaseRecommendedAppCatalogRepository from repositories.step_by_step_tour_repository import SQLAlchemyStepByStepTourStateRepository @@ -55,6 +58,7 @@ from repositories.webapp_access_query_repository import WebAppAccessQueryReposit from repositories.workflow_run_archive_repository import WorkflowRunArchiveBundleQueryRepository from repositories.workspace_member_query_repository import WorkspaceMemberQueryRepository from repositories.workspace_query_repository import WorkspaceQueryRepository +from services.account_access_service import AccountAccessService from services.account_activation_service import AccountActivationService from services.account_adapters import ( BillingAccountActivationEligibility, @@ -211,6 +215,7 @@ def _is_user_allowed_to_access_webapp(user_id: str, app_id: str) -> bool: @dataclass(frozen=True, slots=True) class AccountServices: + access: AccountAccessService authentication: ConsoleAuthenticationService avatar: AccountAvatarService change_email: AccountChangeEmailService @@ -411,6 +416,13 @@ def build_application_services( account_provisioning = SQLAlchemyConsoleAuthProvisioningGateway(session_factory=database_client) return ApplicationServices( accounts=AccountServices( + access=AccountAccessService( + accounts=accounts, + workspaces=workspace_query_repository, + sessions=SQLAlchemyOAuthAccessTokenRepository(session_factory=database_client), + invalidate_token_cache=partial(invalidate_oauth_token_cache, redis), + now=utc_now, + ), authentication=ConsoleAuthenticationService( accounts=accounts, workspaces=workspace_query_repository, diff --git a/api/libs/datetime_utils.py b/api/libs/datetime_utils.py index d962d81e78a..a158bb1d5ea 100644 --- a/api/libs/datetime_utils.py +++ b/api/libs/datetime_utils.py @@ -17,11 +17,16 @@ class _NowFunction(Protocol): _now_func: _NowFunction = datetime.datetime.now +def utc_now() -> datetime.datetime: + """Return a timezone-aware datetime representing the current UTC time.""" + return _now_func(datetime.UTC) + + def naive_utc_now() -> datetime.datetime: """Return a naive datetime object (without timezone information) representing current UTC time. """ - return _now_func(datetime.UTC).replace(tzinfo=None) + return utc_now().replace(tzinfo=None) def ensure_naive_utc(dt: datetime.datetime) -> datetime.datetime: diff --git a/api/libs/oauth_bearer.py b/api/libs/oauth_bearer.py index ed17503bb3b..89ede4b8a12 100644 --- a/api/libs/oauth_bearer.py +++ b/api/libs/oauth_bearer.py @@ -318,6 +318,14 @@ AUDIT_OAUTH_EXPIRED = "oauth.token_expired" ScopeVariant = Literal["account", "external_sso"] +class _TokenCacheClient(Protocol): + def delete(self, *names: str | bytes) -> object: ... + + +def invalidate_oauth_token_cache(client: _TokenCacheClient, token_hash: str) -> None: + client.delete(TOKEN_CACHE_KEY_FMT.format(hash=token_hash)) + + class OAuthAccessTokenResolver: """``.for_account()`` / ``.for_external_sso()`` are variant-scoped views sharing DB + cache plumbing. @@ -385,7 +393,7 @@ class OAuthAccessTokenResolver: row_id, extra={"audit": True, "token_id": str(row_id)}, ) - self._redis.delete(self._cache_key(token_hash)) + invalidate_oauth_token_cache(self._redis, token_hash) self.cache_set_negative(token_hash) diff --git a/api/machinery/context.py b/api/machinery/context.py index 5d6f5436619..4facf6caf9b 100644 --- a/api/machinery/context.py +++ b/api/machinery/context.py @@ -8,3 +8,12 @@ class RequestContext(NamedTuple): trace_id: str | None account_id: str active_workspace_id: str + + +class AccountRequestContext(NamedTuple): + """Stable identity for account-scoped use cases that do not require a workspace.""" + + request_id: str + trace_id: str | None + account_id: str + access_token_id: str | None = None diff --git a/api/repositories/oauth_access_token_repository.py b/api/repositories/oauth_access_token_repository.py new file mode 100644 index 00000000000..87ba4704093 --- /dev/null +++ b/api/repositories/oauth_access_token_repository.py @@ -0,0 +1,89 @@ +"""SQLAlchemy persistence for account-scoped OAuth access sessions.""" + +from __future__ import annotations + +from datetime import datetime +from typing import override + +from sqlalchemy import func, select, update +from sqlalchemy.orm import Session, sessionmaker + +from models.oauth import OAuthAccessToken +from services.account_ports import AccountSessionRepository +from services.entities.account_access_entities import ( + AccountSessionRevocation, + AccountSessionSnapshot, +) + + +class SQLAlchemyOAuthAccessTokenRepository(AccountSessionRepository): + def __init__(self, *, session_factory: sessionmaker[Session]) -> None: + self._session_factory = session_factory + + @override + def list_active( + self, + *, + account_id: str, + active_at: datetime, + offset: int, + limit: int, + ) -> tuple[int, tuple[AccountSessionSnapshot, ...]]: + predicates = ( + OAuthAccessToken.account_id == account_id, + OAuthAccessToken.revoked_at.is_(None), + OAuthAccessToken.token_hash.is_not(None), + OAuthAccessToken.expires_at > active_at, + ) + with self._session_factory() as session: + total = session.scalar(select(func.count()).select_from(OAuthAccessToken).where(*predicates)) or 0 + rows = session.scalars( + select(OAuthAccessToken) + .where(*predicates) + .order_by(OAuthAccessToken.created_at.desc(), OAuthAccessToken.id.desc()) + .offset(offset) + .limit(limit) + ).all() + return int(total), tuple(self._to_snapshot(row) for row in rows) + + @override + def revoke( + self, + *, + account_id: str, + token_id: str, + revoked_at: datetime, + ) -> AccountSessionRevocation: + with self._session_factory.begin() as session: + row = session.execute( + select(OAuthAccessToken.account_id, OAuthAccessToken.token_hash) + .where(OAuthAccessToken.id == token_id) + .with_for_update() + ).one_or_none() + if row is None or row.account_id != account_id: + return AccountSessionRevocation(owned=False) + + token_hash = row.token_hash + if token_hash is not None: + session.execute( + update(OAuthAccessToken) + .where( + OAuthAccessToken.id == token_id, + OAuthAccessToken.account_id == account_id, + OAuthAccessToken.revoked_at.is_(None), + ) + .values(revoked_at=revoked_at, token_hash=None) + ) + return AccountSessionRevocation(owned=True, token_hash=token_hash) + + @staticmethod + def _to_snapshot(row: OAuthAccessToken) -> AccountSessionSnapshot: + return AccountSessionSnapshot( + id=str(row.id), + prefix=row.prefix, + client_id=row.client_id, + device_label=row.device_label, + created_at=row.created_at, + last_used_at=row.last_used_at, + expires_at=row.expires_at, + ) diff --git a/api/repositories/workspace_query_repository.py b/api/repositories/workspace_query_repository.py index 84758d997be..8e83d158c47 100644 --- a/api/repositories/workspace_query_repository.py +++ b/api/repositories/workspace_query_repository.py @@ -7,11 +7,17 @@ from sqlalchemy.orm import Session, sessionmaker from models.account import Tenant, TenantAccountJoin, TenantStatus from services.account_login_service import ConsoleAuthWorkspaceQuery -from services.account_ports import AccountWorkspaceMembershipQuery +from services.account_ports import AccountWorkspaceMembershipQuery, AccountWorkspaceSnapshotQuery +from services.entities.account_access_entities import AccountWorkspaceSnapshot from services.workspace_query_service import WorkspaceQuery, WorkspaceRecord -class WorkspaceQueryRepository(WorkspaceQuery, AccountWorkspaceMembershipQuery, ConsoleAuthWorkspaceQuery): +class WorkspaceQueryRepository( + WorkspaceQuery, + AccountWorkspaceMembershipQuery, + AccountWorkspaceSnapshotQuery, + ConsoleAuthWorkspaceQuery, +): def __init__(self, session_factory: sessionmaker[Session]) -> None: self._session_factory = session_factory @@ -52,6 +58,35 @@ class WorkspaceQueryRepository(WorkspaceQuery, AccountWorkspaceMembershipQuery, with self._session_factory() as session: return tuple(session.scalars(stmt).all()) + @override + def list_account_access_workspaces(self, account_id: str) -> tuple[AccountWorkspaceSnapshot, ...]: + """List every membership for the OpenAPI account identity response. + + Unlike the Console workspace picker, the identity response preserves + its existing behavior of including archived memberships. + """ + stmt = ( + select( + Tenant.id, + Tenant.name, + TenantAccountJoin.role, + TenantAccountJoin.current, + ) + .join(TenantAccountJoin, TenantAccountJoin.tenant_id == Tenant.id) + .where(TenantAccountJoin.account_id == account_id) + .order_by(Tenant.created_at.asc(), Tenant.id.asc()) + ) + with self._session_factory() as session: + return tuple( + AccountWorkspaceSnapshot( + id=workspace_id, + name=name, + role=role.value, + current=current, + ) + for workspace_id, name, role, current in session.execute(stmt).all() + ) + @override def has_active_for_account(self, account_id: str) -> bool: stmt = ( diff --git a/api/services/account_access_service.py b/api/services/account_access_service.py new file mode 100644 index 00000000000..9fb3a768cb3 --- /dev/null +++ b/api/services/account_access_service.py @@ -0,0 +1,77 @@ +"""Application service for account identity and access-session use cases.""" + +from __future__ import annotations + +from collections.abc import Callable +from datetime import datetime + +from machinery.context import AccountRequestContext +from services.account_errors import AccountNotFoundError, AccountSessionNotFoundError +from services.account_ports import ( + AccountSessionRepository, + AccountSnapshotQuery, + AccountTokenCacheInvalidator, + AccountWorkspaceSnapshotQuery, +) +from services.entities.account_access_entities import AccountAccessSnapshot, AccountSessionPage + + +class AccountAccessService: + def __init__( + self, + *, + accounts: AccountSnapshotQuery, + workspaces: AccountWorkspaceSnapshotQuery, + sessions: AccountSessionRepository, + invalidate_token_cache: AccountTokenCacheInvalidator, + now: Callable[[], datetime], + ) -> None: + self._accounts = accounts + self._workspaces = workspaces + self._sessions = sessions + self._invalidate_token_cache = invalidate_token_cache + self._now = now + + def get(self, context: AccountRequestContext) -> AccountAccessSnapshot: + account = self._accounts.get(context.account_id) + if account is None: + raise AccountNotFoundError + + workspaces = tuple(self._workspaces.list_account_access_workspaces(context.account_id)) + default_workspace_id = next((workspace.id for workspace in workspaces if workspace.current), None) + if default_workspace_id is None and workspaces: + default_workspace_id = workspaces[0].id + + return AccountAccessSnapshot( + account=account, + workspaces=workspaces, + default_workspace_id=default_workspace_id, + ) + + def list_sessions(self, context: AccountRequestContext, *, page: int, limit: int) -> AccountSessionPage: + total, sessions = self._sessions.list_active( + account_id=context.account_id, + active_at=self._now(), + offset=(page - 1) * limit, + limit=limit, + ) + return AccountSessionPage(page=page, limit=limit, total=total, items=tuple(sessions)) + + def revoke_current_session(self, context: AccountRequestContext) -> None: + if context.access_token_id is None: + raise RuntimeError("OpenAPI account admission did not resolve an access token") + self._revoke(context, token_id=context.access_token_id, require_owned=False) + + def revoke_session(self, context: AccountRequestContext, *, token_id: str) -> None: + self._revoke(context, token_id=token_id, require_owned=True) + + def _revoke(self, context: AccountRequestContext, *, token_id: str, require_owned: bool) -> None: + revocation = self._sessions.revoke( + account_id=context.account_id, + token_id=token_id, + revoked_at=self._now(), + ) + if require_owned and not revocation.owned: + raise AccountSessionNotFoundError + if revocation.token_hash is not None: + self._invalidate_token_cache(revocation.token_hash) diff --git a/api/services/account_errors.py b/api/services/account_errors.py index 773ae90ea46..9ff9d5ec73a 100644 --- a/api/services/account_errors.py +++ b/api/services/account_errors.py @@ -9,6 +9,10 @@ class AccountNotFoundError(AccountApplicationError): """The admitted account no longer exists.""" +class AccountSessionNotFoundError(AccountApplicationError): + """The requested access session is not owned by the admitted account.""" + + class CurrentAccountPasswordIncorrectError(AccountApplicationError): """The supplied current password does not match the account credential.""" diff --git a/api/services/account_ports.py b/api/services/account_ports.py index 8ea4f06e074..bc0b5d1eeb9 100644 --- a/api/services/account_ports.py +++ b/api/services/account_ports.py @@ -1,9 +1,14 @@ -"""Persistence ports used by account application services.""" +"""Ports used by account application services.""" from collections.abc import Sequence from datetime import datetime from typing import Protocol +from services.entities.account_access_entities import ( + AccountSessionRevocation, + AccountSessionSnapshot, + AccountWorkspaceSnapshot, +) from services.entities.account_entities import ( AccountCredentials, AccountDeletionChallenge, @@ -17,6 +22,10 @@ from services.entities.account_entities import ( ) +class AccountSnapshotQuery(Protocol): + def get(self, account_id: str) -> AccountSnapshot | None: ... + + class AccountRepository(Protocol): def get(self, account_id: str) -> AccountSnapshot | None: ... @@ -58,6 +67,33 @@ class AccountWorkspaceMembershipQuery(Protocol): def has_active_membership(self, account_id: str) -> bool: ... +class AccountWorkspaceSnapshotQuery(Protocol): + def list_account_access_workspaces(self, account_id: str) -> Sequence[AccountWorkspaceSnapshot]: ... + + +class AccountSessionRepository(Protocol): + def list_active( + self, + *, + account_id: str, + active_at: datetime, + offset: int, + limit: int, + ) -> tuple[int, Sequence[AccountSessionSnapshot]]: ... + + def revoke( + self, + *, + account_id: str, + token_id: str, + revoked_at: datetime, + ) -> AccountSessionRevocation: ... + + +class AccountTokenCacheInvalidator(Protocol): + def __call__(self, token_hash: str) -> None: ... + + class AccountAvatarFileGateway(Protocol): def get_owned_signed_url(self, *, account_id: str, upload_file_id: str) -> str | None: ... diff --git a/api/services/account_service.py b/api/services/account_service.py index dae4ea2b4a6..e33af75e385 100644 --- a/api/services/account_service.py +++ b/api/services/account_service.py @@ -321,9 +321,9 @@ class AccountService: @staticmethod def get_account_by_id(account_id: str, *, session: Session) -> Account | None: """Plain ``Account`` getter — no banned check, no tenant rotation, - no ``last_active_at`` write. Use this from read-only identity - endpoints (``/openapi/v1/account``) where ``load_user``'s - side-effects (current-tenant assignment, commit) are unwanted. + no ``last_active_at`` write. Use this from authentication and read + paths where ``load_user``'s current-tenant assignment and commit are + unwanted. ``session`` is injected by the caller so this service stays free of a Flask-scoped session import. @@ -1212,36 +1212,12 @@ class TenantService: ).all() ) - @staticmethod - def get_account_memberships(account_id: str, *, session: Session) -> list[Row[tuple[TenantAccountJoin, Tenant]]]: - """Return ``(TenantAccountJoin, Tenant)`` rows for every workspace - the account belongs to. Unlike :meth:`get_join_tenants` this keeps - the join row so callers can read ``role``/``current`` alongside the - tenant — used by ``/openapi/v1/account`` to render workspace - membership + pick the default workspace. - - ``session`` is injected by the caller so this service stays free - of a Flask-scoped session import. - - No tenant-status filter: parity with the legacy controller query - (the openapi identity endpoint listed all joined tenants). - """ - return ( - session.query(TenantAccountJoin, Tenant) - .join(Tenant, Tenant.id == TenantAccountJoin.tenant_id) - .filter(TenantAccountJoin.account_id == account_id) - .all() - ) - @staticmethod def get_workspaces_for_account(account_id: str, *, session: Session) -> list[Row[tuple[Tenant, TenantAccountJoin]]]: """``(Tenant, TenantAccountJoin)`` rows for every workspace the account belongs to, ordered by ``Tenant.created_at`` ASC — the - canonical ordering for ``/openapi/v1/workspaces``. - - Distinct from :meth:`get_account_memberships`: tuple order is - flipped (tenant first) and rows are sorted, so the workspace - listing is stable across requests. + canonical ordering for ``/openapi/v1/workspaces``. Rows keep the + tenant first so callers can serialize the workspace directly. """ return list( session.execute( diff --git a/api/services/entities/account_access_entities.py b/api/services/entities/account_access_entities.py new file mode 100644 index 00000000000..a59db492259 --- /dev/null +++ b/api/services/entities/account_access_entities.py @@ -0,0 +1,52 @@ +"""Framework-neutral data contracts for account identity and access sessions.""" + +from __future__ import annotations + +from dataclasses import dataclass +from datetime import datetime + +from services.entities.account_entities import AccountSnapshot + + +@dataclass(frozen=True, slots=True) +class AccountWorkspaceSnapshot: + id: str + name: str + role: str + current: bool + + +@dataclass(frozen=True, slots=True) +class AccountAccessSnapshot: + account: AccountSnapshot + workspaces: tuple[AccountWorkspaceSnapshot, ...] + default_workspace_id: str | None + + +@dataclass(frozen=True, slots=True) +class AccountSessionSnapshot: + id: str + prefix: str + client_id: str + device_label: str + created_at: datetime | None + last_used_at: datetime | None + expires_at: datetime | None + + +@dataclass(frozen=True, slots=True) +class AccountSessionPage: + page: int + limit: int + total: int + items: tuple[AccountSessionSnapshot, ...] + + @property + def has_more(self) -> bool: + return self.page * self.limit < self.total + + +@dataclass(frozen=True, slots=True) +class AccountSessionRevocation: + owned: bool + token_hash: str | None = None diff --git a/api/services/entities/account_entities.py b/api/services/entities/account_entities.py index 3e4293871e8..cfc637b09e1 100644 --- a/api/services/entities/account_entities.py +++ b/api/services/entities/account_entities.py @@ -1,4 +1,4 @@ -"""Framework-neutral contracts for Console account use cases.""" +"""Framework-neutral contracts shared by account use cases.""" from __future__ import annotations diff --git a/api/services/oauth_device_flow.py b/api/services/oauth_device_flow.py index 9e59b8c326a..f3967026aad 100644 --- a/api/services/oauth_device_flow.py +++ b/api/services/oauth_device_flow.py @@ -10,12 +10,12 @@ import uuid from dataclasses import asdict, dataclass, field from datetime import UTC, datetime, timedelta from enum import StrEnum -from typing import Any, NotRequired, TypedDict +from typing import NotRequired, TypedDict -from sqlalchemy import and_, func, select, update +from sqlalchemy import func, select, update from sqlalchemy.orm import Session -from libs.oauth_bearer import TOKEN_CACHE_KEY_FMT, AuthContext, SubjectType +from libs.oauth_bearer import SubjectType, invalidate_oauth_token_cache from models.oauth import OAuthAccessToken logger = logging.getLogger(__name__) @@ -382,7 +382,7 @@ def mint_oauth_token( ) if outcome.rotated and outcome.old_hash: - redis_client.delete(TOKEN_CACHE_KEY_FMT.format(hash=outcome.old_hash)) + invalidate_oauth_token_cache(redis_client, outcome.old_hash) return MintResult(token=token, token_id=outcome.token_id, expires_at=expires_at) @@ -487,70 +487,3 @@ def oauth_ttl_days(tenant_id: str | None = None) -> int: logger.warning("%s=%d above max %d; clamping", _TTL_ENV_VAR, value, MAX_TTL_DAYS) return MAX_TTL_DAYS return value - - -def subject_match_clauses(ctx: AuthContext) -> tuple[Any, ...]: - if ctx.subject_type == SubjectType.ACCOUNT: - return (OAuthAccessToken.account_id == str(ctx.account_id),) - return ( - OAuthAccessToken.subject_email == ctx.subject_email, - OAuthAccessToken.subject_issuer == ctx.subject_issuer, - OAuthAccessToken.account_id.is_(None), - ) - - -def list_active_sessions(ctx: AuthContext, now: datetime, *, session: Session) -> list[OAuthAccessToken]: - return list( - session.execute( - select(OAuthAccessToken) - .where( - and_( - *subject_match_clauses(ctx), - OAuthAccessToken.revoked_at.is_(None), - OAuthAccessToken.token_hash.is_not(None), - OAuthAccessToken.expires_at > now, - ) - ) - .order_by(OAuthAccessToken.created_at.desc()) - ) - .scalars() - .all() - ) - - -def token_belongs_to_subject(token_id: str, ctx: AuthContext, *, session: Session) -> bool: - row = session.execute( - select(OAuthAccessToken.id).where( - and_( - OAuthAccessToken.id == token_id, - *subject_match_clauses(ctx), - ) - ) - ).first() - return row is not None - - -def revoke_oauth_token(redis_client: Any, token_id: str, *, session: Session) -> None: - row = ( - session.query(OAuthAccessToken.token_hash) - .filter( - OAuthAccessToken.id == token_id, - OAuthAccessToken.revoked_at.is_(None), - ) - .one_or_none() - ) - pre_revoke_hash = row[0] if row else None - - stmt = ( - update(OAuthAccessToken) - .where( - OAuthAccessToken.id == token_id, - OAuthAccessToken.revoked_at.is_(None), - ) - .values(revoked_at=datetime.now(UTC), token_hash=None) - ) - session.execute(stmt) - session.commit() - - if pre_revoke_hash: - redis_client.delete(TOKEN_CACHE_KEY_FMT.format(hash=pre_revoke_hash)) diff --git a/api/tests/test_containers_integration_tests/controllers/openapi/conftest.py b/api/tests/test_containers_integration_tests/controllers/openapi/conftest.py index d06da5bc018..2047ee765d1 100644 --- a/api/tests/test_containers_integration_tests/controllers/openapi/conftest.py +++ b/api/tests/test_containers_integration_tests/controllers/openapi/conftest.py @@ -1,8 +1,7 @@ from __future__ import annotations import uuid -from collections.abc import Callable, Generator -from contextlib import contextmanager +from collections.abc import Callable from typing import Literal from unittest.mock import patch @@ -12,7 +11,8 @@ from flask import Flask from sqlalchemy.orm import Session from controllers.openapi.auth.data import AuthData -from libs.oauth_bearer import AuthContext, Scope, SubjectType, TokenType, reset_auth_ctx, set_auth_ctx +from libs.oauth_bearer import Scope, TokenType +from machinery.context import AccountRequestContext from models import Account, Tenant from services.account_service import AccountService, TenantService from tests.test_containers_integration_tests.helpers import generate_valid_password @@ -91,34 +91,14 @@ def auth_for( ) -@contextmanager -def account_auth_context( +def request_context_for( account: Account, *, - token_id: uuid.UUID, - client_id: str = "integration-cli", -) -> Generator[AuthContext]: - """Publish an account ``AuthContext`` for handlers that read ``get_auth_ctx()``. - - The auth pipeline normally sets this ContextVar; the integration suite - bypasses the pipeline via ``inspect.unwrap``, so endpoints that resolve the - caller through ``get_auth_ctx()`` (the ``/account/sessions*`` family) need it - set explicitly. Resets on exit so the worker thread can't leak identity. - """ - ctx = AuthContext( - subject_type=SubjectType.ACCOUNT, - subject_email=account.email, - subject_issuer=None, - account_id=uuid.UUID(str(account.id)), - client_id=client_id, - scopes=frozenset({Scope.FULL}), - token_id=token_id, - token_type=TokenType.OAUTH_ACCOUNT, - expires_at=None, - token_hash="integration-test", + token_id: uuid.UUID | None = None, +) -> AccountRequestContext: + return AccountRequestContext( + request_id="integration-request", + trace_id="integration-trace", + account_id=str(account.id), + access_token_id=str(token_id) if token_id is not None else None, ) - reset_token = set_auth_ctx(ctx) - try: - yield ctx - finally: - reset_auth_ctx(reset_token) diff --git a/api/tests/test_containers_integration_tests/controllers/openapi/test_account.py b/api/tests/test_containers_integration_tests/controllers/openapi/test_account.py index 0a04c4fdad0..d2181b8398d 100644 --- a/api/tests/test_containers_integration_tests/controllers/openapi/test_account.py +++ b/api/tests/test_containers_integration_tests/controllers/openapi/test_account.py @@ -9,20 +9,21 @@ from sqlalchemy.orm import Session from controllers.openapi.account import AccountApi from models import Account from models.account import TenantAccountRole -from tests.test_containers_integration_tests.controllers.openapi.conftest import add_tenant_for_account, auth_for +from tests.test_containers_integration_tests.controllers.openapi.conftest import ( + add_tenant_for_account, + request_context_for, +) class TestAccountInfo: - def test_returns_account_and_owner_workspace( - self, app: Flask, db_session_with_containers: Session, make_account: Callable[..., Account] - ) -> None: + def test_returns_account_and_owner_workspace(self, app: Flask, make_account: Callable[..., Account]) -> None: account = make_account() owner_tenant = account.current_tenant assert owner_tenant is not None api = AccountApi() with app.test_request_context("/openapi/v1/account"): - result = unwrap(api.get)(api, db_session_with_containers, auth_data=auth_for(account)) + result = unwrap(api.get)(api, request_context_for(account)) assert result.subject_type == "account" assert result.subject_email == account.email @@ -47,7 +48,7 @@ class TestAccountInfo: api = AccountApi() with app.test_request_context("/openapi/v1/account"): - result = unwrap(api.get)(api, db_session_with_containers, auth_data=auth_for(account)) + result = unwrap(api.get)(api, request_context_for(account)) assert {w.id for w in result.workspaces} == {owner_tenant.id, second.id} roles = {w.id: w.role for w in result.workspaces} diff --git a/api/tests/test_containers_integration_tests/controllers/openapi/test_account_sessions.py b/api/tests/test_containers_integration_tests/controllers/openapi/test_account_sessions.py index 9e6acab5543..d2e339c565e 100644 --- a/api/tests/test_containers_integration_tests/controllers/openapi/test_account_sessions.py +++ b/api/tests/test_containers_integration_tests/controllers/openapi/test_account_sessions.py @@ -18,7 +18,7 @@ from controllers.openapi.account import ( from extensions.ext_redis import redis_client from models import Account from services.oauth_device_flow import PREFIX_OAUTH_ACCOUNT, MintResult, mint_oauth_token -from tests.test_containers_integration_tests.controllers.openapi.conftest import account_auth_context, auth_for +from tests.test_containers_integration_tests.controllers.openapi.conftest import request_context_for def _mint_account_token( @@ -51,13 +51,11 @@ class TestSessionList: api = AccountSessionsApi() with app.test_request_context("/openapi/v1/account/sessions"): - with account_auth_context(account, token_id=mint.token_id): - result = unwrap(api.get)( - api, - db_session_with_containers, - auth_data=auth_for(account, token_id=mint.token_id), - query=SessionListQuery(), - ) + result = unwrap(api.get)( + api, + request_context_for(account, token_id=mint.token_id), + query=SessionListQuery(), + ) assert result.total == 1 row = result.data[0] @@ -76,13 +74,11 @@ class TestSessionList: api = AccountSessionsApi() with app.test_request_context("/openapi/v1/account/sessions"): - with account_auth_context(account, token_id=mine.token_id): - result = unwrap(api.get)( - api, - db_session_with_containers, - auth_data=auth_for(account, token_id=mine.token_id), - query=SessionListQuery(), - ) + result = unwrap(api.get)( + api, + request_context_for(account, token_id=mine.token_id), + query=SessionListQuery(), + ) assert {row.id for row in result.data} == {str(mine.token_id)} @@ -96,23 +92,18 @@ class TestSessionRevoke: revoke_api = AccountSessionsSelfApi() with app.test_request_context("/openapi/v1/account/sessions/self", method="DELETE"): - with account_auth_context(account, token_id=mint.token_id): - result = unwrap(revoke_api.delete)( - revoke_api, db_session_with_containers, auth_data=auth_for(account, token_id=mint.token_id) - ) + result = unwrap(revoke_api.delete)(revoke_api, request_context_for(account, token_id=mint.token_id)) assert result.status == "revoked" # Revocation persisted: the real list path no longer returns it. list_api = AccountSessionsApi() with app.test_request_context("/openapi/v1/account/sessions"): - with account_auth_context(account, token_id=mint.token_id): - listing = unwrap(list_api.get)( - list_api, - db_session_with_containers, - auth_data=auth_for(account, token_id=mint.token_id), - query=SessionListQuery(), - ) + listing = unwrap(list_api.get)( + list_api, + request_context_for(account, token_id=mint.token_id), + query=SessionListQuery(), + ) assert listing.total == 0 def test_revoke_by_id_for_own_session( @@ -124,13 +115,11 @@ class TestSessionRevoke: api = AccountSessionByIdApi() with app.test_request_context(f"/openapi/v1/account/sessions/{session_id}", method="DELETE"): - with account_auth_context(account, token_id=mint.token_id): - result = unwrap(api.delete)( - api, - db_session_with_containers, - session_id=session_id, - auth_data=auth_for(account, token_id=mint.token_id), - ) + result = unwrap(api.delete)( + api, + request_context_for(account, token_id=mint.token_id), + session_id=session_id, + ) assert result.status == "revoked" @@ -146,11 +135,9 @@ class TestSessionRevoke: api = AccountSessionByIdApi() session_id = str(foreign.token_id) with app.test_request_context(f"/openapi/v1/account/sessions/{session_id}", method="DELETE"): - with account_auth_context(outsider, token_id=uuid4()): - with pytest.raises(NotFound): - unwrap(api.delete)( - api, - db_session_with_containers, - session_id=session_id, - auth_data=auth_for(outsider, token_id=uuid4()), - ) + with pytest.raises(NotFound): + unwrap(api.delete)( + api, + request_context_for(outsider, token_id=uuid4()), + session_id=session_id, + ) diff --git a/api/tests/unit_tests/controllers/console/workspace/test_workspace.py b/api/tests/unit_tests/controllers/console/workspace/test_workspace.py index dadfa2a23a5..2bf35fc6917 100644 --- a/api/tests/unit_tests/controllers/console/workspace/test_workspace.py +++ b/api/tests/unit_tests/controllers/console/workspace/test_workspace.py @@ -201,6 +201,7 @@ class TestWorkspaceQueryRepository: TenantAccountJoin( tenant_id=earlier.id, account_id="account-1", + current=True, last_opened_at=last_opened_at, ), TenantAccountJoin(tenant_id=later.id, account_id="account-1"), @@ -214,6 +215,7 @@ class TestWorkspaceQueryRepository: repository = WorkspaceQueryRepository(workspace_session.session_factory) result = repository.list_for_account("account-1") membership_ids = repository.list_ids_for_account("account-1") + access_workspaces = repository.list_account_access_workspaces("account-1") assert repository.has_active_for_account("account-1") is True assert repository.has_active_for_account("missing-account") is False @@ -234,6 +236,10 @@ class TestWorkspaceQueryRepository: ), ) assert set(membership_ids) == {earlier.id, later.id, archived.id} + access_by_id = {workspace.id: workspace for workspace in access_workspaces} + assert set(access_by_id) == {earlier.id, later.id, archived.id} + assert access_by_id[earlier.id].current is True + assert access_by_id[earlier.id].role == "normal" assert repository.has_active_membership("account-1") is True assert repository.has_active_membership("account-3") is False diff --git a/api/tests/unit_tests/controllers/openapi/conftest.py b/api/tests/unit_tests/controllers/openapi/conftest.py index 3a900011f54..9f60b5979a4 100644 --- a/api/tests/unit_tests/controllers/openapi/conftest.py +++ b/api/tests/unit_tests/controllers/openapi/conftest.py @@ -18,6 +18,7 @@ def _stub_execute( scope=None, allowed_token_types=None, edition=None, + require_valid_enterprise_license=False, workspace_membership=False, allowed_roles=None, rbac=None, diff --git a/api/tests/unit_tests/controllers/openapi/test_account.py b/api/tests/unit_tests/controllers/openapi/test_account.py index d4b28fa554e..a623d29ca69 100644 --- a/api/tests/unit_tests/controllers/openapi/test_account.py +++ b/api/tests/unit_tests/controllers/openapi/test_account.py @@ -2,13 +2,12 @@ import builtins import sys -import uuid from types import SimpleNamespace import pytest from flask import Flask from flask.views import MethodView -from werkzeug.exceptions import UnprocessableEntity +from werkzeug.exceptions import NotFound, UnprocessableEntity from controllers.openapi import bp as openapi_bp from controllers.openapi.account import ( @@ -17,8 +16,8 @@ from controllers.openapi.account import ( AccountSessionsApi, AccountSessionsSelfApi, ) -from controllers.openapi.auth.data import AuthData -from libs.oauth_bearer import Scope, TokenType +from machinery.context import AccountRequestContext +from services.entities.account_access_entities import AccountSessionPage if not hasattr(builtins, "MethodView"): builtins.MethodView = MethodView # type: ignore[attr-defined] @@ -88,96 +87,47 @@ def test_session_by_id_dispatches_to_correct_class(openapi_app: Flask): assert "DELETE" in rule.methods -def test_subject_match_for_account_filters_by_account_id(): - """Account subject scopes queries via account_id.""" - import uuid as _uuid - - from libs.oauth_bearer import AuthContext, SubjectType, TokenType - from services.oauth_device_flow import subject_match_clauses - - aid = _uuid.uuid4() - ctx = AuthContext( - subject_type=SubjectType.ACCOUNT, - subject_email="user@example.com", - subject_issuer="dify:account", - account_id=aid, - client_id="difyctl", - scopes=frozenset({"full"}), - token_id=_uuid.uuid4(), - token_type=TokenType.OAUTH_ACCOUNT, - expires_at=None, - token_hash="h1", - verified_tenants={}, - ) - clauses = subject_match_clauses(ctx) - # One predicate, on account_id - assert len(clauses) == 1 - assert "account_id" in str(clauses[0]) - - -def test_subject_match_for_external_sso_filters_by_email_and_issuer(): - """External SSO subject scopes via (subject_email, subject_issuer) - AND account_id IS NULL — so a same-email account row from a - federated tenant cannot be revoked through an SSO bearer. - """ - import uuid as _uuid - - from libs.oauth_bearer import AuthContext, SubjectType, TokenType - from services.oauth_device_flow import subject_match_clauses - - ctx = 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={}, - ) - clauses = subject_match_clauses(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 +def test_session_by_id_rejects_malformed_uuid(app: Flask) -> None: + api = AccountSessionByIdApi() + with app.test_request_context("/openapi/v1/account/sessions/not-a-uuid", method="DELETE"): + with pytest.raises(NotFound, match="session not found"): + api.delete.__wrapped__(api, _request_context(), session_id="not-a-uuid") # --- GET /account/sessions query validation (the handler routes ?page/?limit through -# SessionListQuery so the server enforces the bounds the contract advertises). The auth ctx and -# DB read are stubbed so these exercise only the validation + paging path; __wrapped__ skips the -# auth guard, which is covered separately in auth/. --- +# SessionListQuery so the server enforces the bounds the contract advertises). The application +# service is replaced with a small fake so these exercise only parsing and serialization; +# __wrapped__ skips the complete Admission boundary. --- _ACCOUNT_MOD = "controllers.openapi.account" -def _session_auth_data() -> AuthData: - return AuthData( - token_type=TokenType.OAUTH_ACCOUNT, - account_id=uuid.uuid4(), - token_hash="test", - token_id=uuid.uuid4(), - scopes=frozenset({Scope.FULL}), - required_scope=Scope.FULL, - allowed_roles=None, +def _request_context() -> AccountRequestContext: + return AccountRequestContext( + request_id="request-1", + trace_id="trace-1", + account_id="account-1", + access_token_id="token-1", ) -def _stub_session_deps(monkeypatch: pytest.MonkeyPatch, rows): +class _SessionListService: + def list_sessions(self, _context: AccountRequestContext, *, page: int, limit: int) -> AccountSessionPage: + return AccountSessionPage(page=page, limit=limit, total=0, items=()) + + +def _stub_account_service(monkeypatch: pytest.MonkeyPatch) -> None: mod = sys.modules[_ACCOUNT_MOD] - monkeypatch.setattr(mod, "get_auth_ctx", lambda: SimpleNamespace()) - monkeypatch.setattr(mod, "list_active_sessions", lambda *args, **kwargs: rows) + services = SimpleNamespace(accounts=SimpleNamespace(access=_SessionListService())) + monkeypatch.setattr(mod, "application_services", lambda: services) def test_sessions_list_valid_query_parses_page_and_limit(app: Flask, monkeypatch: pytest.MonkeyPatch): """A valid ?page&limit round-trips through SessionListQuery into the response envelope.""" api = AccountSessionsApi() - _stub_session_deps(monkeypatch, []) + _stub_account_service(monkeypatch) with app.test_request_context("/openapi/v1/account/sessions?page=2&limit=5"): - body, status = api.get.__wrapped__(api, auth_data=_session_auth_data()) + body, status = api.get.__wrapped__(api, _request_context()) assert status == 200 assert body["page"] == 2 assert body["limit"] == 5 @@ -188,9 +138,9 @@ def test_sessions_list_valid_query_parses_page_and_limit(app: Flask, monkeypatch def test_sessions_list_defaults_when_query_omitted(app: Flask, monkeypatch: pytest.MonkeyPatch): """No query → the model's defaults (page=1, limit=100) drive the envelope.""" api = AccountSessionsApi() - _stub_session_deps(monkeypatch, []) + _stub_account_service(monkeypatch) with app.test_request_context("/openapi/v1/account/sessions"): - body, status = api.get.__wrapped__(api, auth_data=_session_auth_data()) + body, status = api.get.__wrapped__(api, _request_context()) assert status == 200 assert body["page"] == 1 assert body["limit"] == 100 @@ -210,7 +160,7 @@ def test_sessions_list_defaults_when_query_omitted(app: Flask, monkeypatch: pyte def test_sessions_list_rejects_out_of_bounds_query(app: Flask, monkeypatch: pytest.MonkeyPatch, query): """Out-of-range / unknown query params raise 422 instead of being silently coerced.""" api = AccountSessionsApi() - _stub_session_deps(monkeypatch, []) + _stub_account_service(monkeypatch) with app.test_request_context(f"/openapi/v1/account/sessions?{query}"): with pytest.raises(UnprocessableEntity): - api.get.__wrapped__(api, auth_data=_session_auth_data()) + api.get.__wrapped__(api, _request_context()) diff --git a/api/tests/unit_tests/controllers/openapi/test_flask_admission.py b/api/tests/unit_tests/controllers/openapi/test_flask_admission.py new file mode 100644 index 00000000000..fb18eea75ae --- /dev/null +++ b/api/tests/unit_tests/controllers/openapi/test_flask_admission.py @@ -0,0 +1,116 @@ +from __future__ import annotations + +from collections.abc import Callable +from functools import wraps +from uuid import UUID + +import pytest +from flask import Flask +from werkzeug.exceptions import Unauthorized + +from controllers.openapi import flask_admission +from controllers.openapi.auth.data import AuthData +from enums import DeploymentEdition +from libs.oauth_bearer import Scope, TokenType +from libs.rate_limit import LIMIT_ME_PER_ACCOUNT +from machinery.context import AccountRequestContext +from models.account import Account, AccountStatus + + +def _auth_data(*, status: AccountStatus = AccountStatus.ACTIVE) -> AuthData: + account = Account(name="Ada", email="ada@example.com", status=status) + account.id = "11111111-1111-1111-1111-111111111111" + return AuthData( + token_type=TokenType.OAUTH_ACCOUNT, + account_id=UUID(account.id), + token_hash="hash-1", + token_id=UUID("22222222-2222-2222-2222-222222222222"), + scopes=frozenset({Scope.FULL}), + caller=account, + ) + + +def _install_fake_transport( + monkeypatch: pytest.MonkeyPatch, + auth_data: AuthData, + captured: dict[str, object], +) -> None: + def guard(**requirements: object) -> Callable[[Callable[..., object]], Callable[..., object]]: + captured.update(requirements) + + def decorator(view: Callable[..., object]) -> Callable[..., object]: + @wraps(view) + def admitted(*args: object, **kwargs: object) -> object: + return view(*args, auth_data=auth_data, **kwargs) + + return admitted + + return decorator + + monkeypatch.setattr(flask_admission.auth_router, "guard", guard) + + +def test_admission_builds_stable_request_context( + monkeypatch: pytest.MonkeyPatch, +) -> None: + captured: dict[str, object] = {} + limited: list[tuple[object, str]] = [] + _install_fake_transport(monkeypatch, _auth_data(), captured) + monkeypatch.setattr(flask_admission, "get_request_id", lambda: "request-1") + monkeypatch.setattr(flask_admission, "get_trace_id", lambda: "trace-1") + monkeypatch.setattr(flask_admission, "enforce", lambda spec, *, key: limited.append((spec, key))) + + @flask_admission.openapi_account_admission( + scope=Scope.FULL, + editions=frozenset({DeploymentEdition.ENTERPRISE}), + rate_limit=LIMIT_ME_PER_ACCOUNT, + ) + def view(_self: object, context: AccountRequestContext) -> AccountRequestContext: + return context + + with Flask(__name__).test_request_context("/openapi/v1/account"): + context = view(object()) + + assert context == AccountRequestContext( + request_id="request-1", + trace_id="trace-1", + account_id="11111111-1111-1111-1111-111111111111", + access_token_id="22222222-2222-2222-2222-222222222222", + ) + assert captured == { + "scope": Scope.FULL, + "allowed_token_types": frozenset({TokenType.OAUTH_ACCOUNT}), + "edition": frozenset({DeploymentEdition.ENTERPRISE}), + "require_valid_enterprise_license": True, + } + assert limited == [(LIMIT_ME_PER_ACCOUNT, "account:11111111-1111-1111-1111-111111111111")] + + +def test_admission_rejects_uninitialized_account(monkeypatch: pytest.MonkeyPatch) -> None: + _install_fake_transport(monkeypatch, _auth_data(status=AccountStatus.UNINITIALIZED), {}) + + @flask_admission.openapi_account_admission(scope=Scope.FULL) + def view(_self: object, _context: AccountRequestContext) -> None: + raise AssertionError("view must not run") + + with Flask(__name__).test_request_context("/openapi/v1/account"): + with pytest.raises(Unauthorized, match="account not initialized"): + view(object()) + + +def test_admission_rejects_missing_auth_data(monkeypatch: pytest.MonkeyPatch) -> None: + def guard(**_requirements: object) -> Callable[[Callable[..., object]], Callable[..., object]]: + def decorator(view: Callable[..., object]) -> Callable[..., object]: + return view + + return decorator + + monkeypatch.setattr(flask_admission.auth_router, "guard", guard) + + @flask_admission.openapi_account_admission(scope=Scope.FULL) + def view(_self: object, _context: AccountRequestContext) -> None: + raise AssertionError("view must not run") + + with Flask(__name__).test_request_context("/openapi/v1/account"): + with pytest.raises(RuntimeError, match="did not provide valid AuthData"): + view(object()) diff --git a/api/tests/unit_tests/libs/test_datetime_utils.py b/api/tests/unit_tests/libs/test_datetime_utils.py index e2dce54d1e0..a4229d20777 100644 --- a/api/tests/unit_tests/libs/test_datetime_utils.py +++ b/api/tests/unit_tests/libs/test_datetime_utils.py @@ -4,7 +4,19 @@ from unittest.mock import patch import pytest import pytz -from libs.datetime_utils import naive_utc_now, parse_time_range, to_utc_timestamp +from libs.datetime_utils import naive_utc_now, parse_time_range, to_utc_timestamp, utc_now + + +def test_utc_now(monkeypatch: pytest.MonkeyPatch): + expected = datetime.datetime(2026, 8, 26, 12, tzinfo=datetime.UTC) + + def _now_func(tz: datetime.timezone | None) -> datetime.datetime: + return expected.astimezone(tz) + + monkeypatch.setattr("libs.datetime_utils._now_func", _now_func) + + assert utc_now() == expected + assert utc_now().tzinfo is datetime.UTC def test_naive_utc_now(monkeypatch: pytest.MonkeyPatch): diff --git a/api/tests/unit_tests/repositories/test_oauth_access_token_repository.py b/api/tests/unit_tests/repositories/test_oauth_access_token_repository.py new file mode 100644 index 00000000000..3fd15186d84 --- /dev/null +++ b/api/tests/unit_tests/repositories/test_oauth_access_token_repository.py @@ -0,0 +1,126 @@ +from __future__ import annotations + +from datetime import UTC, datetime, timedelta + +import pytest +from sqlalchemy.orm import Session, sessionmaker + +from models.oauth import OAuthAccessToken +from repositories.oauth_access_token_repository import SQLAlchemyOAuthAccessTokenRepository + +ACCOUNT_ID = "11111111-1111-1111-1111-111111111111" +OTHER_ACCOUNT_ID = "22222222-2222-2222-2222-222222222222" +TOKEN_ID = "33333333-3333-3333-3333-333333333333" +OTHER_TOKEN_ID = "44444444-4444-4444-4444-444444444444" +NOW = datetime(2026, 8, 25, 12, tzinfo=UTC) + + +def _token( + *, + token_id: str = TOKEN_ID, + account_id: str | None = ACCOUNT_ID, + 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="user@example.com", + subject_issuer="dify:account" if account_id is not None else "https://idp.example.com", + account_id=account_id, + 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 NOW + timedelta(days=1), + revoked_at=revoked_at, + ) + token.id = token_id + if created_at is not None: + token.created_at = created_at + return token + + +@pytest.mark.parametrize("sqlite_session", [(OAuthAccessToken,)], indirect=True) +def test_list_active_is_account_scoped_and_database_paginated( + sqlite_session: Session, + sqlite_session_factory: sessionmaker[Session], +) -> None: + sqlite_session.add_all( + [ + _token(token_id=TOKEN_ID, created_at=NOW - timedelta(minutes=1)), + _token(token_id=OTHER_TOKEN_ID, created_at=NOW - timedelta(minutes=2)), + _token( + token_id="55555555-5555-5555-5555-555555555555", + expires_at=NOW - timedelta(seconds=1), + ), + _token( + token_id="66666666-6666-6666-6666-666666666666", + token_hash=None, + revoked_at=NOW - timedelta(seconds=1), + ), + _token(token_id="77777777-7777-7777-7777-777777777777", account_id=OTHER_ACCOUNT_ID), + _token(token_id="88888888-8888-8888-8888-888888888888", account_id=None), + ] + ) + sqlite_session.commit() + repository = SQLAlchemyOAuthAccessTokenRepository(session_factory=sqlite_session_factory) + + total, rows = repository.list_active(account_id=ACCOUNT_ID, active_at=NOW, offset=1, limit=1) + + assert total == 2 + assert [row.id for row in rows] == [OTHER_TOKEN_ID] + + +@pytest.mark.parametrize("sqlite_session", [(OAuthAccessToken,)], indirect=True) +def test_revoke_returns_hash_and_persists_revocation( + sqlite_session: Session, + sqlite_session_factory: sessionmaker[Session], +) -> None: + sqlite_session.add(_token()) + sqlite_session.commit() + repository = SQLAlchemyOAuthAccessTokenRepository(session_factory=sqlite_session_factory) + + result = repository.revoke(account_id=ACCOUNT_ID, token_id=TOKEN_ID, revoked_at=NOW) + + assert result.owned is True + assert result.token_hash == "live-hash" + sqlite_session.expire_all() + persisted = sqlite_session.get(OAuthAccessToken, TOKEN_ID) + assert persisted is not None + assert persisted.token_hash is None + assert persisted.revoked_at is not None + + +@pytest.mark.parametrize("sqlite_session", [(OAuthAccessToken,)], indirect=True) +def test_revoke_is_idempotent_for_an_owned_session( + sqlite_session: Session, + sqlite_session_factory: sessionmaker[Session], +) -> None: + sqlite_session.add(_token(token_hash=None, revoked_at=NOW - timedelta(minutes=1))) + sqlite_session.commit() + repository = SQLAlchemyOAuthAccessTokenRepository(session_factory=sqlite_session_factory) + + result = repository.revoke(account_id=ACCOUNT_ID, token_id=TOKEN_ID, revoked_at=NOW) + + assert result.owned is True + assert result.token_hash is None + + +@pytest.mark.parametrize("sqlite_session", [(OAuthAccessToken,)], indirect=True) +def test_revoke_does_not_disclose_another_accounts_session( + sqlite_session: Session, + sqlite_session_factory: sessionmaker[Session], +) -> None: + sqlite_session.add(_token(account_id=OTHER_ACCOUNT_ID)) + sqlite_session.commit() + repository = SQLAlchemyOAuthAccessTokenRepository(session_factory=sqlite_session_factory) + + result = repository.revoke(account_id=ACCOUNT_ID, token_id=TOKEN_ID, revoked_at=NOW) + + assert result.owned is False + assert result.token_hash is None + sqlite_session.expire_all() + persisted = sqlite_session.get(OAuthAccessToken, TOKEN_ID) + assert persisted is not None + assert persisted.token_hash == "live-hash" diff --git a/api/tests/unit_tests/services/test_account_access_service.py b/api/tests/unit_tests/services/test_account_access_service.py new file mode 100644 index 00000000000..7121d8cc1da --- /dev/null +++ b/api/tests/unit_tests/services/test_account_access_service.py @@ -0,0 +1,173 @@ +from __future__ import annotations + +from dataclasses import dataclass, field +from datetime import UTC, datetime + +import pytest + +from machinery.context import AccountRequestContext +from services.account_access_service import AccountAccessService +from services.account_errors import AccountNotFoundError, AccountSessionNotFoundError +from services.entities.account_access_entities import ( + AccountSessionRevocation, + AccountSessionSnapshot, + AccountWorkspaceSnapshot, +) +from services.entities.account_entities import AccountSnapshot + +NOW = datetime(2026, 8, 25, 12, tzinfo=UTC) + + +def _context(*, token_id: str | None = "token-1") -> AccountRequestContext: + return AccountRequestContext("request-1", "trace-1", "account-1", token_id) + + +def _account() -> AccountSnapshot: + return AccountSnapshot( + id="account-1", + name="Ada", + email="ada@example.com", + avatar=None, + is_password_set=False, + interface_language=None, + interface_theme=None, + timezone=None, + last_login_at=None, + last_login_ip=None, + status="active", + initialized_at=None, + created_at=NOW, + ) + + +@dataclass +class _Accounts: + account: AccountSnapshot | None = field(default_factory=_account) + + def get(self, account_id: str) -> AccountSnapshot | None: + assert account_id == "account-1" + return self.account + + +@dataclass +class _Workspaces: + items: tuple[AccountWorkspaceSnapshot, ...] = () + + def list_account_access_workspaces(self, account_id: str) -> tuple[AccountWorkspaceSnapshot, ...]: + assert account_id == "account-1" + return self.items + + +@dataclass +class _Sessions: + items: tuple[AccountSessionSnapshot, ...] = () + total: int = 0 + revocation: AccountSessionRevocation = AccountSessionRevocation(owned=True) + list_call: tuple[str, datetime, int, int] | None = None + revoke_call: tuple[str, str, datetime] | None = None + + def list_active( + self, + *, + account_id: str, + active_at: datetime, + offset: int, + limit: int, + ) -> tuple[int, tuple[AccountSessionSnapshot, ...]]: + self.list_call = (account_id, active_at, offset, limit) + return self.total, self.items + + def revoke(self, *, account_id: str, token_id: str, revoked_at: datetime) -> AccountSessionRevocation: + self.revoke_call = (account_id, token_id, revoked_at) + return self.revocation + + +@dataclass +class _TokenCache: + invalidated: list[str] = field(default_factory=list) + + def __call__(self, token_hash: str) -> None: + self.invalidated.append(token_hash) + + +def _service( + *, + accounts: _Accounts | None = None, + workspaces: _Workspaces | None = None, + sessions: _Sessions | None = None, + token_cache: _TokenCache | None = None, +) -> AccountAccessService: + return AccountAccessService( + accounts=accounts or _Accounts(), + workspaces=workspaces or _Workspaces(), + sessions=sessions or _Sessions(), + invalidate_token_cache=token_cache or _TokenCache(), + now=lambda: NOW, + ) + + +def test_get_prefers_current_workspace_as_default() -> None: + workspaces = _Workspaces( + items=( + AccountWorkspaceSnapshot("workspace-1", "First", "normal", False), + AccountWorkspaceSnapshot("workspace-2", "Current", "owner", True), + ) + ) + + snapshot = _service(workspaces=workspaces).get(_context()) + + assert snapshot.account.email == "ada@example.com" + assert snapshot.workspaces == workspaces.items + assert snapshot.default_workspace_id == "workspace-2" + + +def test_get_falls_back_to_first_workspace() -> None: + workspaces = _Workspaces( + items=( + AccountWorkspaceSnapshot("workspace-1", "First", "normal", False), + AccountWorkspaceSnapshot("workspace-2", "Second", "owner", False), + ) + ) + + assert _service(workspaces=workspaces).get(_context()).default_workspace_id == "workspace-1" + + +def test_get_raises_when_admitted_account_disappeared() -> None: + with pytest.raises(AccountNotFoundError): + _service(accounts=_Accounts(account=None)).get(_context()) + + +def test_list_sessions_delegates_database_pagination() -> None: + sessions = _Sessions(total=12) + + page = _service(sessions=sessions).list_sessions(_context(), page=3, limit=5) + + assert sessions.list_call == ("account-1", NOW, 10, 5) + assert page.page == 3 + assert page.total == 12 + assert page.has_more is False + + +def test_revoke_current_session_invalidates_live_token_cache() -> None: + sessions = _Sessions(revocation=AccountSessionRevocation(owned=True, token_hash="hash-1")) + cache = _TokenCache() + + _service(sessions=sessions, token_cache=cache).revoke_current_session(_context()) + + assert sessions.revoke_call == ("account-1", "token-1", NOW) + assert cache.invalidated == ["hash-1"] + + +def test_revoke_foreign_session_does_not_invalidate_cache() -> None: + sessions = _Sessions(revocation=AccountSessionRevocation(owned=False)) + cache = _TokenCache() + + with pytest.raises(AccountSessionNotFoundError): + _service(sessions=sessions, token_cache=cache).revoke_session(_context(), token_id="foreign") + + assert cache.invalidated == [] + + +def test_revoke_current_requires_admitted_token_id() -> None: + with pytest.raises(RuntimeError, match="did not resolve an access token"): + _service().revoke_current_session(_context(token_id=None)) diff --git a/api/tests/unit_tests/services/test_account_service.py b/api/tests/unit_tests/services/test_account_service.py index 64ffe961cd6..75221d81f73 100644 --- a/api/tests/unit_tests/services/test_account_service.py +++ b/api/tests/unit_tests/services/test_account_service.py @@ -2901,22 +2901,6 @@ class TestSessionInjectedGetters: def test_account_belongs_to_tenant_false_when_no_join(self, sqlite_session: Session) -> None: assert TenantService.account_belongs_to_tenant("user-1", "tenant-1", session=sqlite_session) is False - def test_get_account_memberships_returns_join_tenant_pairs(self, sqlite_session: Session) -> None: - """Returns every ``(TenantAccountJoin, Tenant)`` pair for an account.""" - tenant = Tenant(name="Joined Workspace") - other_tenant = Tenant(name="Other Workspace") - sqlite_session.add_all([tenant, other_tenant]) - sqlite_session.flush() - join = self._add_tenant_account_join(sqlite_session, tenant, "user-123", TenantAccountRole.NORMAL, current=True) - self._add_tenant_account_join(sqlite_session, other_tenant, "other-user", TenantAccountRole.NORMAL) - sqlite_session.commit() - - out = TenantService.get_account_memberships("user-123", session=sqlite_session) - - assert len(out) == 1 - assert out[0][0] is join - assert out[0][1] is tenant - def test_get_workspaces_for_account_uses_session_execute(self, sqlite_session: Session) -> None: """The list endpoint orders by ``Tenant.created_at``; the helper returns ``(Tenant, TenantAccountJoin)`` rows in that order. diff --git a/api/tests/unit_tests/services/test_oauth_device_flow.py b/api/tests/unit_tests/services/test_oauth_device_flow.py deleted file mode 100644 index 35c5a691672..00000000000 --- a/api/tests/unit_tests/services/test_oauth_device_flow.py +++ /dev/null @@ -1,219 +0,0 @@ -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