From 16d3302bff26232bc4608038cd01b7e7bff82141 Mon Sep 17 00:00:00 2001 From: "Byron.wang" Date: Wed, 23 Sep 2026 06:36:48 +0000 Subject: [PATCH] refactor(api): isolate OAuth device flow orchestration (#41361) --- api/.importlinter | 21 +- api/constants/oauth_bearer.py | 65 ++ api/controllers/console/flask_admission.py | 8 +- api/controllers/openapi/_models.py | 29 +- api/controllers/openapi/account.py | 2 +- api/controllers/openapi/app_dsl.py | 2 +- api/controllers/openapi/app_run.py | 2 +- api/controllers/openapi/apps.py | 2 +- .../openapi/apps_permitted_external.py | 2 +- api/controllers/openapi/auth/requirements.py | 2 +- api/controllers/openapi/auth/subjects.py | 5 +- api/controllers/openapi/files.py | 2 +- api/controllers/openapi/flask_admission.py | 32 + api/controllers/openapi/human_input_form.py | 2 +- api/controllers/openapi/oauth_device.py | 398 ++----- api/controllers/openapi/oauth_device_sso.py | 448 ++------ api/controllers/openapi/workflow_events.py | 2 +- api/controllers/openapi/workspaces.py | 2 +- api/extensions/ext_application_services.py | 47 + api/extensions/ext_redis.py | 29 +- api/libs/device_flow_security.py | 64 +- api/libs/oauth_bearer.py | 60 +- api/repositories/account_repository.py | 18 +- .../oauth_device_token_repository.py | 193 ++++ .../workspace_query_repository.py | 15 + api/services/account_ports.py | 2 + api/services/account_service.py | 20 +- api/services/feature_query_service.py | 7 + api/services/oauth_device_adapters.py | 336 ++++++ .../oauth_device_application_service.py | 760 +++++++++++++ api/services/oauth_device_contracts.py | 276 +++++ api/services/oauth_device_flow.py | 460 +++----- .../controllers/openapi/conftest.py | 2 +- .../controllers/openapi/conftest.py | 3 +- .../openapi/test_account_sessions.py | 29 +- .../services/test_oauth_device_flow.py | 72 ++ .../controllers/openapi/auth/_world.py | 3 +- .../controllers/openapi/auth/test_router.py | 9 +- .../controllers/openapi/auth/test_subjects.py | 2 +- .../controllers/openapi/conftest.py | 3 +- .../controllers/openapi/test_auth_matrix.py | 9 +- .../openapi/test_device_approve_deny.py | 10 +- .../controllers/openapi/test_device_sso.py | 73 +- .../openapi/test_error_contract.py | 15 +- .../openapi/test_oauth_device_http.py | 383 +++++++ .../openapi/test_oauth_sso_claims.py | 73 -- .../openapi/test_oauth_sso_csrf.py | 22 - .../openapi/test_oauth_sso_host_header.py | 76 -- .../openapi/test_workspaces_members.py | 3 +- api/tests/unit_tests/extensions/test_redis.py | 110 +- .../libs/test_device_flow_security.py | 25 + .../unit_tests/libs/test_oauth_bearer.py | 3 +- .../test_oauth_bearer_rate_limit_ordering.py | 3 +- .../repositories/test_account_repository.py | 15 + .../test_oauth_device_token_repository.py | 236 ++++ .../services/test_account_oauth_service.py | 3 + .../services/test_account_service.py | 5 +- .../services/test_feature_query_service.py | 23 + .../services/test_oauth_device_adapters.py | 189 +++ .../test_oauth_device_application_service.py | 1010 +++++++++++++++++ .../services/test_oauth_device_flow.py | 191 ++++ .../services/test_oauth_device_flow_mint.py | 98 +- 62 files changed, 4616 insertions(+), 1395 deletions(-) create mode 100644 api/constants/oauth_bearer.py create mode 100644 api/controllers/openapi/flask_admission.py create mode 100644 api/repositories/oauth_device_token_repository.py create mode 100644 api/services/oauth_device_adapters.py create mode 100644 api/services/oauth_device_application_service.py create mode 100644 api/services/oauth_device_contracts.py create mode 100644 api/tests/test_containers_integration_tests/services/test_oauth_device_flow.py create mode 100644 api/tests/unit_tests/controllers/openapi/test_oauth_device_http.py delete mode 100644 api/tests/unit_tests/controllers/openapi/test_oauth_sso_claims.py delete mode 100644 api/tests/unit_tests/controllers/openapi/test_oauth_sso_csrf.py delete mode 100644 api/tests/unit_tests/controllers/openapi/test_oauth_sso_host_header.py create mode 100644 api/tests/unit_tests/libs/test_device_flow_security.py create mode 100644 api/tests/unit_tests/repositories/test_oauth_device_token_repository.py create mode 100644 api/tests/unit_tests/services/test_oauth_device_adapters.py create mode 100644 api/tests/unit_tests/services/test_oauth_device_application_service.py create mode 100644 api/tests/unit_tests/services/test_oauth_device_flow.py diff --git a/api/.importlinter b/api/.importlinter index 9a4eaaf0994..c0dde58f038 100644 --- a/api/.importlinter +++ b/api/.importlinter @@ -97,9 +97,6 @@ ignore_imports = core.workflow.nodes.agent_v2.session_store -> services.agent.workspace_service core.workflow.nodes.agent_v2.validators -> services.agent.knowledge_datasets core.workflow.nodes.agent_v2.workspace_retirement_layer -> tasks.collect_agent_resources_task - libs.device_flow_security -> controllers.openapi._models - libs.device_flow_security -> services.entities.feature_entities - libs.device_flow_security -> services.system_feature_service libs.email_i18n -> services.entities.feature_entities libs.email_i18n -> services.system_feature_service libs.external_api -> core @@ -435,6 +432,24 @@ forbidden_modules = sqlalchemy werkzeug +[importlinter:contract:oauth-device-application-service-boundary] +name = OAuth device application service and contracts are framework and persistence neutral +type = forbidden +source_modules = + services.oauth_device_application_service + services.oauth_device_contracts +forbidden_modules = + configs + controllers + extensions + flask + httpx + libs + models + repositories + sqlalchemy + werkzeug + [importlinter:contract:workspace-member-query-service-boundary] name = Workspace member query application service is framework and persistence neutral type = forbidden diff --git a/api/constants/oauth_bearer.py b/api/constants/oauth_bearer.py new file mode 100644 index 00000000000..48c7564ef94 --- /dev/null +++ b/api/constants/oauth_bearer.py @@ -0,0 +1,65 @@ +"""Stable OAuth bearer vocabulary shared across application layers.""" + +from __future__ import annotations + +from enum import StrEnum + + +class Scope(StrEnum): + """Catalog of bearer scopes recognised by the openapi surface. + + `FULL` is the catch-all carried by `dfoa_` account tokens — it satisfies + any per-route scope requirement. `dfoe_` tokens carry the per-feature + scopes (`APPS_RUN`, `APPS_READ_PERMITTED_EXTERNAL`). + """ + + FULL = "full" + APPS_READ = "apps:read" + APPS_READ_PERMITTED_EXTERNAL = "apps:read:permitted-external" + APPS_RUN = "apps:run" + WORKSPACE_READ = "workspace:read" + WORKSPACE_WRITE = "workspace:write" + + +class SubjectType(StrEnum): + # Annotation-only names are not members; they declare what `__new__` attaches. + scopes: frozenset[Scope] + # Whether the subject's token rows carry an `account_id`. A subject that does + # not is keyed by its external issuer instead. + bound_to_account: bool + + ACCOUNT = ("account", frozenset({Scope.FULL}), True) + EXTERNAL_SSO = ("external_sso", frozenset({Scope.APPS_RUN, Scope.APPS_READ_PERMITTED_EXTERNAL}), False) + + def __new__(cls, value: str, scopes: frozenset[Scope], bound_to_account: bool) -> SubjectType: + obj = str.__new__(cls, value) + obj._value_ = value + obj.scopes = scopes + obj.bound_to_account = bound_to_account + return obj + + +class TokenType(StrEnum): + """Many token types may serve one subject; the prefix is the only thing a + bearer string reveals, so it lives here and nowhere else. + """ + + prefix: str + subject: SubjectType + + OAUTH_ACCOUNT = ("oauth_account", "dfoa_", SubjectType.ACCOUNT) + OAUTH_EXTERNAL_SSO = ("oauth_external_sso", "dfoe_", SubjectType.EXTERNAL_SSO) + + def __new__(cls, value: str, prefix: str, subject: SubjectType) -> TokenType: + obj = str.__new__(cls, value) + obj._value_ = value + obj.prefix = prefix + obj.subject = subject + return obj + + @classmethod + def for_token(cls, token: str) -> TokenType | None: + return next((token_type for token_type in cls if token.startswith(token_type.prefix)), None) + + +TOKEN_CACHE_KEY_FMT = "auth:token:{hash}" diff --git a/api/controllers/console/flask_admission.py b/api/controllers/console/flask_admission.py index d9017bf1736..bcf99578a5b 100644 --- a/api/controllers/console/flask_admission.py +++ b/api/controllers/console/flask_admission.py @@ -18,6 +18,7 @@ from controllers.console.wraps import ( from core.logging.context import get_request_id, get_trace_id from enums import DeploymentEdition from libs.login import current_account_with_tenant, login_required +from libs.oauth_bearer import bearer_feature_required from machinery.context import RequestContext from models.account import TenantAccountRole from services.system_feature_service import SystemFeatureService @@ -46,6 +47,7 @@ def console_account_admission[T, **P, R]( require_change_email_enabled: bool = False, require_initialized: bool = True, require_valid_enterprise_license: bool = False, + require_oauth_bearer_enabled: bool = False, allowed_roles: frozenset[TenantAccountRole] | None = None, rbac_checks: Sequence[RBACCheck] | None = None, ) -> Callable[ @@ -56,8 +58,8 @@ def console_account_admission[T, **P, R]( All combinations use this decorator factory. Requirements are data, while the execution order stays fixed: edition, setup, login/CSRF, optional - account initialization, optional enterprise license, role/RBAC checks, then - context construction. + account initialization, optional enterprise license and feature checks, + role/RBAC checks, then context construction. """ def decorator( @@ -91,6 +93,8 @@ def console_account_admission[T, **P, R]( admitted: Callable[Concatenate[T, P], R | Response] = inject_request_context if require_change_email_enabled: admitted = enable_change_email(admitted) + if require_oauth_bearer_enabled: + admitted = bearer_feature_required(admitted) if require_valid_enterprise_license: admitted = enterprise_license_required(admitted) if require_initialized: diff --git a/api/controllers/openapi/_models.py b/api/controllers/openapi/_models.py index 0cc37e77da8..56bed1a8600 100644 --- a/api/controllers/openapi/_models.py +++ b/api/controllers/openapi/_models.py @@ -2,16 +2,17 @@ from __future__ import annotations +from datetime import datetime from enum import StrEnum from typing import Any, Final, Literal, Self from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator +from constants.oauth_bearer import SubjectType from controllers.common.human_input import HumanInputFormSubmitPayload from controllers.openapi._upload import UploadPart, UploadParts from enums import DeploymentEdition from libs.helper import EmailStr, UUIDStr, UUIDStrOrEmpty, uuid_value -from libs.oauth_bearer import SubjectType from models.model import AppMode from services.app_dsl_service import Import @@ -252,6 +253,14 @@ class DeviceMutateResponse(BaseModel): status: str +class DeviceApprovalContextResponse(BaseModel): + subject_email: str + subject_issuer: str + user_code: str + csrf_token: str + expires_at: datetime + + class ServerVersionResponse(BaseModel): """Meta endpoint payload for `GET /openapi/v1/_version` — no auth required.""" @@ -416,24 +425,6 @@ class PermittedExternalAppsListQuery(PageQuery): name: str | None = Field(None, max_length=200) -_EMAIL_FIELD = Field(min_length=3, max_length=320, pattern=r"^[^@\s]+@[^@\s]+$") - - -class ExtSubjectAssertionClaims(BaseModel): - email: str = _EMAIL_FIELD - issuer: str = Field(min_length=1, max_length=255) - user_code: str = Field(min_length=1, max_length=32) - nonce: str = Field(min_length=1, max_length=128) - - -class ApprovalGrantClaimsPayload(BaseModel): - subject_email: str = _EMAIL_FIELD - subject_issuer: str = Field(min_length=1, max_length=255) - user_code: str = Field(min_length=1, max_length=32) - nonce: str = Field(min_length=1, max_length=128) - csrf_token: str = Field(min_length=1, max_length=128) - - # Closed enum for invite/update-role payloads. Owner is intentionally not # assignable through these endpoints — ownership transfer goes through the # console's three-step email-verification flow. diff --git a/api/controllers/openapi/account.py b/api/controllers/openapi/account.py index 4e80879a556..70f32086057 100644 --- a/api/controllers/openapi/account.py +++ b/api/controllers/openapi/account.py @@ -7,6 +7,7 @@ from flask import request from flask_restx import Resource from werkzeug.exceptions import NotFound, Unauthorized +from constants.oauth_bearer import Scope from controllers.openapi import openapi_ns from controllers.openapi._contract import Example, Kind, endpoint from controllers.openapi._models import ( @@ -23,7 +24,6 @@ from controllers.openapi.auth.requirements import CheckScope, CheckSubject from controllers.openapi.auth.subjects import AccountSubject from core.logging.context import get_request_id, get_trace_id from extensions.ext_application_services import application_services -from libs.oauth_bearer import Scope from libs.rate_limit import LIMIT_ME_PER_ACCOUNT, enforce from machinery.context import AccountRequestContext from services.account_errors import AccountNotFoundError, AccountSessionNotFoundError diff --git a/api/controllers/openapi/app_dsl.py b/api/controllers/openapi/app_dsl.py index e9c3bff44e2..b9698bd7221 100644 --- a/api/controllers/openapi/app_dsl.py +++ b/api/controllers/openapi/app_dsl.py @@ -4,6 +4,7 @@ from flask_restx import Resource from sqlalchemy.orm import Session from werkzeug.exceptions import Forbidden +from constants.oauth_bearer import Scope from controllers.common.rbac import PlainApp, RBACCheck, RBACPermission, Workspace from controllers.openapi import openapi_ns from controllers.openapi._contract import Example, Kind, endpoint, op_of @@ -25,7 +26,6 @@ from controllers.openapi.auth.requirements import ( ) from controllers.openapi.auth.subjects import AccountSubject from extensions.ext_database import db -from libs.oauth_bearer import Scope from models.account import TenantAccountRole from services.app_dsl_service import AppDslService, Import from services.entities.dsl_entities import CheckDependenciesResult, ImportStatus diff --git a/api/controllers/openapi/app_run.py b/api/controllers/openapi/app_run.py index ca10e0f61ea..7b5d8d3acfe 100644 --- a/api/controllers/openapi/app_run.py +++ b/api/controllers/openapi/app_run.py @@ -20,6 +20,7 @@ from werkzeug.exceptions import ( ) import services +from constants.oauth_bearer import Scope from controllers.common.fields import EventStreamResponse from controllers.common.rbac import PlainApp, RBACCheck, RBACPermission from controllers.openapi import openapi_ns @@ -70,7 +71,6 @@ from extensions.ext_redis import redis_client from graphon.graph_engine.manager import GraphEngineManager from graphon.model_runtime.errors.invoke import InvokeError from libs import helper -from libs.oauth_bearer import Scope from models.model import App, AppMode from services.app_generate_service import AppGenerateService from services.errors.app import ( diff --git a/api/controllers/openapi/apps.py b/api/controllers/openapi/apps.py index 09f825f7178..f01838e37e7 100644 --- a/api/controllers/openapi/apps.py +++ b/api/controllers/openapi/apps.py @@ -9,6 +9,7 @@ from flask_restx import Resource from sqlalchemy.orm import Session from configs import dify_config +from constants.oauth_bearer import Scope from controllers.common.app_access import AppAccessFilter, resolve_app_access_filter from controllers.common.fields import Parameters from controllers.common.rbac import PlainApp, RBACCheck, RBACPermission @@ -35,7 +36,6 @@ from controllers.openapi.auth.requirements import ( from controllers.openapi.auth.subjects import AccountSubject from controllers.service_api.app.error import AppUnavailableError from core.app.app_config.common.parameters_mapping import get_parameters_from_feature_dict -from libs.oauth_bearer import Scope from models import App from models.enums import AppStatus from models.model import AppMode diff --git a/api/controllers/openapi/apps_permitted_external.py b/api/controllers/openapi/apps_permitted_external.py index 2087eb79d89..f6b094c7338 100644 --- a/api/controllers/openapi/apps_permitted_external.py +++ b/api/controllers/openapi/apps_permitted_external.py @@ -9,6 +9,7 @@ from __future__ import annotations from flask_restx import Resource +from constants.oauth_bearer import Scope from controllers.openapi import openapi_ns from controllers.openapi._contract import Example, Kind, endpoint from controllers.openapi._models import ( @@ -28,7 +29,6 @@ from controllers.openapi.auth.requirements import ( ) from controllers.openapi.auth.subjects import ExternalSsoSubject from enums import DeploymentEdition -from libs.oauth_bearer import Scope from models import App from models.enums import AppStatus from services.account_service import TenantService diff --git a/api/controllers/openapi/auth/requirements.py b/api/controllers/openapi/auth/requirements.py index a6446cadabd..6fc6ab1c0ee 100644 --- a/api/controllers/openapi/auth/requirements.py +++ b/api/controllers/openapi/auth/requirements.py @@ -18,13 +18,13 @@ from sqlalchemy.orm import Session from werkzeug.exceptions import Forbidden from configs import dify_config +from constants.oauth_bearer import Scope from controllers.common.rbac import RBACCheck, enforce_rbac_checks from controllers.openapi._audit import emit_wrong_surface from controllers.openapi.auth.context import Context from controllers.openapi.auth.loaders import load_app, load_caller, load_workspace, load_workspace_role from controllers.openapi.auth.subjects import Subject from enums import DeploymentEdition -from libs.oauth_bearer import Scope from models.account import TenantAccountRole from models.enums import CreatorUserRole from services.enterprise.enterprise_service import EnterpriseService, WebAppAccessMode diff --git a/api/controllers/openapi/auth/subjects.py b/api/controllers/openapi/auth/subjects.py index ba50cbaa16b..84f22e479f2 100644 --- a/api/controllers/openapi/auth/subjects.py +++ b/api/controllers/openapi/auth/subjects.py @@ -1,4 +1,4 @@ -"""`TokenType` in `libs/oauth_bearer` owns the mint-time facts (prefix, subject, +"""`TokenType` in `constants/oauth_bearer` owns the mint-time facts (prefix, subject, scopes); a `Subject` owns the request-time behaviour, so `libs/` never has to import the auth layer. """ @@ -12,11 +12,12 @@ from typing import ClassVar, override from sqlalchemy.orm import Session from werkzeug.exceptions import Unauthorized +from constants.oauth_bearer import Scope, SubjectType from controllers.openapi.auth.context import Context from controllers.openapi.auth.data import ExternalIdentity from controllers.openapi.auth.loaders import load_app, load_workspace, route_has_app from extensions.ext_application_services import application_services -from libs.oauth_bearer import AuthContext, Scope, SubjectType +from libs.oauth_bearer import AuthContext from models.account import Account from models.enums import CreatorUserRole, EndUserType from models.model import EndUser diff --git a/api/controllers/openapi/files.py b/api/controllers/openapi/files.py index 8808835dbf9..69ba44c7414 100644 --- a/api/controllers/openapi/files.py +++ b/api/controllers/openapi/files.py @@ -4,6 +4,7 @@ from __future__ import annotations from flask_restx import Resource +from constants.oauth_bearer import Scope from controllers.openapi import openapi_ns from controllers.openapi._contract import Example, Kind, endpoint from controllers.openapi._files import end_read_transaction, upload @@ -18,7 +19,6 @@ from controllers.openapi.auth.requirements import ( ) from controllers.openapi.auth.subjects import AccountSubject, ExternalSsoSubject from fields.file_fields import FileResponse -from libs.oauth_bearer import Scope @openapi_ns.route("/apps//files") diff --git a/api/controllers/openapi/flask_admission.py b/api/controllers/openapi/flask_admission.py new file mode 100644 index 00000000000..590803a0332 --- /dev/null +++ b/api/controllers/openapi/flask_admission.py @@ -0,0 +1,32 @@ +"""Flask admission adapter for OAuth device SSO endpoints.""" + +from __future__ import annotations + +from collections.abc import Callable +from functools import wraps +from typing import Concatenate + +from flask import request +from werkzeug.exceptions import NotFound + +from core.logging.context import get_request_id, get_trace_id +from extensions.ext_application_services import application_services +from services.oauth_device_contracts import DeviceRequestContext + + +def oauth_device_sso_admission[**P, R]( + view: Callable[Concatenate[DeviceRequestContext, P], R], +) -> Callable[P, R]: + """Require an active Enterprise license and inject stable request metadata.""" + + @wraps(view) + def admitted(*args: P.args, **kwargs: P.kwargs) -> R: + if not application_services().feature_queries.has_valid_enterprise_license(): + raise NotFound() + context = DeviceRequestContext( + request_id=get_request_id(), + trace_id=get_trace_id() or request.headers.get("X-Trace-Id"), + ) + return view(context, *args, **kwargs) + + return admitted diff --git a/api/controllers/openapi/human_input_form.py b/api/controllers/openapi/human_input_form.py index e71d4cd88cc..e4fdff91108 100644 --- a/api/controllers/openapi/human_input_form.py +++ b/api/controllers/openapi/human_input_form.py @@ -10,6 +10,7 @@ from flask_restx import Resource from sqlalchemy.orm import Session from werkzeug.exceptions import BadRequest +from constants.oauth_bearer import Scope from controllers.common.human_input import stringify_form_default_values from controllers.common.rbac import PlainApp, RBACCheck, RBACPermission from controllers.openapi import openapi_ns @@ -40,7 +41,6 @@ from core.db.session_factory import session_factory from core.workflow.human_input_policy import HumanInputSurface, is_recipient_type_allowed_for_surface from extensions.ext_database import db from libs.helper import to_timestamp -from libs.oauth_bearer import Scope from models.enums import CreatorUserRole from models.model import App from services.human_input_service import FormNotFoundError, HumanInputService diff --git a/api/controllers/openapi/oauth_device.py b/api/controllers/openapi/oauth_device.py index 2c418150fe3..686bc686a9f 100644 --- a/api/controllers/openapi/oauth_device.py +++ b/api/controllers/openapi/oauth_device.py @@ -1,40 +1,22 @@ -"""Device-flow endpoints under /openapi/v1/oauth/device/*. Two -sub-groups in one module: +"""Device-flow HTTP adapters under ``/openapi/v1/oauth/device/*``. - Protocol (RFC 8628, public + rate-limited): - POST /oauth/device/code - POST /oauth/device/token - GET /oauth/device/lookup - - Approval (account branch, console-cookie authed): - POST /oauth/device/approve - POST /oauth/device/deny - -SSO branch lives in oauth_device_sso.py. +Protocol endpoints are public and rate-limited. Approval endpoints use +Console account admission. The SSO branch lives in ``oauth_device_sso.py``. """ from __future__ import annotations -import logging -from typing import Any +from http import HTTPStatus from flask import request -from flask_login import login_required from flask_restx import Resource from pydantic import BaseModel, ValidationError from werkzeug.exceptions import BadRequest -from configs import dify_config from controllers.common.schema import query_params_from_model -from controllers.console.wraps import ( - account_initialization_required, - setup_required, - with_current_tenant_id, - with_current_user, -) +from controllers.console.flask_admission import console_account_admission from controllers.openapi import openapi_ns from controllers.openapi._models import ( - AccountPayload, DeviceCodeRequest, DeviceCodeResponse, DeviceLookupQuery, @@ -43,354 +25,146 @@ from controllers.openapi._models import ( DeviceMutateResponse, DevicePollRequest, DeviceTokenResponse, - WorkspacePayload, ) -from extensions.ext_database import db -from extensions.ext_redis import redis_client -from libs.helper import extract_remote_ip -from libs.oauth_bearer import SubjectType, TokenType, bearer_feature_required +from extensions.ext_application_services import application_services +from libs.helper import dump_response, extract_remote_ip from libs.rate_limit import ( LIMIT_DEVICE_CODE_PER_IP, LIMIT_DEVICE_FLOW_APPROVE, LIMIT_LOOKUP_PUBLIC, rate_limit, ) -from models import Account -from services.account_service import TenantService -from services.oauth_device_flow import ( - ACCOUNT_ISSUER_SENTINEL, - DEFAULT_POLL_INTERVAL_SECONDS, - DEVICE_FLOW_TTL_SECONDS, - DeviceFlowRedis, - DeviceFlowStatus, - InvalidTransitionError, - PollPayload, - SlowDownDecision, - StateNotFoundError, - mint_oauth_token, - oauth_ttl_days, +from machinery.context import RequestContext +from services.oauth_device_contracts import ( + AccessDeniedError, + AlreadyResolvedError, + ApprovalInProgressError, + ApprovalOutcomeUnknownError, + AuthorizationPendingError, + DeviceStateLostError, + ExpiredOrUnknownError, + ExpiredTokenError, + OAuthDeviceError, + PollTooFastError, + UnsupportedClientError, ) -logger = logging.getLogger(__name__) - - -# ========================================================================= -# Validation helpers -# ========================================================================= - def _validate_json[M: BaseModel](model: type[M]) -> M: body = request.get_json(silent=True) or {} try: return model.model_validate(body) - except ValidationError as exc: - raise BadRequest(str(exc)) + except ValidationError as error: + raise BadRequest(str(error)) from error def _validate_query[M: BaseModel](model: type[M]) -> M: try: return model.model_validate(request.args.to_dict(flat=True)) - except ValidationError as exc: - raise BadRequest(str(exc)) + except ValidationError as error: + raise BadRequest(str(error)) from error -# ========================================================================= -# Protocol endpoints — RFC 8628 (public + per-IP rate limit) -# ========================================================================= +def _error_response(error: OAuthDeviceError) -> tuple[dict[str, str], int]: + match error: + case UnsupportedClientError(): + return {"error": "unsupported_client"}, HTTPStatus.BAD_REQUEST + case PollTooFastError(): + return {"error": "slow_down"}, HTTPStatus.BAD_REQUEST + case ExpiredTokenError(): + return {"error": "expired_token"}, HTTPStatus.BAD_REQUEST + case AuthorizationPendingError(): + return {"error": "authorization_pending"}, HTTPStatus.BAD_REQUEST + case AccessDeniedError(): + return {"error": "access_denied"}, HTTPStatus.BAD_REQUEST + case ExpiredOrUnknownError(): + return {"error": "expired_or_unknown"}, HTTPStatus.NOT_FOUND + case AlreadyResolvedError(): + return {"error": "already_resolved"}, HTTPStatus.CONFLICT + case ApprovalInProgressError(): + return {"error": "approve_in_progress"}, HTTPStatus.CONFLICT + case ApprovalOutcomeUnknownError(): + return {"error": "approval_outcome_unknown"}, HTTPStatus.SERVICE_UNAVAILABLE + case DeviceStateLostError(): + return {"error": "state_lost"}, HTTPStatus.CONFLICT + raise RuntimeError(f"unmapped OAuth device error: {type(error).__name__}") @openapi_ns.route("/oauth/device/code") class OAuthDeviceCodeApi(Resource): @openapi_ns.expect(openapi_ns.models[DeviceCodeRequest.__name__]) - @openapi_ns.response(200, "Device code created", openapi_ns.models[DeviceCodeResponse.__name__]) + @openapi_ns.response(HTTPStatus.OK, "Device code created", openapi_ns.models[DeviceCodeResponse.__name__]) @rate_limit(LIMIT_DEVICE_CODE_PER_IP) def post(self): payload = _validate_json(DeviceCodeRequest) - client_id = payload.client_id - device_label = payload.device_label - - if client_id not in dify_config.OPENAPI_KNOWN_CLIENT_IDS: - return {"error": "unsupported_client"}, 400 - - store = DeviceFlowRedis(redis_client) - ip = extract_remote_ip(request) - device_code, user_code, expires_in = store.start(client_id, device_label, created_ip=ip) - - return { - "device_code": device_code, - "user_code": user_code, - "verification_uri": _verification_uri(), - "expires_in": expires_in, - "interval": DEFAULT_POLL_INTERVAL_SECONDS, - }, 200 + try: + authorization = application_services().oauth_device.start( + client_id=payload.client_id, + device_label=payload.device_label, + created_ip=extract_remote_ip(request), + request_origin=request.host_url, + ) + except OAuthDeviceError as error: + return _error_response(error) + return dump_response(DeviceCodeResponse, authorization), HTTPStatus.OK @openapi_ns.route("/oauth/device/token") class OAuthDeviceTokenApi(Resource): - """RFC 8628 poll.""" + """RFC 8628 poll endpoint.""" @openapi_ns.expect(openapi_ns.models[DevicePollRequest.__name__]) - @openapi_ns.response(200, "Device token", openapi_ns.models[DeviceTokenResponse.__name__]) + @openapi_ns.response(HTTPStatus.OK, "Device token", openapi_ns.models[DeviceTokenResponse.__name__]) def post(self): payload = _validate_json(DevicePollRequest) - device_code = payload.device_code - - store = DeviceFlowRedis(redis_client) - - # slow_down beats every other branch — polling-too-fast clients - # see only that response regardless of underlying state. - if store.record_poll(device_code, DEFAULT_POLL_INTERVAL_SECONDS) is SlowDownDecision.SLOW_DOWN: - return {"error": "slow_down"}, 400 - - state = store.load_by_device_code(device_code) - if state is None: - return {"error": "expired_token"}, 400 - - if state.status is DeviceFlowStatus.PENDING: - return {"error": "authorization_pending"}, 400 - - terminal = store.consume_on_poll(device_code) - if terminal is None: - return {"error": "expired_token"}, 400 - - if terminal.status is DeviceFlowStatus.DENIED: - return {"error": "access_denied"}, 400 - - poll_payload: PollPayload | dict[str, Any] = terminal.poll_payload or {} - if "token" not in poll_payload: - logger.error("device_flow: approved state missing poll_payload for %s", device_code) - return {"error": "expired_token"}, 400 - - _audit_cross_ip_if_needed(state) - return poll_payload, 200 + try: + token = application_services().oauth_device.poll( + device_code=payload.device_code, + poll_ip=extract_remote_ip(request), + ) + except OAuthDeviceError as error: + return _error_response(error) + return dump_response(DeviceTokenResponse, token, exclude_unset=True), HTTPStatus.OK @openapi_ns.route("/oauth/device/lookup") class OAuthDeviceLookupApi(Resource): - """Read-only — public for pre-validate before login. user_code is - high-entropy + short-TTL; per-IP rate limit blocks enumeration. - """ + """Public pre-login validation for a high-entropy, short-lived user code.""" @openapi_ns.doc(params=query_params_from_model(DeviceLookupQuery)) - @openapi_ns.response(200, "Device lookup result", openapi_ns.models[DeviceLookupResponse.__name__]) + @openapi_ns.response(HTTPStatus.OK, "Device lookup result", openapi_ns.models[DeviceLookupResponse.__name__]) @rate_limit(LIMIT_LOOKUP_PUBLIC) def get(self): payload = _validate_query(DeviceLookupQuery) - user_code = payload.user_code.strip().upper() - - store = DeviceFlowRedis(redis_client) - found = store.load_by_user_code(user_code) - if found is None: - return {"valid": False, "expires_in_remaining": 0, "client_id": None}, 200 - - _device_code, state = found - if state.status is not DeviceFlowStatus.PENDING: - return {"valid": False, "expires_in_remaining": 0, "client_id": state.client_id}, 200 - - return { - "valid": True, - "expires_in_remaining": DEVICE_FLOW_TTL_SECONDS, - "client_id": state.client_id, - }, 200 - - -# ========================================================================= -# Approval endpoints — account branch (cookie-authed) -# ========================================================================= - - -_APPROVE_GUARD_KEY_FMT = "device_code:{code}:approving" -_APPROVE_GUARD_TTL_SECONDS = 10 + lookup = application_services().oauth_device.lookup(user_code=payload.user_code) + return dump_response(DeviceLookupResponse, lookup), HTTPStatus.OK @openapi_ns.route("/oauth/device/approve") class DeviceApproveApi(Resource): @openapi_ns.expect(openapi_ns.models[DeviceMutateRequest.__name__]) - @openapi_ns.response(200, "Approved", openapi_ns.models[DeviceMutateResponse.__name__]) - @setup_required - @login_required - @account_initialization_required - @bearer_feature_required + @openapi_ns.response(HTTPStatus.OK, "Approved", openapi_ns.models[DeviceMutateResponse.__name__]) + @console_account_admission(require_oauth_bearer_enabled=True) @rate_limit(LIMIT_DEVICE_FLOW_APPROVE) - @with_current_user - @with_current_tenant_id - def post(self, tenant: str, account: Account): + def post(self, request_context: RequestContext): payload = _validate_json(DeviceMutateRequest) - user_code = payload.user_code.strip().upper() - - store = DeviceFlowRedis(redis_client) - - found = store.load_by_user_code(user_code) - if found is None: - return {"error": "expired_or_unknown"}, 404 - device_code, state = found - if state.status is not DeviceFlowStatus.PENDING: - return {"error": "already_resolved"}, 409 - - # SET NX guard — without it, two in-flight approves both pass - # PENDING, both mint, and the second upsert silently rotates the - # first caller into an already-revoked token. - guard_key = _APPROVE_GUARD_KEY_FMT.format(code=device_code) - if not redis_client.set(guard_key, "1", nx=True, ex=_APPROVE_GUARD_TTL_SECONDS): - return {"error": "approve_in_progress"}, 409 - try: - ttl_days = oauth_ttl_days(tenant_id=tenant) - mint = mint_oauth_token( - redis_client, - subject_email=account.email, - subject_issuer=ACCOUNT_ISSUER_SENTINEL, - account_id=str(account.id), - client_id=state.client_id, - device_label=state.device_label, - token_type=TokenType.OAUTH_ACCOUNT, - ttl_days=ttl_days, - session=db.session(), - ) - - poll_payload = _build_account_poll_payload(account, tenant, mint) - try: - store.approve( - device_code, - subject_email=account.email, - account_id=str(account.id), - subject_issuer=ACCOUNT_ISSUER_SENTINEL, - minted_token=mint.token, - token_id=str(mint.token_id), - poll_payload=poll_payload, - ) - except (StateNotFoundError, InvalidTransitionError): - # Row minted but state vanished — roll forward; the orphan - # token is revocable via auth devices list / Authorized Apps. - logger.exception("device_flow: approve raced on %s", device_code) - return {"error": "state_lost"}, 409 - finally: - redis_client.delete(guard_key) - - _emit_approve_audit(state, account, tenant, mint) - return {"status": "approved"}, 200 + result = application_services().oauth_device.approve(request_context, user_code=payload.user_code) + except OAuthDeviceError as error: + return _error_response(error) + return dump_response(DeviceMutateResponse, result), HTTPStatus.OK @openapi_ns.route("/oauth/device/deny") class DeviceDenyApi(Resource): @openapi_ns.expect(openapi_ns.models[DeviceMutateRequest.__name__]) - @openapi_ns.response(200, "Denied", openapi_ns.models[DeviceMutateResponse.__name__]) - @setup_required - @login_required - @account_initialization_required - @bearer_feature_required + @openapi_ns.response(HTTPStatus.OK, "Denied", openapi_ns.models[DeviceMutateResponse.__name__]) + @console_account_admission(require_oauth_bearer_enabled=True) @rate_limit(LIMIT_DEVICE_FLOW_APPROVE) - def post(self): + def post(self, _request_context: RequestContext): payload = _validate_json(DeviceMutateRequest) - user_code = payload.user_code.strip().upper() - - store = DeviceFlowRedis(redis_client) - found = store.load_by_user_code(user_code) - if found is None: - return {"error": "expired_or_unknown"}, 404 - device_code, state = found - if state.status is not DeviceFlowStatus.PENDING: - return {"error": "already_resolved"}, 409 - try: - store.deny(device_code) - except (StateNotFoundError, InvalidTransitionError): - logger.exception("device_flow: deny raced on %s", device_code) - return {"error": "state_lost"}, 409 - - _emit_deny_audit(state) - return {"status": "denied"}, 200 - - -# ========================================================================= -# Helpers -# ========================================================================= - - -def _verification_uri() -> str: - base = getattr(dify_config, "CONSOLE_WEB_URL", None) - if base: - return f"{base.rstrip('/')}/device" - return f"{request.host_url.rstrip('/')}/device" - - -def _audit_cross_ip_if_needed(state) -> None: - poll_ip = extract_remote_ip(request) - if state.created_ip and poll_ip and poll_ip != state.created_ip: - logger.warning( - "audit: oauth.device_code_cross_ip_poll token_id=%s creation_ip=%s poll_ip=%s", - state.token_id, - state.created_ip, - poll_ip, - extra={ - "audit": True, - "token_id": state.token_id, - "creation_ip": state.created_ip, - "poll_ip": poll_ip, - }, - ) - - -def _build_account_poll_payload(account, tenant, mint) -> PollPayload: - rows = TenantService.get_workspaces_for_account(str(account.id), session=db.session()) - workspaces = [WorkspacePayload(id=str(t.id), name=t.name, role=getattr(m, "role", "")) for t, m in rows] - # Prefer active session tenant → DB-flagged current join → first membership. - default_ws_id = None - if tenant and any(w.id == str(tenant) for w in workspaces): - default_ws_id = str(tenant) - if default_ws_id is None: - for _t, m in rows: - if getattr(m, "current", False): - default_ws_id = str(m.tenant_id) - break - if default_ws_id is None and workspaces: - default_ws_id = workspaces[0].id - - payload: PollPayload = { - "token": mint.token, - "expires_at": mint.expires_at.isoformat(), - "subject_type": SubjectType.ACCOUNT, - "account": AccountPayload(id=str(account.id), email=account.email, name=account.name).model_dump(mode="json"), - "workspaces": [w.model_dump(mode="json") for w in workspaces], - "default_workspace_id": default_ws_id, - "token_id": str(mint.token_id), - } - return payload - - -def _emit_approve_audit(state, account, tenant, mint) -> None: - logger.warning( - "audit: oauth.device_flow_approved token_id=%s subject=%s client_id=%s device_label=%s rotated=? expires_at=%s", - mint.token_id, - account.email, - state.client_id, - state.device_label, - mint.expires_at, - extra={ - "audit": True, - "event": "oauth.device_flow_approved", - "token_id": str(mint.token_id), - "subject_type": SubjectType.ACCOUNT, - "subject_email": account.email, - "account_id": str(account.id), - "tenant_id": tenant, - "client_id": state.client_id, - "device_label": state.device_label, - "scopes": ["full"], - "expires_at": mint.expires_at.isoformat(), - }, - ) - - -def _emit_deny_audit(state) -> None: - logger.warning( - "audit: oauth.device_flow_denied client_id=%s device_label=%s", - state.client_id, - state.device_label, - extra={ - "audit": True, - "event": "oauth.device_flow_denied", - "client_id": state.client_id, - "device_label": state.device_label, - }, - ) + result = application_services().oauth_device.deny(user_code=payload.user_code) + except OAuthDeviceError as error: + return _error_response(error) + return dump_response(DeviceMutateResponse, result), HTTPStatus.OK diff --git a/api/controllers/openapi/oauth_device_sso.py b/api/controllers/openapi/oauth_device_sso.py index 99519fa1890..7de457501d0 100644 --- a/api/controllers/openapi/oauth_device_sso.py +++ b/api/controllers/openapi/oauth_device_sso.py @@ -1,384 +1,186 @@ -"""SSO-branch device-flow endpoints under /openapi/v1/oauth/device/*. -EE-only. Browser flow: - - GET /oauth/device/sso-initiate → 302 to IdP authorize URL - GET /oauth/device/sso-complete → ACS callback, sets approval-grant cookie - GET /oauth/device/approval-context → SPA reads cookie claims (idempotent) - POST /oauth/device/approve-external → mints dfoe_ token + clears cookie - -Function-based (raw @bp.route) rather than Resource classes because the -handlers do redirects + cookie kwargs that don't fit the Resource shape. -""" +"""HTTP adapters for the external-SSO branch of the OAuth device flow.""" from __future__ import annotations import logging import re -import secrets -from dataclasses import dataclass +from http import HTTPStatus +from typing import Never from urllib.parse import urlencode from flask import jsonify, make_response, redirect, request -from pydantic import ValidationError +from pydantic import BaseModel, ValidationError from werkzeug.exceptions import ( BadGateway, BadRequest, Conflict, Forbidden, NotFound, + ServiceUnavailable, + TooManyRequests, Unauthorized, ) -from configs import dify_config from controllers.openapi import bp -from controllers.openapi._models import ExtSubjectAssertionClaims -from extensions.ext_database import db -from extensions.ext_redis import redis_client -from libs import jws +from controllers.openapi._models import ( + DeviceApprovalContextResponse, + DeviceMutateRequest, + DeviceMutateResponse, +) +from controllers.openapi.flask_admission import oauth_device_sso_admission +from extensions.ext_application_services import application_services from libs.device_flow_security import ( APPROVAL_GRANT_COOKIE_NAME, - ApprovalGrantClaims, approval_grant_cleared_cookie_kwargs, approval_grant_cookie_kwargs, - consume_approval_grant_nonce, - consume_sso_assertion_nonce, - enterprise_only, - mint_approval_grant, - verify_approval_grant, ) -from libs.oauth_bearer import SubjectType, TokenType -from libs.rate_limit import ( - LIMIT_APPROVE_EXT_PER_EMAIL, - LIMIT_SSO_INITIATE_PER_IP, - enforce, - rate_limit, -) -from services.account_service import AccountService -from services.enterprise.enterprise_service import EnterpriseService -from services.oauth_device_flow import ( - DeviceFlowRedis, - DeviceFlowStatus, - InvalidTransitionError, - PollPayload, - StateNotFoundError, - mint_oauth_token, - oauth_ttl_days, +from libs.helper import dump_response +from libs.rate_limit import LIMIT_SSO_INITIATE_PER_IP, rate_limit +from services.oauth_device_contracts import ( + AlreadyResolvedError, + ApprovalInProgressError, + ApprovalOutcomeUnknownError, + ApprovalSessionConsumedError, + DeviceRequestContext, + DeviceSSOCompletion, + DeviceStateLostError, + ExternalApprovalCSRFError, + ExternalApprovalRateLimitError, + ExternalIdentityConflictError, + ExternalUserCodeMismatchError, + ExternalUserCodeNotFoundError, + InvalidApprovalSessionError, + InvalidUserCodeError, + OAuthDeviceError, + OAuthDeviceSSOConfigurationError, + OAuthDeviceSSOInitiationError, ) logger = logging.getLogger(__name__) - -# Matches DEVICE_FLOW_TTL_SECONDS so the signed state can't outlive the -# device_code it references. -STATE_ENVELOPE_TTL_SECONDS = 15 * 60 - -# Canonical sso-complete path. IdP-side ACS callback URL must point here. -_SSO_COMPLETE_PATH = "/openapi/v1/oauth/device/sso-complete" - -_ALLOWED_SSO_ERRORS = {"sso_failed", "email_belongs_to_dify_account"} - -# user_code only ever reaches the redirect as a urlencoded query value; the -# charset bound additionally forbids the path/scheme separators a redirection -# attack would need, so an untrusted value cannot escape the fixed /device path. +_ALLOWED_SSO_ERRORS = frozenset({"sso_failed", "email_belongs_to_dify_account"}) _USER_CODE_RE = re.compile(r"\A[A-Z0-9-]{1,16}\Z") +def _validate_json[M: BaseModel](model: type[M]) -> M: + try: + return model.model_validate(request.get_json(silent=True) or {}) + except ValidationError as error: + raise BadRequest(str(error)) from error + + def _device_error_redirect(code: str, user_code: str | None = None): safe_code = code if code in _ALLOWED_SSO_ERRORS else "sso_failed" params: dict[str, str] = {"sso_error": safe_code} - if user_code and _USER_CODE_RE.match(user_code): + if user_code and _USER_CODE_RE.fullmatch(user_code): params["user_code"] = user_code - return redirect(f"/device?{urlencode(params)}", code=302) + return redirect(f"/device?{urlencode(params)}", code=HTTPStatus.FOUND) -def _trusted_origin() -> str: - base = (dify_config.CONSOLE_API_URL or "").rstrip("/") - if not base: - raise BadGateway("console_api_url_unset") - return base +def _completion_response(result: DeviceSSOCompletion): + if result.error_code: + return _device_error_redirect(result.error_code, result.user_code) + if not result.approval_grant: + return _device_error_redirect("sso_failed", result.user_code) + + response = redirect("/device?sso_verified=1", code=HTTPStatus.FOUND) + response.set_cookie(**approval_grant_cookie_kwargs(result.approval_grant)) + return response + + +def _raise_http_error(error: OAuthDeviceError, *, session_error: str = "invalid_session") -> Never: + match error: + case InvalidUserCodeError(): + raise BadRequest("invalid_user_code") from error + case OAuthDeviceSSOConfigurationError(): + raise BadGateway("console_api_url_unset") from error + case OAuthDeviceSSOInitiationError(): + raise BadGateway(str(error) or "sso_initiate_failed") from error + case InvalidApprovalSessionError(): + raise Unauthorized(session_error) from error + case ExternalApprovalRateLimitError(): + raise TooManyRequests("rate_limited") from error + case ExternalApprovalCSRFError(): + raise Forbidden("csrf_mismatch") from error + case ExternalUserCodeMismatchError(): + raise BadRequest("user_code_mismatch") from error + case ExternalUserCodeNotFoundError(): + raise NotFound("user_code_not_pending") from error + case AlreadyResolvedError(): + raise Conflict("user_code_not_pending") from error + case ApprovalInProgressError(): + raise Conflict("approve_in_progress") from error + case ApprovalOutcomeUnknownError(): + raise ServiceUnavailable("approval_outcome_unknown") from error + case ExternalIdentityConflictError(): + raise Forbidden("email_belongs_to_dify_account") from error + case ApprovalSessionConsumedError(): + raise Unauthorized("session_already_consumed") from error + case DeviceStateLostError(): + raise Conflict("state_lost") from error + raise RuntimeError(f"unmapped OAuth device SSO error: {type(error).__name__}") @bp.route("/oauth/device/sso-initiate", methods=["GET"]) -@enterprise_only +@oauth_device_sso_admission @rate_limit(LIMIT_SSO_INITIATE_PER_IP) -def sso_initiate(): - user_code = (request.args.get("user_code") or "").strip().upper() +def sso_initiate(context: DeviceRequestContext): + user_code = (request.args.get("user_code") or "").strip() if not user_code: raise BadRequest("user_code required") - - store = DeviceFlowRedis(redis_client) - found = store.load_by_user_code(user_code) - if found is None: - raise BadRequest("invalid_user_code") - _, state = found - if state.status is not DeviceFlowStatus.PENDING: - raise BadRequest("invalid_user_code") - - origin = _trusted_origin() - keyset = jws.KeySet.from_shared_secret() - signed_state = jws.sign( - keyset, - payload={ - "redirect_url": "", - "app_code": "", - "intent": "device_flow", - "user_code": user_code, - "nonce": secrets.token_urlsafe(16), - "return_to": "", - "idp_callback_url": f"{origin}{_SSO_COMPLETE_PATH}", - }, - aud=jws.AUD_STATE_ENVELOPE, - ttl_seconds=STATE_ENVELOPE_TTL_SECONDS, - ) - try: - reply = EnterpriseService.initiate_device_flow_sso(signed_state) - except Exception as e: - logger.warning("sso-initiate: enterprise call failed: %s", e) - raise BadGateway("sso_initiate_failed") from e + result = application_services().oauth_device.initiate_sso(context, user_code=user_code) + except OAuthDeviceError as error: + _raise_http_error(error) - url = (reply or {}).get("url") - if not url: - raise BadGateway("sso_initiate_missing_url") - - # Clear stale approval-grant — defends against cross-tab/back-button mixing. - resp = redirect(url, code=302) - resp.set_cookie(**approval_grant_cleared_cookie_kwargs()) - return resp + response = redirect(result.redirect_url, code=HTTPStatus.FOUND) + response.set_cookie(**approval_grant_cleared_cookie_kwargs()) + return response @bp.route("/oauth/device/sso-complete", methods=["GET"]) -@enterprise_only -def sso_complete(): +@oauth_device_sso_admission +def sso_complete(context: DeviceRequestContext): try: - return _sso_complete_impl() + result = application_services().oauth_device.complete_sso( + context, + inbound_error=request.args.get("sso_error"), + inbound_user_code=request.args.get("user_code"), + assertion=request.args.get("sso_assertion"), + ) except Exception: logger.exception("sso-complete: unhandled") return _device_error_redirect("sso_failed") - - -def _sso_complete_impl(): - inbound_error = request.args.get("sso_error") - if inbound_error: - return _device_error_redirect(inbound_error, request.args.get("user_code")) - - blob = request.args.get("sso_assertion") - if not blob: - return _device_error_redirect("sso_failed") - - keyset = jws.KeySet.from_shared_secret() - - try: - raw_claims = jws.verify(keyset, blob, expected_aud=jws.AUD_EXT_SUBJECT_ASSERTION) - except jws.VerifyError as e: - logger.warning("sso-complete: rejected assertion: %s", e) - return _device_error_redirect("sso_failed") - - try: - claims = ExtSubjectAssertionClaims.model_validate(raw_claims) - except ValidationError as e: - logger.warning("sso-complete: claim shape invalid: %s", e) - return _device_error_redirect("sso_failed") - - user_code = claims.user_code.strip().upper() - - if not consume_sso_assertion_nonce(redis_client, claims.nonce): - return _device_error_redirect("sso_failed", user_code) - - store = DeviceFlowRedis(redis_client) - found = store.load_by_user_code(user_code) - if found is None: - return _device_error_redirect("sso_failed", user_code) - _, state = found - if state.status is not DeviceFlowStatus.PENDING: - return _device_error_redirect("sso_failed", user_code) - - if AccountService.has_active_account_with_email(claims.email, session=db.session()): - _emit_external_rejection_audit( - state, - _RejectedClaims(subject_email=claims.email, subject_issuer=claims.issuer), - reason="email_belongs_to_dify_account", - ) - return _device_error_redirect("email_belongs_to_dify_account", user_code) - - iss = _trusted_origin() - cookie_value, _ = mint_approval_grant( - keyset=keyset, - iss=iss, - subject_email=claims.email, - subject_issuer=claims.issuer, - user_code=user_code, - ) - - resp = redirect("/device?sso_verified=1", code=302) - resp.set_cookie(**approval_grant_cookie_kwargs(cookie_value)) - return resp + return _completion_response(result) @bp.route("/oauth/device/approval-context", methods=["GET"]) -@enterprise_only -def approval_context(): - token = request.cookies.get(APPROVAL_GRANT_COOKIE_NAME) - if not token: - raise Unauthorized("no_session") - - keyset = jws.KeySet.from_shared_secret() +@oauth_device_sso_admission +def approval_context(context: DeviceRequestContext): try: - claims = verify_approval_grant(keyset, token) - except jws.VerifyError as e: - logger.warning("approval-context: bad cookie: %s", e) - raise Unauthorized("no_session") from e - - return jsonify( - { - "subject_email": claims.subject_email, - "subject_issuer": claims.subject_issuer, - "user_code": claims.user_code, - "csrf_token": claims.csrf_token, - "expires_at": claims.expires_at.isoformat(), - } - ), 200 + result = application_services().oauth_device.get_approval_context( + context, + approval_grant=request.cookies.get(APPROVAL_GRANT_COOKIE_NAME, ""), + ) + except OAuthDeviceError as error: + _raise_http_error(error, session_error="no_session") + return dump_response(DeviceApprovalContextResponse, result), HTTPStatus.OK @bp.route("/oauth/device/approve-external", methods=["POST"]) -@enterprise_only -def approve_external(): - token = request.cookies.get(APPROVAL_GRANT_COOKIE_NAME) - if not token: - raise Unauthorized("invalid_session") - - keyset = jws.KeySet.from_shared_secret() +@oauth_device_sso_admission +def approve_external(context: DeviceRequestContext): + payload = _validate_json(DeviceMutateRequest) try: - claims: ApprovalGrantClaims = verify_approval_grant(keyset, token) - except jws.VerifyError as e: - logger.warning("approve-external: bad cookie: %s", e) - raise Unauthorized("invalid_session") from e - - enforce(LIMIT_APPROVE_EXT_PER_EMAIL, key=f"subject:{claims.subject_email}") - - csrf_header = request.headers.get("X-CSRF-Token", "") - if not csrf_header or not secrets.compare_digest(csrf_header, claims.csrf_token): - raise Forbidden("csrf_mismatch") - - data = request.get_json(silent=True) or {} - body_user_code = (data.get("user_code") or "").strip().upper() - if body_user_code != claims.user_code: - raise BadRequest("user_code_mismatch") - - store = DeviceFlowRedis(redis_client) - found = store.load_by_user_code(claims.user_code) - if found is None: - raise NotFound("user_code_not_pending") - device_code, state = found - if state.status is not DeviceFlowStatus.PENDING: - raise Conflict("user_code_not_pending") - - if AccountService.has_active_account_with_email(claims.subject_email, session=db.session()): - _emit_external_rejection_audit(state, claims, reason="email_belongs_to_dify_account") - raise Forbidden("email_belongs_to_dify_account") - - if not consume_approval_grant_nonce(redis_client, claims.nonce): - raise Unauthorized("session_already_consumed") - - ttl_days = oauth_ttl_days(tenant_id=None) - mint = mint_oauth_token( - redis_client, - subject_email=claims.subject_email, - subject_issuer=claims.subject_issuer, - account_id=None, - client_id=state.client_id, - device_label=state.device_label, - token_type=TokenType.OAUTH_EXTERNAL_SSO, - ttl_days=ttl_days, - session=db.session(), - ) - - # SSO branch of the shared PollPayload contract: account/workspace - # fields are zero-filled (`None` / `[]`) for parity with the account - # branch in `oauth_device._build_account_poll_payload`. - poll_payload: PollPayload = { - "token": mint.token, - "expires_at": mint.expires_at.isoformat(), - "subject_type": SubjectType.EXTERNAL_SSO, - "subject_email": claims.subject_email, - "subject_issuer": claims.subject_issuer, - "account": None, - "workspaces": [], - "default_workspace_id": None, - "token_id": str(mint.token_id), - } - - try: - store.approve( - device_code, - subject_email=claims.subject_email, - account_id=None, - subject_issuer=claims.subject_issuer, - minted_token=mint.token, - token_id=str(mint.token_id), - poll_payload=poll_payload, + result = application_services().oauth_device.approve_external( + context, + approval_grant=request.cookies.get(APPROVAL_GRANT_COOKIE_NAME, ""), + csrf_token=request.headers.get("X-CSRF-Token", ""), + user_code=payload.user_code, ) - except (StateNotFoundError, InvalidTransitionError) as e: - logger.exception("approve-external: state transition raced") - raise Conflict("state_lost") from e + except OAuthDeviceError as error: + _raise_http_error(error) - _emit_approve_external_audit(state, claims, mint) - - resp = make_response(jsonify({"status": "approved"}), 200) - resp.set_cookie(**approval_grant_cleared_cookie_kwargs()) - return resp - - -@dataclass(frozen=True) -class _RejectedClaims: - """Minimal subject shape consumed by `_emit_external_rejection_audit`. - - Mirrors the attributes used from `ApprovalGrantClaims` so callers holding - only a raw JWS claims dict (e.g. `sso_complete`) can emit the same audit - event without reaching for the full dataclass. - """ - - subject_email: str - subject_issuer: str - - -def _emit_external_rejection_audit(state, claims, *, reason: str) -> None: - logger.warning( - "audit: oauth.device_flow_rejected subject_type=%s subject_email=%s subject_issuer=%s reason=%s", - SubjectType.EXTERNAL_SSO, - claims.subject_email, - claims.subject_issuer, - reason, - extra={ - "audit": True, - "event": "oauth.device_flow_rejected", - "subject_type": SubjectType.EXTERNAL_SSO, - "subject_email": claims.subject_email, - "subject_issuer": claims.subject_issuer, - "reason": reason, - "client_id": state.client_id, - "device_label": state.device_label, - }, - ) - - -def _emit_approve_external_audit(state, claims, mint) -> None: - logger.warning( - "audit: oauth.device_flow_approved subject_type=%s subject_email=%s subject_issuer=%s token_id=%s", - SubjectType.EXTERNAL_SSO, - claims.subject_email, - claims.subject_issuer, - mint.token_id, - extra={ - "audit": True, - "event": "oauth.device_flow_approved", - "subject_type": SubjectType.EXTERNAL_SSO, - "subject_email": claims.subject_email, - "subject_issuer": claims.subject_issuer, - "token_id": str(mint.token_id), - "client_id": state.client_id, - "device_label": state.device_label, - "scopes": ["apps:run"], - "expires_at": mint.expires_at.isoformat(), - }, - ) + response = make_response(jsonify(dump_response(DeviceMutateResponse, result)), HTTPStatus.OK) + response.set_cookie(**approval_grant_cleared_cookie_kwargs()) + return response diff --git a/api/controllers/openapi/workflow_events.py b/api/controllers/openapi/workflow_events.py index 4377844d476..154d7fe13d7 100644 --- a/api/controllers/openapi/workflow_events.py +++ b/api/controllers/openapi/workflow_events.py @@ -16,6 +16,7 @@ from flask_restx import Resource from pydantic import BaseModel, Field from werkzeug.exceptions import NotFound, UnprocessableEntity +from constants.oauth_bearer import Scope from controllers.common.fields import EventStreamResponse from controllers.common.rbac import PlainApp, RBACCheck, RBACPermission from controllers.openapi import openapi_ns @@ -39,7 +40,6 @@ from core.app.apps.workflow.app_generator import WorkflowAppGenerator from core.app.entities.task_entities import StreamEvent from core.db.session_factory import session_factory from core.workflow.human_input_policy import HumanInputSurface -from libs.oauth_bearer import Scope from models.model import AppMode from repositories.factory import DifyAPIRepositoryFactory from services.workflow_event_snapshot_service import build_workflow_event_stream diff --git a/api/controllers/openapi/workspaces.py b/api/controllers/openapi/workspaces.py index d0aba47e73d..007575460af 100644 --- a/api/controllers/openapi/workspaces.py +++ b/api/controllers/openapi/workspaces.py @@ -19,6 +19,7 @@ from flask_restx import Resource from werkzeug.exceptions import BadRequest, NotFound from configs import dify_config +from constants.oauth_bearer import Scope from controllers.common.rbac import RBACCheck, RBACPermission, Workspace from controllers.openapi import openapi_ns from controllers.openapi._contract import Example, Kind, endpoint @@ -46,7 +47,6 @@ from controllers.openapi.auth.requirements import ( ) from controllers.openapi.auth.subjects import AccountSubject from enums import DeploymentEdition -from libs.oauth_bearer import Scope from models import Account, Tenant, TenantAccountJoin from models.account import TenantAccountRole from services.account_service import AccountService, RegisterService, TenantService diff --git a/api/extensions/ext_application_services.py b/api/extensions/ext_application_services.py index 168ad732d0c..808d8222414 100644 --- a/api/extensions/ext_application_services.py +++ b/api/extensions/ext_application_services.py @@ -56,6 +56,7 @@ from repositories.installation_state_repository import InstallationStateReposito from repositories.installed_app_access_repository import SQLAlchemyInstalledAppAccessRepository from repositories.message_file_preview_repository import MessageFilePreviewQueryRepository from repositories.oauth_access_token_repository import SQLAlchemyOAuthAccessTokenRepository +from repositories.oauth_device_token_repository import SQLAlchemyOAuthDeviceTokenRepository from repositories.oauth_server_repository import RedisOAuthServerTokenRepository, SQLAlchemyOAuthServerRepository from repositories.plugin_file_upload_repository import SQLAlchemyPluginFileUploadOwnerRepository from repositories.recommended_app_catalog_repository import DatabaseRecommendedAppCatalogRepository @@ -183,6 +184,19 @@ from services.message_suggested_questions_service import MessageSuggestedQuestio from services.notification_gateway import BillingNotificationGateway from services.notification_service import NotificationService from services.notion_data_source_gateway import NotionDataSourceGateway +from services.oauth_device_adapters import ( + DifyConfigOAuthDeviceSettings, + EnterpriseOAuthDeviceSSOGateway, + EnvironmentOAuthDeviceTokenTTLPolicy, + OAuthDeviceTokenIssuanceGateway, + RedisExternalApprovalLimiter, +) +from services.oauth_device_application_service import ( + DeviceWorkspaceQuery, + OAuthDeviceAccountQuery, + OAuthDeviceApplicationService, +) +from services.oauth_device_flow import DeviceFlowRedis from services.oauth_server_service import OAUTH_ACCESS_TOKEN_EXPIRES_IN, OAuthServerService from services.partner_tenant_binding_service import PartnerTenantBindingService from services.plugin_file_upload_gateway import ToolFilePluginUploadGateway @@ -309,6 +323,7 @@ class ApplicationServices: tool_file_downloads: ToolFileDownloadService upload_file_delivery: UploadFileDeliveryService oauth_server: OAuthServerService + oauth_device: OAuthDeviceApplicationService init_validation: InitValidationService installed_app_access: InstalledAppAccessService notifications: NotificationService @@ -373,6 +388,32 @@ def _build_oauth_server_service( ) +def _build_oauth_device_service( + *, + database_client: sessionmaker[Session], + redis: RedisClientWrapper, + accounts: OAuthDeviceAccountQuery, + workspaces: DeviceWorkspaceQuery, +) -> OAuthDeviceApplicationService: + token_repository = SQLAlchemyOAuthDeviceTokenRepository(session_factory=database_client, redis=redis) + return OAuthDeviceApplicationService( + store=DeviceFlowRedis(redis), + accounts=accounts, + workspaces=workspaces, + tokens=OAuthDeviceTokenIssuanceGateway( + tokens=token_repository, + ttl_policy=EnvironmentOAuthDeviceTokenTTLPolicy(), + ), + sessions=token_repository, + sso=EnterpriseOAuthDeviceSSOGateway( + redis=redis, + enterprise_service=EnterpriseService(), + ), + external_approval_limiter=RedisExternalApprovalLimiter(redis=redis), + settings=DifyConfigOAuthDeviceSettings(), + ) + + def _build_file_grant_service(*, database_client: sessionmaker[Session]) -> FileGrantService: repository = FileGrantRepository(session_factory=database_client) return FileGrantService( @@ -752,6 +793,12 @@ def build_application_services( storage=storage, ), oauth_server=_build_oauth_server_service(database_client=database_client, redis=redis), + oauth_device=_build_oauth_device_service( + database_client=database_client, + redis=redis, + accounts=accounts, + workspaces=workspace_query_repository, + ), init_validation=InitValidationService( state=installation_state, validation_required=(deployment_edition != DeploymentEdition.CLOUD and bool(initialization_password)), diff --git a/api/extensions/ext_redis.py b/api/extensions/ext_redis.py index 7cdce7191fd..9b271178163 100644 --- a/api/extensions/ext_redis.py +++ b/api/extensions/ext_redis.py @@ -5,7 +5,7 @@ import ssl import sys from collections.abc import Callable from datetime import timedelta -from typing import Any, Union, cast +from typing import Any, Protocol, Union, cast, runtime_checkable import redis from redis import RedisError @@ -41,6 +41,11 @@ _serialize_redis_name_arg = serialize_redis_name_arg _serialize_redis_name_args = serialize_redis_name_args +@runtime_checkable +class _ScriptRegistrar(Protocol): + def register_script(self, script: str) -> Callable[..., Any]: ... + + class RedisClientWrapper: """ A wrapper class for the Redis client that addresses the issue where the global @@ -217,6 +222,28 @@ class RedisClientWrapper: def pipeline(self, transaction: bool = True, shard_hint: str | None = None) -> Any: return self._require_client().pipeline(transaction=transaction, shard_hint=shard_hint) + def register_script(self, script: str) -> Callable[..., Any]: + """Register a Lua script whose key arguments use logical Redis names.""" + client = self._require_client() + if not isinstance(client, _ScriptRegistrar): + raise RuntimeError("Redis client does not support Lua script registration") + registered_script = client.register_script(script) + + def execute( + keys: list[str | bytes] | tuple[str | bytes, ...] | None = None, + args: list[Any] | tuple[Any, ...] | None = None, + client: RedisClientWrapper | redis.Redis | RedisCluster | None = None, + ) -> Any: + redis_keys = _serialize_redis_name_args(tuple(keys or ()), self._get_prefix()) + execution_client = client._require_client() if isinstance(client, RedisClientWrapper) else client + return registered_script( + keys=redis_keys, + args=args, + client=execution_client or self._require_client(), + ) + + return execute + def __getattr__(self, item: str) -> Any: return getattr(self._require_client(), item) diff --git a/api/libs/device_flow_security.py b/api/libs/device_flow_security.py index 1e7ec3683c3..cbd90506272 100644 --- a/api/libs/device_flow_security.py +++ b/api/libs/device_flow_security.py @@ -1,52 +1,21 @@ -"""Device-flow security primitives: enterprise_only gate, approval-grant -cookie mint/verify/consume, and anti-framing headers. -""" +"""Device-flow approval-grant security and anti-framing primitives.""" from __future__ import annotations import logging import secrets -from collections.abc import Callable from dataclasses import dataclass from datetime import UTC, datetime, timedelta -from functools import wraps from flask import Blueprint -from pydantic import ValidationError -from werkzeug.exceptions import NotFound +from pydantic import BaseModel, Field, ValidationError from libs import jws from libs.token import is_secure -from services.entities.feature_entities import LicenseStatus -from services.system_feature_service import SystemFeatureService logger = logging.getLogger(__name__) -# ============================================================================ -# enterprise_only decorator -# ============================================================================ - - -# Fail-closed: any non-EE-active status (default NONE on CE, plus INACTIVE / EXPIRED / LOST) -# is denied. Future LicenseStatus values default to denial unless explicitly admitted. -_EE_ENABLED_STATUSES = {LicenseStatus.ACTIVE, LicenseStatus.EXPIRING} - - -def enterprise_only[**P, R](view: Callable[P, R]) -> Callable[P, R]: - """404 on CE, passthrough on EE. Apply before rate-limit so CE - responses don't consume the bucket. - """ - - @wraps(view) - def decorated(*args: P.args, **kwargs: P.kwargs): - if SystemFeatureService.get_license_status() not in _EE_ENABLED_STATUSES: - raise NotFound() - return view(*args, **kwargs) - - return decorated - - # ============================================================================ # approval_grant cookie # ============================================================================ @@ -70,6 +39,17 @@ class ApprovalGrantClaims: expires_at: datetime +_EMAIL_FIELD = Field(min_length=3, max_length=320, pattern=r"^[^@\s]+@[^@\s]+$") + + +class _ApprovalGrantClaimsPayload(BaseModel): + subject_email: str = _EMAIL_FIELD + subject_issuer: str = Field(min_length=1, max_length=255) + user_code: str = Field(min_length=1, max_length=32) + nonce: str = Field(min_length=1, max_length=128) + csrf_token: str = Field(min_length=1, max_length=128) + + def mint_approval_grant( *, keyset: jws.KeySet, @@ -108,12 +88,9 @@ def mint_approval_grant( def verify_approval_grant(keyset: jws.KeySet, token: str) -> ApprovalGrantClaims: """Sig + aud + exp only — nonce consumption is the caller's job.""" - # lazy import: breaks libs → controllers cycle - from controllers.openapi._models import ApprovalGrantClaimsPayload - raw = jws.verify(keyset, token, expected_aud=jws.AUD_APPROVAL_GRANT) try: - parsed = ApprovalGrantClaimsPayload.model_validate(raw) + parsed = _ApprovalGrantClaimsPayload.model_validate(raw) except ValidationError as e: raise jws.VerifyError(f"claim shape invalid: {e}") from e @@ -127,19 +104,6 @@ def verify_approval_grant(keyset: jws.KeySet, token: str) -> ApprovalGrantClaims ) -def consume_approval_grant_nonce(redis_client, nonce: str) -> bool: - if not nonce: - return False - return bool( - redis_client.set( - NONCE_KEY_FMT.format(nonce=nonce), - "1", - nx=True, - ex=NONCE_TTL_SECONDS, - ) - ) - - def consume_sso_assertion_nonce(redis_client, nonce: str) -> bool: if not nonce: return False diff --git a/api/libs/oauth_bearer.py b/api/libs/oauth_bearer.py index f29064bda51..0832cd8580f 100644 --- a/api/libs/oauth_bearer.py +++ b/api/libs/oauth_bearer.py @@ -10,7 +10,6 @@ from collections.abc import Callable, Mapping from contextvars import ContextVar, Token from dataclasses import dataclass from datetime import UTC, datetime -from enum import StrEnum from functools import wraps from typing import Literal, Protocol @@ -19,6 +18,7 @@ from sqlalchemy.orm import Session from werkzeug.exceptions import ServiceUnavailable from configs import dify_config +from constants.oauth_bearer import TOKEN_CACHE_KEY_FMT, Scope, SubjectType, TokenType from libs.rate_limit import enforce_bearer_rate_limit from models import OAuthAccessToken @@ -30,63 +30,6 @@ logger = logging.getLogger(__name__) # ============================================================================ -class Scope(StrEnum): - """Catalog of bearer scopes recognised by the openapi surface. - - `FULL` is the catch-all carried by `dfoa_` account tokens — it satisfies - any per-route scope requirement. `dfoe_` tokens carry the per-feature - scopes (`APPS_RUN`, `APPS_READ_PERMITTED_EXTERNAL`). - """ - - FULL = "full" - APPS_READ = "apps:read" - APPS_READ_PERMITTED_EXTERNAL = "apps:read:permitted-external" - APPS_RUN = "apps:run" - WORKSPACE_READ = "workspace:read" - WORKSPACE_WRITE = "workspace:write" - - -class SubjectType(StrEnum): - # Annotation-only names are not members; they declare what `__new__` attaches. - scopes: frozenset[Scope] - # Whether the subject's token rows carry an `account_id`. A subject that does - # not is keyed by its external issuer instead. - bound_to_account: bool - - ACCOUNT = ("account", frozenset({Scope.FULL}), True) - EXTERNAL_SSO = ("external_sso", frozenset({Scope.APPS_RUN, Scope.APPS_READ_PERMITTED_EXTERNAL}), False) - - def __new__(cls, value: str, scopes: frozenset[Scope], bound_to_account: bool) -> SubjectType: - obj = str.__new__(cls, value) - obj._value_ = value - obj.scopes = scopes - obj.bound_to_account = bound_to_account - return obj - - -class TokenType(StrEnum): - """Many token types may serve one subject; the prefix is the only thing a - bearer string reveals, so it lives here and nowhere else. - """ - - prefix: str - subject: SubjectType - - OAUTH_ACCOUNT = ("oauth_account", "dfoa_", SubjectType.ACCOUNT) - OAUTH_EXTERNAL_SSO = ("oauth_external_sso", "dfoe_", SubjectType.EXTERNAL_SSO) - - def __new__(cls, value: str, prefix: str, subject: SubjectType) -> TokenType: - obj = str.__new__(cls, value) - obj._value_ = value - obj.prefix = prefix - obj.subject = subject - return obj - - @classmethod - def for_token(cls, token: str) -> TokenType | None: - return next((token_type for token_type in cls if token.startswith(token_type.prefix)), None) - - @dataclass(frozen=True, slots=True) class AuthContext: """Per-request identity published via :data:`_auth_ctx_var`. Subject and @@ -212,7 +155,6 @@ class BearerAuthenticator: # OAuth access token resolver (PAT resolver would be a sibling class) # ============================================================================ -TOKEN_CACHE_KEY_FMT = "auth:token:{hash}" POSITIVE_TTL_SECONDS = 60 NEGATIVE_TTL_SECONDS = 10 AUDIT_OAUTH_EXPIRED = "oauth.token_expired" diff --git a/api/repositories/account_repository.py b/api/repositories/account_repository.py index 65c352384d6..46500d04638 100644 --- a/api/repositories/account_repository.py +++ b/api/repositories/account_repository.py @@ -3,7 +3,7 @@ from datetime import datetime from typing import override -from sqlalchemy import case, delete, select +from sqlalchemy import case, delete, func, select from sqlalchemy.orm import Session, sessionmaker from models.account import Account, AccountIntegrate, AccountStatus, InvitationCode, InvitationCodeStatus @@ -185,6 +185,22 @@ class SQLAlchemyAccountRepository(AccountRepository, ConsoleAuthAccountRepositor with self._session_factory() as session: return session.scalar(select(Account.id).where(Account.email == email).limit(1)) is not None + @override + def has_active_email(self, email: str) -> bool: + normalized = email.strip().lower() + if not normalized: + return False + stmt = ( + select(Account.id) + .where( + func.lower(Account.email) == normalized, + Account.status == AccountStatus.ACTIVE, + ) + .limit(1) + ) + with self._session_factory() as session: + return session.scalar(stmt) is not None + @override def reset_email( self, diff --git a/api/repositories/oauth_device_token_repository.py b/api/repositories/oauth_device_token_repository.py new file mode 100644 index 00000000000..c338c92f7d9 --- /dev/null +++ b/api/repositories/oauth_device_token_repository.py @@ -0,0 +1,193 @@ +"""SQLAlchemy persistence adapter for OAuth device tokens and sessions.""" + +import logging +from datetime import UTC, datetime + +from sqlalchemy import func, select, update +from sqlalchemy.orm import Session, sessionmaker + +from constants.oauth_bearer import TOKEN_CACHE_KEY_FMT +from extensions.ext_redis import RedisClientWrapper +from models.oauth import OAuthAccessToken +from services.oauth_device_contracts import ( + OAuthDeviceSession, + OAuthDeviceSessionPage, + OAuthDeviceTokenRotation, + OAuthDeviceTokenWrite, +) + +logger = logging.getLogger(__name__) + + +class SQLAlchemyOAuthDeviceTokenRepository: + def __init__(self, *, session_factory: sessionmaker[Session], redis: RedisClientWrapper) -> None: + self._session_factory = session_factory + self._redis = redis + + def rotate_token(self, token: OAuthDeviceTokenWrite) -> OAuthDeviceTokenRotation: + old_hash: str | None = None + replaced_token_id: str | None = None + with self._session_factory() as session: + prior = session.execute( + select(OAuthAccessToken.id, OAuthAccessToken.token_hash) + .where( + OAuthAccessToken.subject_email == token.subject_email, + OAuthAccessToken.subject_issuer == token.subject_issuer, + OAuthAccessToken.client_id == token.client_id, + OAuthAccessToken.device_label == token.device_label, + OAuthAccessToken.revoked_at.is_(None), + ) + .limit(1) + .with_for_update() + ).first() + if prior is not None: + old_hash = prior.token_hash + replaced_token_id = str(prior.id) + session.execute( + update(OAuthAccessToken).where(OAuthAccessToken.id == prior.id).values(revoked_at=datetime.now(UTC)) + ) + + record = OAuthAccessToken( + subject_email=token.subject_email, + subject_issuer=token.subject_issuer, + account_id=token.account_id, + client_id=token.client_id, + device_label=token.device_label, + prefix=token.prefix, + token_hash=token.token_hash, + expires_at=token.expires_at, + ) + session.add(record) + session.flush() + token_id = str(record.id) + session.commit() + + rotation = OAuthDeviceTokenRotation( + token_id=token_id, + replaced_token_id=replaced_token_id, + replaced_token_hash=old_hash, + ) + if old_hash: + try: + self._redis.delete(TOKEN_CACHE_KEY_FMT.format(hash=old_hash)) + except Exception: + self.rollback_rotation(rotation) + raise + return rotation + + def rollback_rotation(self, rotation: OAuthDeviceTokenRotation) -> bool: + new_hash: str | None = None + with self._session_factory() as session: + current = session.scalar( + select(OAuthAccessToken) + .where( + OAuthAccessToken.id == rotation.token_id, + OAuthAccessToken.revoked_at.is_(None), + ) + .with_for_update() + ) + # A later successful rotation already superseded this token. In + # that case restoring its predecessor would create two live rows. + if current is None: + return False + + new_hash = current.token_hash + session.delete(current) + session.flush() + if rotation.replaced_token_id is not None: + predecessor = session.scalar( + select(OAuthAccessToken) + .where( + OAuthAccessToken.id == rotation.replaced_token_id, + OAuthAccessToken.revoked_at.is_not(None), + ) + .with_for_update() + ) + if predecessor is None: + raise RuntimeError("OAuth token rotation predecessor could not be restored") + predecessor.revoked_at = None + session.commit() + + cache_keys = [ + TOKEN_CACHE_KEY_FMT.format(hash=token_hash) + for token_hash in (new_hash, rotation.replaced_token_hash) + if token_hash + ] + if cache_keys: + try: + self._redis.delete(*cache_keys) + except Exception: + logger.exception("failed to invalidate rolled-back OAuth token caches") + return True + + def list_account_sessions( + self, + *, + account_id: str, + page: int, + limit: int, + ) -> OAuthDeviceSessionPage: + now = datetime.now(UTC) + filters = ( + OAuthAccessToken.account_id == account_id, + OAuthAccessToken.revoked_at.is_(None), + OAuthAccessToken.token_hash.is_not(None), + OAuthAccessToken.expires_at > now, + ) + with self._session_factory() as session: + total = session.scalar(select(func.count()).select_from(OAuthAccessToken).where(*filters)) or 0 + records = session.scalars( + select(OAuthAccessToken) + .where(*filters) + .order_by(OAuthAccessToken.created_at.desc()) + .offset((page - 1) * limit) + .limit(limit) + ).all() + + return OAuthDeviceSessionPage( + page=page, + limit=limit, + total=total, + items=tuple( + OAuthDeviceSession( + id=str(record.id), + prefix=record.prefix, + client_id=record.client_id, + device_label=record.device_label, + created_at=record.created_at, + last_used_at=record.last_used_at, + expires_at=record.expires_at, + ) + for record in records + ), + ) + + def revoke_account_session(self, *, account_id: str, token_id: str) -> bool: + old_hash: str | None + with self._session_factory() as session: + prior = session.execute( + select(OAuthAccessToken.id, OAuthAccessToken.token_hash) + .where( + OAuthAccessToken.id == token_id, + OAuthAccessToken.account_id == account_id, + OAuthAccessToken.revoked_at.is_(None), + ) + .with_for_update() + ).first() + if prior is None: + return False + + old_hash = prior.token_hash + session.execute( + update(OAuthAccessToken) + .where( + OAuthAccessToken.id == prior.id, + OAuthAccessToken.revoked_at.is_(None), + ) + .values(revoked_at=datetime.now(UTC), token_hash=None) + ) + session.commit() + + if old_hash: + self._redis.delete(TOKEN_CACHE_KEY_FMT.format(hash=old_hash)) + return True diff --git a/api/repositories/workspace_query_repository.py b/api/repositories/workspace_query_repository.py index 8e83d158c47..eba63eef4f1 100644 --- a/api/repositories/workspace_query_repository.py +++ b/api/repositories/workspace_query_repository.py @@ -9,6 +9,8 @@ from models.account import Tenant, TenantAccountJoin, TenantStatus from services.account_login_service import ConsoleAuthWorkspaceQuery from services.account_ports import AccountWorkspaceMembershipQuery, AccountWorkspaceSnapshotQuery from services.entities.account_access_entities import AccountWorkspaceSnapshot +from services.oauth_device_application_service import DeviceWorkspaceQuery +from services.oauth_device_contracts import DeviceWorkspace from services.workspace_query_service import WorkspaceQuery, WorkspaceRecord @@ -17,6 +19,7 @@ class WorkspaceQueryRepository( AccountWorkspaceMembershipQuery, AccountWorkspaceSnapshotQuery, ConsoleAuthWorkspaceQuery, + DeviceWorkspaceQuery, ): def __init__(self, session_factory: sessionmaker[Session]) -> None: self._session_factory = session_factory @@ -87,6 +90,18 @@ class WorkspaceQueryRepository( for workspace_id, name, role, current in session.execute(stmt).all() ) + @override + def list_for_device_flow(self, account_id: str) -> tuple[DeviceWorkspace, ...]: + return tuple( + DeviceWorkspace( + id=workspace.id, + name=workspace.name, + role=workspace.role, + current=workspace.current, + ) + for workspace in self.list_account_access_workspaces(account_id) + ) + @override def has_active_for_account(self, account_id: str) -> bool: stmt = ( diff --git a/api/services/account_ports.py b/api/services/account_ports.py index bc0b5d1eeb9..6914dc3fdf6 100644 --- a/api/services/account_ports.py +++ b/api/services/account_ports.py @@ -50,6 +50,8 @@ class AccountRepository(Protocol): def email_exists(self, email: str) -> bool: ... + def has_active_email(self, email: str) -> bool: ... + def reset_email(self, account_id: str, *, expected_old_email: str, new_email: str) -> AccountEmailResetResult: ... diff --git a/api/services/account_service.py b/api/services/account_service.py index 8f5e36c3e73..5764229fde6 100644 --- a/api/services/account_service.py +++ b/api/services/account_service.py @@ -285,27 +285,9 @@ class AccountService: @staticmethod def get_account_by_email(email: str, *, session: Session) -> Account | None: - """Plain ``Account`` getter keyed by email. Case-sensitive — use - :meth:`has_active_account_with_email` for the case-insensitive - existence check that backs the SSO collision rule. - """ + """Plain ``Account`` getter keyed by case-sensitive email.""" return session.execute(select(Account).where(Account.email == email)).scalar_one_or_none() - @staticmethod - def has_active_account_with_email(email: str, *, session: Session) -> bool: - if not email: - return False - normalized = email.strip().lower() - if not normalized: - return False - row = session.execute( - select(Account.id).where( - func.lower(Account.email) == normalized, - Account.status == AccountStatus.ACTIVE, - ) - ).scalar_one_or_none() - return row is not None - @staticmethod def has_account_with_normalized_email(email: str, *, session: Session) -> bool: """Check the normalized identity through its indexed column.""" diff --git a/api/services/feature_query_service.py b/api/services/feature_query_service.py index 83fac59fd87..089b87ddb22 100644 --- a/api/services/feature_query_service.py +++ b/api/services/feature_query_service.py @@ -6,10 +6,13 @@ from machinery.context import RequestContext from services.entities.feature_entities import ( FeatureModel, LicenseModel, + LicenseStatus, SystemFeatureModel, VectorSpaceLimitationModel, ) +_VALID_ENTERPRISE_LICENSE_STATUSES = frozenset({LicenseStatus.ACTIVE, LicenseStatus.EXPIRING}) + class FeatureQueryGateway(Protocol): """Read dynamic feature resources without exposing their current implementation.""" @@ -58,3 +61,7 @@ class FeatureQueryService: def get_license(self) -> LicenseModel: return self._features.get_license() + + def has_valid_enterprise_license(self) -> bool: + status = self._features.get_public_system_features().license.status + return status in _VALID_ENTERPRISE_LICENSE_STATUSES diff --git a/api/services/oauth_device_adapters.py b/api/services/oauth_device_adapters.py new file mode 100644 index 00000000000..354f8a9208e --- /dev/null +++ b/api/services/oauth_device_adapters.py @@ -0,0 +1,336 @@ +"""Outer adapters for the OAuth device external-SSO use cases.""" + +import logging +import os +import secrets +from collections.abc import Mapping +from datetime import UTC, datetime, timedelta +from typing import Protocol, override + +from pydantic import BaseModel, Field, ValidationError + +from configs import dify_config +from constants.oauth_bearer import TokenType +from extensions.ext_redis import RedisClientWrapper +from libs import jws +from libs.device_flow_security import ( + NONCE_KEY_FMT, + NONCE_TTL_SECONDS, + consume_sso_assertion_nonce, + mint_approval_grant, + verify_approval_grant, +) +from libs.helper import RateLimiter +from libs.oauth_bearer import sha256_hex +from libs.rate_limit import LIMIT_APPROVE_EXT_PER_EMAIL +from services.entities.account_entities import AccountSnapshot +from services.oauth_device_application_service import ( + ExternalApprovalLimiter, + OAuthDeviceSettings, + OAuthDeviceSSOGateway, + OAuthDeviceTokenIssuer, + OAuthDeviceTokenPersistence, + OAuthDeviceTokenTTLPolicy, +) +from services.oauth_device_contracts import ( + ACCOUNT_ISSUER_SENTINEL, + ExternalApprovalGrant, + ExternalSubjectAssertion, + InvalidApprovalSessionError, + InvalidSSOAssertionError, + IssuedOAuthToken, + OAuthDeviceSSOInitiationError, + OAuthDeviceTokenWrite, +) + +logger = logging.getLogger(__name__) + +_EMAIL_FIELD = Field(min_length=3, max_length=320, pattern=r"^[^@\s]+@[^@\s]+$") +_DEFAULT_OAUTH_TTL_DAYS = 14 +_MIN_OAUTH_TTL_DAYS = 1 +_MAX_OAUTH_TTL_DAYS = 365 +_TTL_ENV_VAR = "OAUTH_TTL_DAYS" +_OAUTH_TOKEN_BODY_BYTES = 32 +_RESERVE_APPROVAL_NONCE_LUA = """ +local current = redis.call('GET', KEYS[1]) +if current == ARGV[1] then return 1 end +if current then return 0 end +local stored = redis.call('SET', KEYS[1], ARGV[1], 'EX', ARGV[2], 'NX') +if stored then return 1 end +return 0 +""" +_RELEASE_APPROVAL_NONCE_LUA = """ +if redis.call('GET', KEYS[1]) == ARGV[1] then + return redis.call('DEL', KEYS[1]) +end +return 0 +""" + + +class EnterpriseDeviceSSOService(Protocol): + def initiate_device_flow_sso(self, signed_state: str) -> Mapping[str, object] | None: ... + + +class DifyConfigOAuthDeviceSettings(OAuthDeviceSettings): + @property + @override + def known_client_ids(self) -> frozenset[str]: + return dify_config.OPENAPI_KNOWN_CLIENT_IDS + + @property + @override + def verification_base_url(self) -> str | None: + return dify_config.CONSOLE_WEB_URL + + @property + @override + def sso_base_url(self) -> str | None: + return dify_config.CONSOLE_API_URL + + +class EnvironmentOAuthDeviceTokenTTLPolicy(OAuthDeviceTokenTTLPolicy): + @override + def ttl_days(self, workspace_id: str | None) -> int: + # Reserved for a future tenant-specific policy, which should take + # precedence over the deployment-wide environment value. + _ = workspace_id + raw = os.environ.get(_TTL_ENV_VAR) + if raw is None: + return _DEFAULT_OAUTH_TTL_DAYS + try: + value = int(raw) + except ValueError: + logger.warning( + "%s=%r is not an int; falling back to %d", + _TTL_ENV_VAR, + raw, + _DEFAULT_OAUTH_TTL_DAYS, + ) + return _DEFAULT_OAUTH_TTL_DAYS + if value < _MIN_OAUTH_TTL_DAYS: + logger.warning("%s=%d below min %d; clamping", _TTL_ENV_VAR, value, _MIN_OAUTH_TTL_DAYS) + return _MIN_OAUTH_TTL_DAYS + if value > _MAX_OAUTH_TTL_DAYS: + logger.warning("%s=%d above max %d; clamping", _TTL_ENV_VAR, value, _MAX_OAUTH_TTL_DAYS) + return _MAX_OAUTH_TTL_DAYS + return value + + +class OAuthDeviceTokenIssuanceGateway(OAuthDeviceTokenIssuer): + def __init__( + self, + *, + tokens: OAuthDeviceTokenPersistence, + ttl_policy: OAuthDeviceTokenTTLPolicy, + ) -> None: + self._tokens = tokens + self._ttl_policy = ttl_policy + + @override + def issue_account_token( + self, + *, + account: AccountSnapshot, + workspace_id: str, + client_id: str, + device_label: str, + ) -> IssuedOAuthToken: + return self._issue_token( + token_type=TokenType.OAUTH_ACCOUNT, + subject_email=account.email, + subject_issuer=ACCOUNT_ISSUER_SENTINEL, + account_id=account.id, + client_id=client_id, + device_label=device_label, + workspace_id=workspace_id, + ) + + @override + def issue_external_token( + self, + *, + subject_email: str, + subject_issuer: str, + client_id: str, + device_label: str, + ) -> IssuedOAuthToken: + if not subject_issuer.strip(): + raise ValueError("external-SSO token requires non-empty subject_issuer") + return self._issue_token( + token_type=TokenType.OAUTH_EXTERNAL_SSO, + subject_email=subject_email, + subject_issuer=subject_issuer, + account_id=None, + client_id=client_id, + device_label=device_label, + workspace_id=None, + ) + + def _issue_token( + self, + *, + token_type: TokenType, + subject_email: str, + subject_issuer: str, + account_id: str | None, + client_id: str, + device_label: str, + workspace_id: str | None, + ) -> IssuedOAuthToken: + plaintext = token_type.prefix + secrets.token_urlsafe(_OAUTH_TOKEN_BODY_BYTES) + expires_at = datetime.now(UTC) + timedelta(days=self._ttl_policy.ttl_days(workspace_id)) + rotation = self._tokens.rotate_token( + OAuthDeviceTokenWrite( + subject_email=subject_email, + subject_issuer=subject_issuer, + account_id=account_id, + client_id=client_id, + device_label=device_label, + prefix=token_type.prefix, + token_hash=sha256_hex(plaintext), + expires_at=expires_at, + ) + ) + return IssuedOAuthToken(token=plaintext, expires_at=expires_at.isoformat(), rotation=rotation) + + @override + def rollback_token(self, token: IssuedOAuthToken) -> bool: + return self._tokens.rollback_rotation(token.rotation) + + +class _ExternalSubjectAssertionPayload(BaseModel): + email: str = _EMAIL_FIELD + issuer: str = Field(min_length=1, max_length=255) + user_code: str = Field(min_length=1, max_length=32) + nonce: str = Field(min_length=1, max_length=128) + + +class EnterpriseOAuthDeviceSSOGateway(OAuthDeviceSSOGateway): + def __init__( + self, + *, + redis: RedisClientWrapper, + enterprise_service: EnterpriseDeviceSSOService, + ) -> None: + self._redis = redis + self._enterprise_service = enterprise_service + + @override + def initiate(self, *, user_code: str, callback_url: str, ttl_seconds: int) -> str: + keyset = jws.KeySet.from_shared_secret() + signed_state = jws.sign( + keyset, + payload={ + "redirect_url": "", + "app_code": "", + "intent": "device_flow", + "user_code": user_code, + "nonce": secrets.token_urlsafe(16), + "return_to": "", + "idp_callback_url": callback_url, + }, + aud=jws.AUD_STATE_ENVELOPE, + ttl_seconds=ttl_seconds, + ) + try: + reply = self._enterprise_service.initiate_device_flow_sso(signed_state) + except Exception as error: + logger.warning("oauth device SSO initiation failed: %s", error) + raise OAuthDeviceSSOInitiationError("sso_initiate_failed") from error + + redirect_url = (reply or {}).get("url") + if not isinstance(redirect_url, str) or not redirect_url: + raise OAuthDeviceSSOInitiationError("sso_initiate_missing_url") + return redirect_url + + @override + def verify_assertion(self, assertion: str) -> ExternalSubjectAssertion: + keyset = jws.KeySet.from_shared_secret() + try: + raw_claims = jws.verify(keyset, assertion, expected_aud=jws.AUD_EXT_SUBJECT_ASSERTION) + claims = _ExternalSubjectAssertionPayload.model_validate(raw_claims) + except (jws.VerifyError, ValidationError) as error: + raise InvalidSSOAssertionError(str(error)) from error + return ExternalSubjectAssertion( + subject_email=claims.email, + subject_issuer=claims.issuer, + user_code=claims.user_code, + nonce=claims.nonce, + ) + + @override + def mint_approval_grant( + self, + *, + issuer: str, + subject_email: str, + subject_issuer: str, + user_code: str, + ) -> str: + token, _claims = mint_approval_grant( + keyset=jws.KeySet.from_shared_secret(), + iss=issuer, + subject_email=subject_email, + subject_issuer=subject_issuer, + user_code=user_code, + ) + return token + + @override + def verify_approval_grant(self, token: str) -> ExternalApprovalGrant: + try: + claims = verify_approval_grant(jws.KeySet.from_shared_secret(), token) + except jws.VerifyError as error: + raise InvalidApprovalSessionError(str(error)) from error + return ExternalApprovalGrant( + subject_email=claims.subject_email, + subject_issuer=claims.subject_issuer, + user_code=claims.user_code, + nonce=claims.nonce, + csrf_token=claims.csrf_token, + expires_at=claims.expires_at, + ) + + @override + def consume_assertion_nonce(self, nonce: str) -> bool: + return consume_sso_assertion_nonce(self._redis, nonce) + + @override + def reserve_approval_nonce(self, nonce: str, reservation_id: str) -> bool: + if not nonce or not reservation_id: + return False + reserve_script = self._redis.register_script(_RESERVE_APPROVAL_NONCE_LUA) + return bool( + reserve_script( + keys=[NONCE_KEY_FMT.format(nonce=nonce)], + args=[reservation_id, NONCE_TTL_SECONDS], + ) + ) + + @override + def release_approval_nonce(self, nonce: str, reservation_id: str) -> None: + if not nonce or not reservation_id: + return + release_script = self._redis.register_script(_RELEASE_APPROVAL_NONCE_LUA) + release_script( + keys=[NONCE_KEY_FMT.format(nonce=nonce)], + args=[reservation_id], + ) + + +class RedisExternalApprovalLimiter(ExternalApprovalLimiter): + def __init__(self, *, redis: RedisClientWrapper) -> None: + self._rate_limiter = RateLimiter( + prefix="rl:subject_email", + max_attempts=LIMIT_APPROVE_EXT_PER_EMAIL.limit, + time_window=int(LIMIT_APPROVE_EXT_PER_EMAIL.window.total_seconds()), + redis_client=redis, + ) + + @override + def is_rate_limited(self, subject_email: str) -> bool: + return self._rate_limiter.is_rate_limited(f"subject:{subject_email}") + + @override + def record(self, subject_email: str) -> None: + self._rate_limiter.increment_rate_limit(f"subject:{subject_email}") diff --git a/api/services/oauth_device_application_service.py b/api/services/oauth_device_application_service.py new file mode 100644 index 00000000000..c77ba84e3fe --- /dev/null +++ b/api/services/oauth_device_application_service.py @@ -0,0 +1,760 @@ +"""Application service for OAuth device-flow use cases.""" + +from __future__ import annotations + +import logging +import secrets +from collections.abc import Sequence +from hashlib import sha256 +from typing import Protocol + +from machinery.context import RequestContext +from services.entities.account_entities import AccountSnapshot +from services.oauth_device_contracts import ( + ACCOUNT_ISSUER_SENTINEL, + DEFAULT_POLL_INTERVAL_SECONDS, + DEVICE_FLOW_TTL_SECONDS, + AccessDeniedError, + AlreadyResolvedError, + ApprovalInProgressError, + ApprovalOutcomeUnknownError, + ApprovalSessionConsumedError, + ApprovalTransitionConfirmation, + AuthorizationPendingError, + DeviceApprovalContext, + DeviceAuthorization, + DeviceFlowStateView, + DeviceFlowStatus, + DeviceLookup, + DeviceMutation, + DeviceRequestContext, + DeviceSSOCompletion, + DeviceSSOInitiation, + DeviceStateLostError, + DeviceWorkspace, + ExpiredOrUnknownError, + ExpiredTokenError, + ExternalApprovalCSRFError, + ExternalApprovalGrant, + ExternalApprovalRateLimitError, + ExternalIdentityConflictError, + ExternalSubjectAssertion, + ExternalUserCodeMismatchError, + ExternalUserCodeNotFoundError, + InvalidApprovalSessionError, + InvalidSSOAssertionError, + InvalidTransitionError, + InvalidUserCodeError, + IssuedOAuthToken, + OAuthDeviceSessionNotFoundError, + OAuthDeviceSessionPage, + OAuthDeviceSSOConfigurationError, + OAuthDeviceSSOInitiationError, + OAuthDeviceTokenRotation, + OAuthDeviceTokenWrite, + PollPayload, + PollTooFastError, + SlowDownDecision, + StateNotFoundError, + UnsupportedClientError, +) + +logger = logging.getLogger(__name__) + +_APPROVE_GUARD_TTL_SECONDS = 60 +_SSO_COMPLETE_PATH = "/openapi/v1/oauth/device/sso-complete" + + +class DeviceFlowStore(Protocol): + def start(self, client_id: str, device_label: str, created_ip: str) -> tuple[str, str, int]: ... + + def load_by_user_code(self, user_code: str) -> tuple[str, DeviceFlowStateView] | None: ... + + def load_by_device_code(self, device_code: str) -> DeviceFlowStateView | None: ... + + def approve( + self, + device_code: str, + transition_id: str, + token_id: str, + poll_payload: PollPayload, + ) -> None: ... + + def deny(self, device_code: str, transition_id: str) -> None: ... + + def confirm_approval( + self, + device_code: str, + transition_id: str, + token_id: str, + ) -> ApprovalTransitionConfirmation: ... + + def consume_on_poll(self, device_code: str) -> DeviceFlowStateView | None: ... + + def record_poll(self, device_code: str, interval_seconds: int) -> SlowDownDecision: ... + + def try_acquire_approval(self, guard_id: str, owner_id: str, ttl_seconds: int) -> bool: ... + + def release_approval(self, guard_id: str, owner_id: str) -> None: ... + + +class DeviceWorkspaceQuery(Protocol): + def list_for_device_flow(self, account_id: str) -> Sequence[DeviceWorkspace]: ... + + +class OAuthDeviceAccountQuery(Protocol): + def get(self, account_id: str) -> AccountSnapshot | None: ... + + def has_active_email(self, email: str) -> bool: ... + + +class OAuthDeviceTokenIssuer(Protocol): + def issue_account_token( + self, + *, + account: AccountSnapshot, + workspace_id: str, + client_id: str, + device_label: str, + ) -> IssuedOAuthToken: ... + + def issue_external_token( + self, + *, + subject_email: str, + subject_issuer: str, + client_id: str, + device_label: str, + ) -> IssuedOAuthToken: ... + + def rollback_token(self, token: IssuedOAuthToken) -> bool: ... + + +class OAuthDeviceSessionRepository(Protocol): + def list_account_sessions( + self, + *, + account_id: str, + page: int, + limit: int, + ) -> OAuthDeviceSessionPage: ... + + def revoke_account_session(self, *, account_id: str, token_id: str) -> bool: ... + + +class OAuthDeviceTokenPersistence(Protocol): + def rotate_token(self, token: OAuthDeviceTokenWrite) -> OAuthDeviceTokenRotation: ... + + def rollback_rotation(self, rotation: OAuthDeviceTokenRotation) -> bool: ... + + +class OAuthDeviceTokenTTLPolicy(Protocol): + def ttl_days(self, workspace_id: str | None) -> int: ... + + +class OAuthDeviceSSOGateway(Protocol): + def initiate(self, *, user_code: str, callback_url: str, ttl_seconds: int) -> str: ... + + def verify_assertion(self, assertion: str) -> ExternalSubjectAssertion: ... + + def mint_approval_grant( + self, + *, + issuer: str, + subject_email: str, + subject_issuer: str, + user_code: str, + ) -> str: ... + + def verify_approval_grant(self, token: str) -> ExternalApprovalGrant: ... + + def consume_assertion_nonce(self, nonce: str) -> bool: ... + + def reserve_approval_nonce(self, nonce: str, reservation_id: str) -> bool: ... + + def release_approval_nonce(self, nonce: str, reservation_id: str) -> None: ... + + +class ExternalApprovalLimiter(Protocol): + def is_rate_limited(self, subject_email: str) -> bool: ... + + def record(self, subject_email: str) -> None: ... + + +class OAuthDeviceSettings(Protocol): + @property + def known_client_ids(self) -> frozenset[str]: ... + + @property + def verification_base_url(self) -> str | None: ... + + @property + def sso_base_url(self) -> str | None: ... + + +class OAuthDeviceApplicationService: + def __init__( + self, + *, + store: DeviceFlowStore, + accounts: OAuthDeviceAccountQuery, + workspaces: DeviceWorkspaceQuery, + tokens: OAuthDeviceTokenIssuer, + sessions: OAuthDeviceSessionRepository, + sso: OAuthDeviceSSOGateway, + external_approval_limiter: ExternalApprovalLimiter, + settings: OAuthDeviceSettings, + ) -> None: + self._store = store + self._accounts = accounts + self._workspaces = workspaces + self._tokens = tokens + self._sessions = sessions + self._sso = sso + self._external_approval_limiter = external_approval_limiter + self._settings = settings + + def list_account_sessions(self, *, account_id: str, page: int, limit: int) -> OAuthDeviceSessionPage: + return self._sessions.list_account_sessions(account_id=account_id, page=page, limit=limit) + + def revoke_account_session(self, *, account_id: str, token_id: str) -> DeviceMutation: + if not self._sessions.revoke_account_session(account_id=account_id, token_id=token_id): + raise OAuthDeviceSessionNotFoundError(token_id) + return DeviceMutation(status="revoked") + + def start( + self, + *, + client_id: str, + device_label: str, + created_ip: str, + request_origin: str, + ) -> DeviceAuthorization: + if client_id not in self._settings.known_client_ids: + raise UnsupportedClientError + + device_code, user_code, expires_in = self._store.start(client_id, device_label, created_ip) + base_url = self._settings.verification_base_url or request_origin + return DeviceAuthorization( + device_code=device_code, + user_code=user_code, + verification_uri=f"{base_url.rstrip('/')}/device", + expires_in=expires_in, + interval=DEFAULT_POLL_INTERVAL_SECONDS, + ) + + def poll(self, *, device_code: str, poll_ip: str) -> PollPayload: + if self._store.record_poll(device_code, DEFAULT_POLL_INTERVAL_SECONDS) is SlowDownDecision.SLOW_DOWN: + raise PollTooFastError + + state = self._store.load_by_device_code(device_code) + if state is None: + raise ExpiredTokenError + if state.status is DeviceFlowStatus.PENDING: + raise AuthorizationPendingError + + terminal = self._store.consume_on_poll(device_code) + if terminal is None: + raise ExpiredTokenError + if terminal.status is DeviceFlowStatus.DENIED: + raise AccessDeniedError + + poll_payload = terminal.poll_payload + if poll_payload is None or "token" not in poll_payload: + logger.error("device_flow: approved state missing poll_payload for %s", device_code) + raise ExpiredTokenError + + self._audit_cross_ip_if_needed(state, poll_ip) + return poll_payload + + def lookup(self, *, user_code: str) -> DeviceLookup: + found = self._store.load_by_user_code(user_code.strip().upper()) + if found is None: + return DeviceLookup(valid=False, expires_in_remaining=0, client_id=None) + + _device_code, state = found + if state.status is not DeviceFlowStatus.PENDING: + return DeviceLookup(valid=False, expires_in_remaining=0, client_id=state.client_id) + return DeviceLookup(valid=True, expires_in_remaining=DEVICE_FLOW_TTL_SECONDS, client_id=state.client_id) + + def approve(self, context: RequestContext, *, user_code: str) -> DeviceMutation: + device_code, state = self._pending_state(user_code) + account = self._accounts.get(context.account_id) + if account is None or context.active_workspace_id is None: + raise DeviceStateLostError + workspaces = tuple(self._workspaces.list_for_device_flow(account.id)) + transition_id = secrets.token_urlsafe(16) + guard_id = self._rotation_guard_id( + subject_email=account.email, + subject_issuer=ACCOUNT_ISSUER_SENTINEL, + client_id=state.client_id, + device_label=state.device_label, + ) + if not self._store.try_acquire_approval(guard_id, transition_id, _APPROVE_GUARD_TTL_SECONDS): + raise ApprovalInProgressError + + try: + current = self._store.load_by_user_code(user_code.strip().upper()) + if current is None or current[0] != device_code: + raise DeviceStateLostError + state = current[1] + if state.status is not DeviceFlowStatus.PENDING: + raise AlreadyResolvedError + mint = self._tokens.issue_account_token( + account=account, + workspace_id=context.active_workspace_id, + client_id=state.client_id, + device_label=state.device_label, + ) + try: + poll_payload = self._build_account_poll_payload( + account, + context.active_workspace_id, + workspaces, + mint, + ) + self._publish_approval( + device_code, + transition_id=transition_id, + token_id=mint.token_id, + poll_payload=poll_payload, + ) + except ApprovalOutcomeUnknownError: + raise + except Exception as error: + self._rollback_issued_token(mint) + if isinstance(error, (StateNotFoundError, InvalidTransitionError)): + logger.exception("device_flow: approve raced on %s", device_code) + raise DeviceStateLostError from error + raise + finally: + self._store.release_approval(guard_id, transition_id) + + self._emit_approve_audit(state, account, context.active_workspace_id, mint) + return DeviceMutation(status="approved") + + def deny(self, *, user_code: str) -> DeviceMutation: + device_code, state = self._pending_state(user_code) + try: + self._store.deny(device_code, secrets.token_urlsafe(16)) + except (StateNotFoundError, InvalidTransitionError) as error: + logger.exception("device_flow: deny raced on %s", device_code) + raise DeviceStateLostError from error + + self._emit_deny_audit(state) + return DeviceMutation(status="denied") + + def initiate_sso(self, context: DeviceRequestContext, *, user_code: str) -> DeviceSSOInitiation: + normalized_user_code = user_code.strip().upper() + self._pending_external_state(normalized_user_code) + base_url = self._sso_origin() + redirect_url = self._sso.initiate( + user_code=normalized_user_code, + callback_url=f"{base_url}{_SSO_COMPLETE_PATH}", + ttl_seconds=DEVICE_FLOW_TTL_SECONDS, + ) + if not redirect_url: + raise OAuthDeviceSSOInitiationError + logger.info( + "oauth device SSO initiated", + extra={"request_id": context.request_id}, + ) + return DeviceSSOInitiation(redirect_url=redirect_url) + + def complete_sso( + self, + context: DeviceRequestContext, + *, + inbound_error: str | None, + inbound_user_code: str | None, + assertion: str | None, + ) -> DeviceSSOCompletion: + if inbound_error: + return DeviceSSOCompletion(error_code=inbound_error, user_code=inbound_user_code) + if not assertion: + return DeviceSSOCompletion(error_code="sso_failed") + + try: + claims = self._sso.verify_assertion(assertion) + except InvalidSSOAssertionError as error: + logger.warning( + "oauth device SSO assertion rejected: %s", + error, + extra={"request_id": context.request_id}, + ) + return DeviceSSOCompletion(error_code="sso_failed") + + user_code = claims.user_code.strip().upper() + if not self._sso.consume_assertion_nonce(claims.nonce): + return DeviceSSOCompletion(error_code="sso_failed", user_code=user_code) + + try: + _device_code, state = self._pending_external_state(user_code) + except InvalidUserCodeError: + return DeviceSSOCompletion(error_code="sso_failed", user_code=user_code) + + if self._accounts.has_active_email(claims.subject_email): + self._emit_external_rejection_audit( + context, + state, + claims.subject_email, + claims.subject_issuer, + reason="email_belongs_to_dify_account", + ) + return DeviceSSOCompletion( + error_code="email_belongs_to_dify_account", + user_code=user_code, + ) + + try: + approval_grant = self._sso.mint_approval_grant( + issuer=self._sso_origin(), + subject_email=claims.subject_email, + subject_issuer=claims.subject_issuer, + user_code=user_code, + ) + except OAuthDeviceSSOConfigurationError: + return DeviceSSOCompletion(error_code="sso_failed", user_code=user_code) + return DeviceSSOCompletion(user_code=user_code, approval_grant=approval_grant) + + def get_approval_context( + self, + _context: DeviceRequestContext, + *, + approval_grant: str, + ) -> DeviceApprovalContext: + if not approval_grant: + raise InvalidApprovalSessionError + claims = self._sso.verify_approval_grant(approval_grant) + return DeviceApprovalContext( + subject_email=claims.subject_email, + subject_issuer=claims.subject_issuer, + user_code=claims.user_code, + csrf_token=claims.csrf_token, + expires_at=claims.expires_at, + ) + + def approve_external( + self, + context: DeviceRequestContext, + *, + approval_grant: str, + csrf_token: str, + user_code: str, + ) -> DeviceMutation: + if not approval_grant: + raise InvalidApprovalSessionError + claims = self._sso.verify_approval_grant(approval_grant) + + if self._external_approval_limiter.is_rate_limited(claims.subject_email): + raise ExternalApprovalRateLimitError + self._external_approval_limiter.record(claims.subject_email) + + if not csrf_token or not secrets.compare_digest(csrf_token, claims.csrf_token): + raise ExternalApprovalCSRFError + if user_code.strip().upper() != claims.user_code: + raise ExternalUserCodeMismatchError + + found = self._store.load_by_user_code(claims.user_code) + if found is None: + raise ExternalUserCodeNotFoundError + device_code, state = found + if state.status is not DeviceFlowStatus.PENDING: + raise AlreadyResolvedError + + transition_id = secrets.token_urlsafe(16) + guard_id = self._rotation_guard_id( + subject_email=claims.subject_email, + subject_issuer=claims.subject_issuer, + client_id=state.client_id, + device_label=state.device_label, + ) + if not self._store.try_acquire_approval(guard_id, transition_id, _APPROVE_GUARD_TTL_SECONDS): + raise ApprovalInProgressError + + locked_device_code = device_code + reservation_id = sha256(f"oauth-device-approval\0{claims.nonce}".encode()).hexdigest() + nonce_reserved = False + mint: IssuedOAuthToken | None = None + try: + current = self._store.load_by_user_code(claims.user_code) + if current is None: + raise ExternalUserCodeNotFoundError + device_code, state = current + if device_code != locked_device_code: + raise DeviceStateLostError + if state.status is not DeviceFlowStatus.PENDING: + raise AlreadyResolvedError + + if self._accounts.has_active_email(claims.subject_email): + self._emit_external_rejection_audit( + context, + state, + claims.subject_email, + claims.subject_issuer, + reason="email_belongs_to_dify_account", + ) + raise ExternalIdentityConflictError + if not self._sso.reserve_approval_nonce(claims.nonce, reservation_id): + raise ApprovalSessionConsumedError + nonce_reserved = True + + mint = self._tokens.issue_external_token( + subject_email=claims.subject_email, + subject_issuer=claims.subject_issuer, + client_id=state.client_id, + device_label=state.device_label, + ) + poll_payload: PollPayload = { + "token": mint.token, + "expires_at": mint.expires_at, + "subject_type": "external_sso", + "subject_email": claims.subject_email, + "subject_issuer": claims.subject_issuer, + "account": None, + "workspaces": [], + "default_workspace_id": None, + "token_id": mint.token_id, + } + self._publish_approval( + device_code, + transition_id=transition_id, + token_id=mint.token_id, + poll_payload=poll_payload, + ) + except ApprovalOutcomeUnknownError: + raise + except Exception as error: + if mint is not None: + self._rollback_issued_token(mint) + if nonce_reserved: + self._release_approval_nonce(claims.nonce, reservation_id) + if isinstance(error, (StateNotFoundError, InvalidTransitionError)): + logger.exception("approve-external: state transition raced") + raise DeviceStateLostError from error + raise + finally: + self._store.release_approval(guard_id, transition_id) + + assert mint is not None + self._emit_approve_external_audit(context, state, claims, mint) + return DeviceMutation(status="approved") + + def _sso_origin(self) -> str: + base_url = (self._settings.sso_base_url or "").rstrip("/") + if not base_url: + raise OAuthDeviceSSOConfigurationError + return base_url + + def _pending_external_state(self, user_code: str) -> tuple[str, DeviceFlowStateView]: + found = self._store.load_by_user_code(user_code.strip().upper()) + if found is None or found[1].status is not DeviceFlowStatus.PENDING: + raise InvalidUserCodeError + return found + + def _pending_state(self, user_code: str) -> tuple[str, DeviceFlowStateView]: + found = self._store.load_by_user_code(user_code.strip().upper()) + if found is None: + raise ExpiredOrUnknownError + device_code, state = found + if state.status is not DeviceFlowStatus.PENDING: + raise AlreadyResolvedError + return device_code, state + + def _build_account_poll_payload( + self, + account: AccountSnapshot, + active_workspace_id: str, + workspaces: Sequence[DeviceWorkspace], + mint: IssuedOAuthToken, + ) -> PollPayload: + default_workspace_id = next( + (workspace.id for workspace in workspaces if workspace.id == active_workspace_id), + None, + ) + if default_workspace_id is None: + 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 { + "token": mint.token, + "expires_at": mint.expires_at, + "subject_type": "account", + "account": {"id": account.id, "email": account.email, "name": account.name}, + "workspaces": [ + {"id": workspace.id, "name": workspace.name, "role": workspace.role} for workspace in workspaces + ], + "default_workspace_id": default_workspace_id, + "token_id": mint.token_id, + } + + def _rollback_issued_token(self, mint: IssuedOAuthToken) -> None: + try: + if not self._tokens.rollback_token(mint): + logger.warning("OAuth token compensation skipped because rotation was already superseded") + except Exception: + logger.exception("OAuth token compensation failed", extra={"token_id": mint.token_id}) + + def _publish_approval( + self, + device_code: str, + *, + transition_id: str, + token_id: str, + poll_payload: PollPayload, + ) -> None: + try: + self._store.approve( + device_code, + transition_id=transition_id, + token_id=token_id, + poll_payload=poll_payload, + ) + except (StateNotFoundError, InvalidTransitionError): + raise + except Exception as publish_error: + try: + confirmation = self._store.confirm_approval(device_code, transition_id, token_id) + except Exception: + logger.exception("OAuth approval transition could not be confirmed", extra={"token_id": token_id}) + raise ApprovalOutcomeUnknownError from publish_error + if confirmation is ApprovalTransitionConfirmation.PUBLISHED: + logger.warning("OAuth approval transition succeeded after an ambiguous Redis response") + return + if confirmation is ApprovalTransitionConfirmation.UNKNOWN: + raise ApprovalOutcomeUnknownError from publish_error + raise + + @staticmethod + def _rotation_guard_id( + *, + subject_email: str, + subject_issuer: str, + client_id: str, + device_label: str, + ) -> str: + identity = "\0".join((subject_email, subject_issuer, client_id, device_label)) + return sha256(identity.encode()).hexdigest() + + def _release_approval_nonce(self, nonce: str, reservation_id: str) -> None: + try: + self._sso.release_approval_nonce(nonce, reservation_id) + except Exception: + logger.exception("OAuth approval nonce compensation failed") + + @staticmethod + def _audit_cross_ip_if_needed(state: DeviceFlowStateView, poll_ip: str) -> None: + if state.created_ip and poll_ip and poll_ip != state.created_ip: + token_id = state.poll_payload["token_id"] if state.poll_payload is not None else state.token_id + logger.warning( + "audit: oauth.device_code_cross_ip_poll token_id=%s creation_ip=%s poll_ip=%s", + token_id, + state.created_ip, + poll_ip, + extra={ + "audit": True, + "token_id": token_id, + "creation_ip": state.created_ip, + "poll_ip": poll_ip, + }, + ) + + @staticmethod + def _emit_approve_audit( + state: DeviceFlowStateView, + account: AccountSnapshot, + workspace_id: str, + mint: IssuedOAuthToken, + ) -> None: + logger.warning( + "audit: oauth.device_flow_approved token_id=%s subject=%s client_id=%s device_label=%s " + "rotated=? expires_at=%s", + mint.token_id, + account.email, + state.client_id, + state.device_label, + mint.expires_at, + extra={ + "audit": True, + "event": "oauth.device_flow_approved", + "token_id": mint.token_id, + "subject_type": "account", + "subject_email": account.email, + "account_id": account.id, + "tenant_id": workspace_id, + "client_id": state.client_id, + "device_label": state.device_label, + "scopes": ["full"], + "expires_at": mint.expires_at, + }, + ) + + @staticmethod + def _emit_deny_audit(state: DeviceFlowStateView) -> None: + logger.warning( + "audit: oauth.device_flow_denied client_id=%s device_label=%s", + state.client_id, + state.device_label, + extra={ + "audit": True, + "event": "oauth.device_flow_denied", + "client_id": state.client_id, + "device_label": state.device_label, + }, + ) + + @staticmethod + def _emit_external_rejection_audit( + context: DeviceRequestContext, + state: DeviceFlowStateView, + subject_email: str, + subject_issuer: str, + *, + reason: str, + ) -> None: + logger.warning( + "audit: oauth.device_flow_rejected subject_type=external_sso subject_email=%s subject_issuer=%s reason=%s", + subject_email, + subject_issuer, + reason, + extra={ + "audit": True, + "event": "oauth.device_flow_rejected", + "request_id": context.request_id, + "subject_type": "external_sso", + "subject_email": subject_email, + "subject_issuer": subject_issuer, + "reason": reason, + "client_id": state.client_id, + "device_label": state.device_label, + }, + ) + + @staticmethod + def _emit_approve_external_audit( + context: DeviceRequestContext, + state: DeviceFlowStateView, + claims: ExternalApprovalGrant, + mint: IssuedOAuthToken, + ) -> None: + logger.warning( + "audit: oauth.device_flow_approved subject_type=external_sso subject_email=%s " + "subject_issuer=%s token_id=%s", + claims.subject_email, + claims.subject_issuer, + mint.token_id, + extra={ + "audit": True, + "event": "oauth.device_flow_approved", + "request_id": context.request_id, + "subject_type": "external_sso", + "subject_email": claims.subject_email, + "subject_issuer": claims.subject_issuer, + "token_id": mint.token_id, + "client_id": state.client_id, + "device_label": state.device_label, + "scopes": ["apps:run"], + "expires_at": mint.expires_at, + }, + ) diff --git a/api/services/oauth_device_contracts.py b/api/services/oauth_device_contracts.py new file mode 100644 index 00000000000..3f3ca2683a9 --- /dev/null +++ b/api/services/oauth_device_contracts.py @@ -0,0 +1,276 @@ +"""Framework-neutral contracts shared by the OAuth device-flow use cases.""" + +from dataclasses import dataclass +from datetime import datetime +from enum import StrEnum +from typing import NotRequired, Protocol, TypedDict + +DEVICE_FLOW_TTL_SECONDS = 15 * 60 +DEFAULT_POLL_INTERVAL_SECONDS = 5 + +# PostgreSQL partial unique indexes treat NULLs as distinct. Account tokens use +# a non-null issuer so the active-token uniqueness key remains stable. +ACCOUNT_ISSUER_SENTINEL = "dify:account" + + +class DeviceFlowStatus(StrEnum): + PENDING = "pending" + APPROVED = "approved" + DENIED = "denied" + + +class SlowDownDecision(StrEnum): + OK = "ok" + SLOW_DOWN = "slow_down" + + +class ApprovalTransitionConfirmation(StrEnum): + PUBLISHED = "published" + NOT_PUBLISHED = "not_published" + UNKNOWN = "unknown" + + +class PollPayload(TypedDict): + """Payload returned once by the unauthenticated token poll endpoint.""" + + token: str + expires_at: str + subject_type: str + account: dict[str, object] | None + workspaces: list[dict[str, object]] + default_workspace_id: str | None + token_id: str + subject_email: NotRequired[str] + subject_issuer: NotRequired[str] + + +@dataclass(frozen=True, slots=True) +class DeviceAuthorization: + device_code: str + user_code: str + verification_uri: str + expires_in: int + interval: int + + +@dataclass(frozen=True, slots=True) +class DeviceLookup: + valid: bool + expires_in_remaining: int + client_id: str | None + + +@dataclass(frozen=True, slots=True) +class DeviceMutation: + status: str + + +@dataclass(frozen=True, slots=True) +class DeviceWorkspace: + id: str + name: str + role: str + current: bool + + +@dataclass(frozen=True, slots=True) +class OAuthDeviceTokenRotation: + token_id: str + replaced_token_id: str | None + replaced_token_hash: str | None + + +@dataclass(frozen=True, slots=True) +class IssuedOAuthToken: + token: str + expires_at: str + rotation: OAuthDeviceTokenRotation + + @property + def token_id(self) -> str: + return self.rotation.token_id + + +@dataclass(frozen=True, slots=True) +class OAuthDeviceTokenWrite: + subject_email: str + subject_issuer: str + account_id: str | None + client_id: str + device_label: str + prefix: str + token_hash: str + expires_at: datetime + + +@dataclass(frozen=True, slots=True) +class OAuthDeviceSession: + 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 OAuthDeviceSessionPage: + page: int + limit: int + total: int + items: tuple[OAuthDeviceSession, ...] + + +@dataclass(frozen=True, slots=True) +class DeviceRequestContext: + request_id: str + trace_id: str | None + + +@dataclass(frozen=True, slots=True) +class ExternalSubjectAssertion: + subject_email: str + subject_issuer: str + user_code: str + nonce: str + + +@dataclass(frozen=True, slots=True) +class ExternalApprovalGrant: + subject_email: str + subject_issuer: str + user_code: str + nonce: str + csrf_token: str + expires_at: datetime + + +@dataclass(frozen=True, slots=True) +class DeviceSSOInitiation: + redirect_url: str + + +@dataclass(frozen=True, slots=True) +class DeviceSSOCompletion: + error_code: str | None = None + user_code: str | None = None + approval_grant: str | None = None + + +@dataclass(frozen=True, slots=True) +class DeviceApprovalContext: + subject_email: str + subject_issuer: str + user_code: str + csrf_token: str + expires_at: datetime + + +class DeviceFlowStateView(Protocol): + client_id: str + device_label: str + status: DeviceFlowStatus + token_id: str | None + created_ip: str + poll_payload: PollPayload | None + + +class OAuthDeviceError(Exception): + """Base use-case error translated to HTTP by the transport adapter.""" + + +class UnsupportedClientError(OAuthDeviceError): + pass + + +class PollTooFastError(OAuthDeviceError): + pass + + +class ExpiredTokenError(OAuthDeviceError): + pass + + +class AuthorizationPendingError(OAuthDeviceError): + pass + + +class AccessDeniedError(OAuthDeviceError): + pass + + +class ExpiredOrUnknownError(OAuthDeviceError): + pass + + +class AlreadyResolvedError(OAuthDeviceError): + pass + + +class ApprovalInProgressError(OAuthDeviceError): + pass + + +class ApprovalOutcomeUnknownError(OAuthDeviceError): + """The Redis transition may have committed and must not be compensated.""" + + +class DeviceStateLostError(OAuthDeviceError): + pass + + +class OAuthDeviceSessionNotFoundError(OAuthDeviceError): + pass + + +class InvalidUserCodeError(OAuthDeviceError): + pass + + +class OAuthDeviceSSOConfigurationError(OAuthDeviceError): + pass + + +class OAuthDeviceSSOInitiationError(OAuthDeviceError): + pass + + +class InvalidSSOAssertionError(OAuthDeviceError): + pass + + +class InvalidApprovalSessionError(OAuthDeviceError): + pass + + +class ExternalApprovalCSRFError(OAuthDeviceError): + pass + + +class ExternalApprovalRateLimitError(OAuthDeviceError): + pass + + +class ExternalUserCodeMismatchError(OAuthDeviceError): + pass + + +class ExternalUserCodeNotFoundError(OAuthDeviceError): + pass + + +class ExternalIdentityConflictError(OAuthDeviceError): + pass + + +class ApprovalSessionConsumedError(OAuthDeviceError): + pass + + +class StateNotFoundError(Exception): + pass + + +class InvalidTransitionError(Exception): + pass diff --git a/api/services/oauth_device_flow.py b/api/services/oauth_device_flow.py index 2704fade27d..8371b2dfdcb 100644 --- a/api/services/oauth_device_flow.py +++ b/api/services/oauth_device_flow.py @@ -1,22 +1,31 @@ +"""Redis state machine for OAuth device flow. + +Deployment contract: the CAS writer must not run alongside the pre-CAS +GET/SETEX writer. Old approve and deny operations can overwrite a terminal +state from this implementation because they do not participate in the same +transition protocol. Drain API workers running the old writer before deploying +this version. The JSON representation remains reader-compatible so pending +device codes can survive that restart. +""" + from __future__ import annotations -import hashlib import json import logging -import os import secrets import time -import uuid from dataclasses import asdict, dataclass, field -from datetime import UTC, datetime, timedelta -from enum import StrEnum -from typing import NotRequired, TypedDict +from datetime import UTC, datetime -from sqlalchemy import func, select, update -from sqlalchemy.orm import Session - -from libs.oauth_bearer import SubjectType, TokenType, invalidate_oauth_token_cache -from models.oauth import OAuthAccessToken +from services.oauth_device_contracts import ( + DEVICE_FLOW_TTL_SECONDS, + ApprovalTransitionConfirmation, + DeviceFlowStatus, + InvalidTransitionError, + PollPayload, + SlowDownDecision, + StateNotFoundError, +) logger = logging.getLogger(__name__) @@ -31,87 +40,83 @@ _USER_CODE_KEY_PREFIX = "user_code:" DEVICE_CODE_KEY_FMT = _DEVICE_CODE_KEY_PREFIX + "{code}" USER_CODE_KEY_FMT = _USER_CODE_KEY_PREFIX + "{code}" -# Atomic GET → status-check → DEL(both keys). Two concurrent pollers must +# Atomic GET → status-check → DEL(device key). Two concurrent pollers must # not both observe APPROVED — only the winner gets the plaintext token, -# the loser sees nil and the caller maps that to expired_token. +# the loser sees nil and the caller maps that to expired_token. The user-code +# mapping is cleaned up separately so Redis Cluster only sees one script key. _CONSUME_ON_POLL_LUA = """ local raw = redis.call('GET', KEYS[1]) if not raw then return nil end local ok, decoded = pcall(cjson.decode, raw) if not ok then return nil end if decoded.status == 'pending' then return nil end -if decoded.user_code then - redis.call('DEL', ARGV[1] .. decoded.user_code) -end redis.call('DEL', KEYS[1]) return raw """ -DEVICE_FLOW_TTL_SECONDS = 15 * 60 # RFC 8628 expires_in +# The transition ID makes an approved write idempotent and lets callers +# distinguish a committed write from a connection failure after SETEX. +_TRANSITION_LUA = """ +local raw = redis.call('GET', KEYS[1]) +if not raw then return 0 end +local ok, decoded = pcall(cjson.decode, raw) +if not ok then return -2 end + +local target = ARGV[1] +local transition_marker = ARGV[2] +if decoded.status ~= 'pending' then + if decoded.status == target and decoded.token_id == transition_marker then + return 2 + end + return -1 +end + +decoded.status = target +decoded.token_id = transition_marker +decoded.poll_payload = nil +if target == 'approved' then + local payload_ok = pcall(cjson.decode, ARGV[3]) + if not payload_ok then return -2 end +end + +-- cjson loses the distinction between [] and {} when decoding empty tables. +-- Preserve the validated payload JSON verbatim while encoding the state fields. +local encoded = cjson.encode(decoded) +if target == 'approved' then + encoded = string.sub(encoded, 1, -2) .. ',"poll_payload":' .. ARGV[3] .. '}' +end + +local ttl = redis.call('TTL', KEYS[1]) +local floor = tonumber(ARGV[4]) +if ttl < floor then ttl = floor end +redis.call('SETEX', KEYS[1], ttl, encoded) +return 1 +""" + +_RELEASE_GUARD_LUA = """ +if redis.call('GET', KEYS[1]) == ARGV[1] then + return redis.call('DEL', KEYS[1]) +end +return 0 +""" + APPROVED_TTL_SECONDS_MIN = 60 # plaintext-token lifetime floor USER_CODE_ALPHABET = "ABCDEFGHJKLMNPQRSTUVWXY3456789" # ambiguous chars dropped USER_CODE_SEGMENT_LEN = 4 USER_CODE_MAX_CLAIM_ATTEMPTS = 5 - -DEFAULT_POLL_INTERVAL_SECONDS = 5 # RFC 8628 minimum - - -class DeviceFlowStatus(StrEnum): - PENDING = "pending" - APPROVED = "approved" - DENIED = "denied" - - -class SlowDownDecision(StrEnum): - OK = "ok" - SLOW_DOWN = "slow_down" - - -class PollPayload(TypedDict): - """Body served by the unauthenticated poll endpoint - (`POST /openapi/v1/oauth/device/token`) once approve has run. - - A single shape across both branches so the CLI/SPA can parse one - contract: - - - ``account`` branch (built in `controllers.openapi.oauth_device. - _build_account_poll_payload`) populates ``account`` + ``workspaces`` - + ``default_workspace_id`` and omits the SSO-only fields. - - ``external_sso`` branch (built in - `controllers.openapi.oauth_device_sso.approve_external`) populates - ``subject_email`` + ``subject_issuer`` and zero-fills the - account/workspace fields (``None`` / ``[]``). - - Pre-rendering here means the unauthenticated poll handler doesn't - re-query accounts/tenants for authz data. - """ - - token: str - expires_at: str - subject_type: SubjectType - account: dict[str, object] | None - workspaces: list[dict[str, object]] - default_workspace_id: str | None - token_id: str - subject_email: NotRequired[str] - subject_issuer: NotRequired[str] +_TRANSITION_TOKEN_PREFIX = "transition:" +_LEGACY_APPROVAL_FIELDS = ("subject_email", "account_id", "subject_issuer", "minted_token", "transition_id") @dataclass class DeviceFlowState: - """``minted_token`` is plaintext between approve and the next poll; - DEL'd after the poll reads it. - """ + """Ephemeral device state; approved payload is consumed exactly once.""" user_code: str client_id: str device_label: str status: DeviceFlowStatus - subject_email: str | None = None - account_id: str | None = None - subject_issuer: str | None = None - minted_token: str | None = None token_id: str | None = None created_at: str = "" created_ip: str = "" @@ -124,6 +129,10 @@ class DeviceFlowState: @classmethod def from_json(cls, raw: str) -> DeviceFlowState: data = json.loads(raw) + # Device codes live for at most 15 minutes, but tolerate state written + # by the previous schema during a rolling deployment. + for field_name in _LEGACY_APPROVAL_FIELDS: + data.pop(field_name, None) if "status" in data: data["status"] = DeviceFlowStatus(data["status"]) return cls(**data) @@ -141,14 +150,6 @@ def _random_user_code() -> str: return f"{_random_user_code_segment()}-{_random_user_code_segment()}" -class StateNotFoundError(Exception): - pass - - -class InvalidTransitionError(Exception): - pass - - class UserCodeExhaustedError(Exception): pass @@ -157,6 +158,8 @@ class DeviceFlowRedis: def __init__(self, redis_client) -> None: self._redis = redis_client self._consume_on_poll_script = redis_client.register_script(_CONSUME_ON_POLL_LUA) + self._transition_script = redis_client.register_script(_TRANSITION_LUA) + self._release_guard_script = redis_client.register_script(_RELEASE_GUARD_LUA) def start(self, client_id: str, device_label: str, created_ip: str) -> tuple[str, str, int]: device_code = _random_device_code() @@ -212,43 +215,44 @@ class DeviceFlowRedis: def approve( self, device_code: str, - subject_email: str, - account_id: str | None, - minted_token: str, + transition_id: str, token_id: str, - subject_issuer: str | None = None, - poll_payload: PollPayload | None = None, + poll_payload: PollPayload, ) -> None: - state = self._load_state(device_code) - if state is None: - raise StateNotFoundError(device_code) - if state.status is not DeviceFlowStatus.PENDING: - raise InvalidTransitionError(f"cannot approve {state.status}") - - state.status = DeviceFlowStatus.APPROVED - state.subject_email = subject_email - state.account_id = account_id - state.subject_issuer = subject_issuer - state.minted_token = minted_token - state.token_id = token_id - state.poll_payload = poll_payload - - new_ttl = self._remaining_ttl(device_code, floor=APPROVED_TTL_SECONDS_MIN) - self._redis.setex(DEVICE_CODE_KEY_FMT.format(code=device_code), new_ttl, state.to_json()) - - def deny(self, device_code: str) -> None: - state = self._load_state(device_code) - if state is None: - raise StateNotFoundError(device_code) - if state.status is not DeviceFlowStatus.PENDING: - raise InvalidTransitionError(f"cannot deny {state.status}") - state.status = DeviceFlowStatus.DENIED - self._redis.setex( - DEVICE_CODE_KEY_FMT.format(code=device_code), - self._remaining_ttl(device_code, floor=1), - state.to_json(), + self._transition( + device_code=device_code, + target=DeviceFlowStatus.APPROVED, + transition_id=transition_id, + token_id=token_id, + poll_payload=poll_payload, + ttl_floor=APPROVED_TTL_SECONDS_MIN, ) + def deny(self, device_code: str, transition_id: str) -> None: + self._transition( + device_code=device_code, + target=DeviceFlowStatus.DENIED, + transition_id=transition_id, + token_id="", + poll_payload=None, + ttl_floor=1, + ) + + def confirm_approval( + self, + device_code: str, + transition_id: str, + token_id: str, + ) -> ApprovalTransitionConfirmation: + state = self._load_state(device_code) + if state is None: + return ApprovalTransitionConfirmation.UNKNOWN + if state.status is DeviceFlowStatus.APPROVED and state.token_id == self._transition_marker( + transition_id, token_id + ): + return ApprovalTransitionConfirmation.PUBLISHED + return ApprovalTransitionConfirmation.NOT_PUBLISHED + def consume_on_poll(self, device_code: str) -> DeviceFlowState | None: """Race-safe via Lua EVAL: GET + status-check + DEL execute in a single Redis transaction so only one of N concurrent pollers @@ -257,16 +261,20 @@ class DeviceFlowRedis: """ raw = self._consume_on_poll_script( keys=[DEVICE_CODE_KEY_FMT.format(code=device_code)], - args=[_USER_CODE_KEY_PREFIX], ) if raw is None: return None text_ = raw.decode() if isinstance(raw, (bytes, bytearray)) else raw try: - return DeviceFlowState.from_json(text_) + state = DeviceFlowState.from_json(text_) except (ValueError, KeyError): logger.exception("device_flow: corrupt state on consume %s", device_code) return None + try: + self._redis.delete(USER_CODE_KEY_FMT.format(code=state.user_code)) + except Exception: + logger.exception("device_flow: failed to clean consumed user-code mapping") + return state def record_poll(self, device_code: str, interval_seconds: int) -> SlowDownDecision: now = time.time() @@ -284,201 +292,45 @@ class DeviceFlowRedis: return SlowDownDecision.SLOW_DOWN return SlowDownDecision.OK - def _remaining_ttl(self, device_code: str, floor: int) -> int: - """``max(remaining, floor)`` — guarantees the CLI has at least - ``floor`` seconds to poll after a near-expiry approve. - """ - ttl = self._redis.ttl(DEVICE_CODE_KEY_FMT.format(code=device_code)) - if ttl is None or ttl < 0: - return floor - return max(int(ttl), floor) + def try_acquire_approval(self, guard_id: str, owner_id: str, ttl_seconds: int) -> bool: + return bool(self._redis.set(self._approval_guard_key(guard_id), owner_id, nx=True, ex=ttl_seconds)) - -# ============================================================================ -# Token mint — generate + upsert -# ============================================================================ - - -OAUTH_BODY_BYTES = 32 # ~256 bits entropy - -# Sentinel issuer for account-flow rows. Postgres' default partial unique -# index treats NULLs as distinct, which would let two live `dfoa_` rows -# share (email, client, device) and break rotate-in-place. Storing a -# non-empty literal makes the composite key collide as intended. -ACCOUNT_ISSUER_SENTINEL = "dify:account" - - -@dataclass(frozen=True, slots=True) -class MintResult: - """Plaintext token surfaces to the caller once.""" - - token: str - token_id: uuid.UUID - expires_at: datetime - - -@dataclass(frozen=True, slots=True) -class UpsertOutcome: - token_id: uuid.UUID - rotated: bool - old_hash: str | None - - -def generate_token(prefix: str) -> str: - return prefix + secrets.token_urlsafe(OAUTH_BODY_BYTES) - - -def sha256_hex(token: str) -> str: - return hashlib.sha256(token.encode("utf-8")).hexdigest() - - -def mint_oauth_token( - redis_client, - *, - subject_email: str, - subject_issuer: str | None, - account_id: str | None, - client_id: str, - device_label: str, - token_type: TokenType, - ttl_days: int, - session: Session, -) -> MintResult: - """Live row rotates in place via partial unique index - ``uq_oauth_active_per_device``; hard-expired rows are excluded by the - index predicate so re-login INSERTs fresh. Pre-rotate Redis entry is - deleted so stale AuthContext drops immediately. - """ - if token_type.subject.bound_to_account: - # Account-bound flows always write the sentinel — caller may pass None - # (for clarity) or the sentinel itself; nothing else is valid. - if subject_issuer not in (None, ACCOUNT_ISSUER_SENTINEL): - raise ValueError(f"account-flow token must use ACCOUNT_ISSUER_SENTINEL, got {subject_issuer!r}") - subject_issuer = ACCOUNT_ISSUER_SENTINEL - elif not subject_issuer or not subject_issuer.strip(): - # Defense in depth: enterprise canonicalises + rejects empty, - # but a regression there must not yield a NULL composite key here. - raise ValueError("external-SSO token requires non-empty subject_issuer") - - token = generate_token(token_type.prefix) - new_hash = sha256_hex(token) - expires_at = datetime.now(UTC) + timedelta(days=ttl_days) - - outcome = _upsert( - session, - subject_email=subject_email, - subject_issuer=subject_issuer, - account_id=account_id, - client_id=client_id, - device_label=device_label, - prefix=token_type.prefix, - new_hash=new_hash, - expires_at=expires_at, - ) - - if outcome.rotated and 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) - - -def _upsert( - session: Session, - *, - subject_email: str, - subject_issuer: str | None, - account_id: str | None, - client_id: str, - device_label: str, - prefix: str, - new_hash: str, - expires_at: datetime, -) -> UpsertOutcome: - # Snapshot prior live row's hash for Redis invalidation post-rotate. - # subject_issuer is always non-null here (account flow uses sentinel, - # external-SSO is validated upstream), so equality matches the index. - # FOR UPDATE locks the row so that concurrent logins for the same - # (subject, client, device) serialize here rather than both reading - # the same prior and producing two active tokens (TOCTOU race). - prior = session.execute( - select(OAuthAccessToken.id, OAuthAccessToken.token_hash) - .where( - OAuthAccessToken.subject_email == subject_email, - OAuthAccessToken.subject_issuer == subject_issuer, - OAuthAccessToken.client_id == client_id, - OAuthAccessToken.device_label == device_label, - OAuthAccessToken.revoked_at.is_(None), + def release_approval(self, guard_id: str, owner_id: str) -> None: + self._release_guard_script( + keys=[self._approval_guard_key(guard_id)], + args=[owner_id], ) - .limit(1) - .with_for_update() - ).first() - old_hash = prior.token_hash if prior else None - # Revoke any existing active token for this (subject, client, device) combination. - # PostgreSQL's ON CONFLICT doesn't support partial unique indexes (those with WHERE clauses), - # so we use a manual revoke-then-insert pattern instead. - if prior: - session.execute(update(OAuthAccessToken).where(OAuthAccessToken.id == prior.id).values(revoked_at=func.now())) - - # Insert the new token. - new_token = OAuthAccessToken( - subject_email=subject_email, - subject_issuer=subject_issuer, - account_id=account_id, - client_id=client_id, - device_label=device_label, - prefix=prefix, - token_hash=new_hash, - expires_at=expires_at, - ) - session.add(new_token) - session.flush() - - token_id = new_token.id - session.commit() - - return UpsertOutcome( - token_id=uuid.UUID(str(token_id)), - rotated=prior is not None, - old_hash=old_hash, - ) - - -# ============================================================================ -# TTL policy — days new OAuth tokens live -# ============================================================================ - - -DEFAULT_OAUTH_TTL_DAYS = 14 -MIN_TTL_DAYS = 1 -MAX_TTL_DAYS = 365 - -_TTL_ENV_VAR = "OAUTH_TTL_DAYS" - - -def oauth_ttl_days(tenant_id: str | None = None) -> int: - """``OAUTH_TTL_DAYS`` env, else default. EE tenant-level lookup - is deferred; when it lands it wins over the env (Redis-cached 60s). - """ - _ = tenant_id - - raw = os.environ.get(_TTL_ENV_VAR) - if raw is None: - return DEFAULT_OAUTH_TTL_DAYS - try: - value = int(raw) - except ValueError: - logger.warning( - "%s=%r is not an int; falling back to %d", - _TTL_ENV_VAR, - raw, - DEFAULT_OAUTH_TTL_DAYS, + def _transition( + self, + *, + device_code: str, + target: DeviceFlowStatus, + transition_id: str, + token_id: str, + poll_payload: PollPayload | None, + ttl_floor: int, + ) -> None: + result = int( + self._transition_script( + keys=[DEVICE_CODE_KEY_FMT.format(code=device_code)], + args=[ + target.value, + self._transition_marker(transition_id, token_id), + json.dumps(poll_payload) if poll_payload is not None else "", + ttl_floor, + ], + ) ) - return DEFAULT_OAUTH_TTL_DAYS - if value < MIN_TTL_DAYS: - logger.warning("%s=%d below min %d; clamping", _TTL_ENV_VAR, value, MIN_TTL_DAYS) - return MIN_TTL_DAYS - if value > MAX_TTL_DAYS: - logger.warning("%s=%d above max %d; clamping", _TTL_ENV_VAR, value, MAX_TTL_DAYS) - return MAX_TTL_DAYS - return value + if result in (0, -2): + raise StateNotFoundError(device_code) + if result == -1: + raise InvalidTransitionError(f"cannot transition device flow to {target.value}") + + @staticmethod + def _approval_guard_key(guard_id: str) -> str: + return f"oauth_device:approval_guard:{guard_id}" + + @staticmethod + def _transition_marker(transition_id: str, token_id: str) -> str: + return f"{_TRANSITION_TOKEN_PREFIX}{transition_id}:{token_id}" diff --git a/api/tests/integration_tests/controllers/openapi/conftest.py b/api/tests/integration_tests/controllers/openapi/conftest.py index 86e1e30b5fc..6228db125f1 100644 --- a/api/tests/integration_tests/controllers/openapi/conftest.py +++ b/api/tests/integration_tests/controllers/openapi/conftest.py @@ -10,11 +10,11 @@ from datetime import UTC, datetime, timedelta import pytest from flask import Flask +from constants.oauth_bearer import TokenType from controllers.openapi._catalog import CATALOG_HEADER, catalog_for from enums import DeploymentEdition from extensions.ext_database import db from extensions.ext_redis import redis_client -from libs.oauth_bearer import TokenType from models import Account, App, OAuthAccessToken, Tenant, TenantAccountJoin from models.account import AccountStatus 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 91612174a21..69bf01341c6 100644 --- a/api/tests/test_containers_integration_tests/controllers/openapi/conftest.py +++ b/api/tests/test_containers_integration_tests/controllers/openapi/conftest.py @@ -9,11 +9,12 @@ from faker import Faker from flask import Flask from sqlalchemy.orm import Session +from constants.oauth_bearer import TokenType from controllers.openapi.auth.context import Context from controllers.openapi.auth.loaders import PathParam, load_app from controllers.openapi.auth.requirements import ResolveCaller from controllers.openapi.auth.subjects import subject_from_auth -from libs.oauth_bearer import AuthContext, TokenType +from libs.oauth_bearer import AuthContext from models import Account, Tenant from services.account_service import AccountService, TenantService from tests.test_containers_integration_tests.helpers import generate_valid_password 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 bed322888a3..c7284c2898a 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 @@ -1,44 +1,53 @@ from __future__ import annotations from collections.abc import Callable +from dataclasses import dataclass +from datetime import UTC, datetime, timedelta +from uuid import UUID, uuid4 import pytest from flask import Flask from sqlalchemy.orm import Session from werkzeug.exceptions import NotFound +from constants.oauth_bearer import TokenType from controllers.openapi._models import SessionListQuery from controllers.openapi.account import ( AccountSessionByIdApi, AccountSessionsApi, AccountSessionsSelfApi, ) -from extensions.ext_redis import redis_client -from libs.oauth_bearer import TokenType from models import Account -from services.oauth_device_flow import MintResult, mint_oauth_token +from models.oauth import OAuthAccessToken from tests.test_containers_integration_tests.controllers.openapi.conftest import context_for +@dataclass(frozen=True, slots=True) +class _MintResult: + token_id: UUID + + def _mint_account_token( db_session: Session, account: Account, *, client_id: str = "integration-cli", device_label: str = "Test Device", -) -> MintResult: +) -> _MintResult: """Mint a real, persisted ``dfoa_`` access token for ``account``.""" - return mint_oauth_token( - redis_client, + record = OAuthAccessToken( subject_email=account.email, - subject_issuer=None, + subject_issuer="dify:account", account_id=str(account.id), client_id=client_id, device_label=device_label, - token_type=TokenType.OAUTH_ACCOUNT, - ttl_days=14, - session=db_session, + prefix=TokenType.OAUTH_ACCOUNT.prefix, + token_hash=f"integration-{uuid4().hex}", + expires_at=datetime.now(UTC) + timedelta(days=14), ) + db_session.add(record) + db_session.commit() + return _MintResult(token_id=UUID(str(record.id))) class TestSessionList: diff --git a/api/tests/test_containers_integration_tests/services/test_oauth_device_flow.py b/api/tests/test_containers_integration_tests/services/test_oauth_device_flow.py new file mode 100644 index 00000000000..1e32bdfe392 --- /dev/null +++ b/api/tests/test_containers_integration_tests/services/test_oauth_device_flow.py @@ -0,0 +1,72 @@ +"""Device-flow state transitions against real Redis, including Lua JSON encoding.""" + +import pytest + +from extensions.ext_redis import redis_client +from services.oauth_device_contracts import DeviceFlowStatus, InvalidTransitionError, PollPayload +from services.oauth_device_flow import DEVICE_CODE_KEY_FMT, USER_CODE_KEY_FMT, DeviceFlowRedis + +pytestmark = pytest.mark.usefixtures("flask_app_with_containers") + + +@pytest.mark.parametrize("subject_type", ["external_sso", "account"]) +@pytest.mark.parametrize( + "workspaces", + [[], [{"id": "workspace-1", "name": 'Workspace "quoted"\\path\n测试', "role": "owner"}]], + ids=["empty-workspaces", "nonempty-workspaces"], +) +def test_approval_preserves_payload_and_is_consumed_once( + subject_type: str, workspaces: list[dict[str, object]] +) -> None: + store = DeviceFlowRedis(redis_client) + device_code, user_code, _ = store.start("difyctl", "CLI", "127.0.0.1") + payload: PollPayload = { + "token": "test-device-token", + "expires_at": "2030-01-01T00:00:00+00:00", + "subject_type": subject_type, + "account": ( + {"id": "account-1", "email": "user@example.com", "name": "Example"} if subject_type == "account" else None + ), + "workspaces": workspaces, + "default_workspace_id": "workspace-1" if workspaces else None, + "token_id": "token-1", + } + if subject_type == "external_sso": + payload.update(subject_email="user@example.com", subject_issuer="https://idp.example") + + try: + store.approve(device_code, "approve-1", "token-1", payload) + # Retrying the same transition must not rewrite or corrupt the approved payload. + store.approve(device_code, "approve-1", "token-1", payload) + with pytest.raises(InvalidTransitionError): + store.deny(device_code, "deny-1") + + stored = store.load_by_device_code(device_code) + assert stored is not None + assert stored.status is DeviceFlowStatus.APPROVED + assert stored.poll_payload == payload + + consumed = store.consume_on_poll(device_code) + assert consumed is not None + assert consumed.poll_payload == payload + assert store.consume_on_poll(device_code) is None + assert store.load_by_user_code(user_code) is None + finally: + redis_client.delete(DEVICE_CODE_KEY_FMT.format(code=device_code), USER_CODE_KEY_FMT.format(code=user_code)) + + +def test_denied_state_has_no_payload_and_is_consumed_once() -> None: + store = DeviceFlowRedis(redis_client) + device_code, user_code, _ = store.start("difyctl", "CLI", "127.0.0.1") + try: + store.deny(device_code, "deny-1") + store.deny(device_code, "deny-1") + + consumed = store.consume_on_poll(device_code) + assert consumed is not None + assert consumed.status is DeviceFlowStatus.DENIED + assert consumed.poll_payload is None + assert store.consume_on_poll(device_code) is None + assert store.load_by_user_code(user_code) is None + finally: + redis_client.delete(DEVICE_CODE_KEY_FMT.format(code=device_code), USER_CODE_KEY_FMT.format(code=user_code)) diff --git a/api/tests/unit_tests/controllers/openapi/auth/_world.py b/api/tests/unit_tests/controllers/openapi/auth/_world.py index 1337cea7872..e1db34fa8f2 100644 --- a/api/tests/unit_tests/controllers/openapi/auth/_world.py +++ b/api/tests/unit_tests/controllers/openapi/auth/_world.py @@ -12,10 +12,11 @@ from typing import NoReturn from sqlalchemy.orm import Session +from constants.oauth_bearer import SubjectType, TokenType from controllers.openapi.auth.context import Context from controllers.openapi.auth.subjects import AccountSubject, ExternalSsoSubject, Subject from enums import DeploymentEdition -from libs.oauth_bearer import AuthContext, SubjectType, TokenType +from libs.oauth_bearer import AuthContext from models import Account, App, Tenant, TenantAccountJoin from models.account import AccountStatus, TenantAccountRole, TenantStatus from models.enums import AppStatus diff --git a/api/tests/unit_tests/controllers/openapi/auth/test_router.py b/api/tests/unit_tests/controllers/openapi/auth/test_router.py index d4d6b64dbf2..91161dc0d9f 100644 --- a/api/tests/unit_tests/controllers/openapi/auth/test_router.py +++ b/api/tests/unit_tests/controllers/openapi/auth/test_router.py @@ -12,6 +12,7 @@ from sqlalchemy.orm import Session, sessionmaker from werkzeug.exceptions import Forbidden, NotFound, ServiceUnavailable, Unauthorized import libs.rate_limit as rate_limit_module +from constants.oauth_bearer import TokenType from controllers.openapi._catalog import CATALOG_HEADER, catalog_for from controllers.openapi.auth.context import Context from controllers.openapi.auth.requirements import Requirement @@ -19,13 +20,7 @@ from controllers.openapi.auth.router import subject_router from controllers.openapi.auth.spec import CatalogMeta, EndpointSpec, Kind from controllers.openapi.auth.subjects import AccountSubject from enums import DeploymentEdition -from libs.oauth_bearer import ( - AuthContext, - BearerAuthenticator, - OAuthAccessTokenResolver, - Resolver, - TokenType, -) +from libs.oauth_bearer import AuthContext, BearerAuthenticator, OAuthAccessTokenResolver, Resolver from models import Account from models.oauth import OAuthAccessToken from services.entities.feature_entities import ( diff --git a/api/tests/unit_tests/controllers/openapi/auth/test_subjects.py b/api/tests/unit_tests/controllers/openapi/auth/test_subjects.py index a3bd212dae6..63c87b0d42b 100644 --- a/api/tests/unit_tests/controllers/openapi/auth/test_subjects.py +++ b/api/tests/unit_tests/controllers/openapi/auth/test_subjects.py @@ -7,8 +7,8 @@ from sqlalchemy import select from sqlalchemy.orm import Session, sessionmaker from werkzeug.exceptions import Unauthorized +from constants.oauth_bearer import TokenType from controllers.openapi.auth.subjects import AccountSubject, ExternalSsoSubject -from libs.oauth_bearer import TokenType from models import Account, EndUser, TenantAccountJoin from models.account import TenantAccountRole from models.enums import EndUserType diff --git a/api/tests/unit_tests/controllers/openapi/conftest.py b/api/tests/unit_tests/controllers/openapi/conftest.py index 14b9681faac..9eaf7ba3e9f 100644 --- a/api/tests/unit_tests/controllers/openapi/conftest.py +++ b/api/tests/unit_tests/controllers/openapi/conftest.py @@ -11,10 +11,11 @@ from flask.testing import FlaskClient from sqlalchemy.orm import Session from app_factory import create_flask_app_with_configs +from constants.oauth_bearer import TokenType from controllers.openapi import bp as openapi_bp from controllers.openapi._catalog import CATALOG_HEADER, catalog_for from enums import DeploymentEdition -from libs.oauth_bearer import AuthContext, TokenType +from libs.oauth_bearer import AuthContext from models import Account, App, Tenant, TenantAccountJoin from models.account import AccountStatus, TenantAccountRole, TenantStatus from models.enums import AppStatus diff --git a/api/tests/unit_tests/controllers/openapi/test_auth_matrix.py b/api/tests/unit_tests/controllers/openapi/test_auth_matrix.py index 7fd823c714e..5e3b75e580c 100644 --- a/api/tests/unit_tests/controllers/openapi/test_auth_matrix.py +++ b/api/tests/unit_tests/controllers/openapi/test_auth_matrix.py @@ -75,6 +75,7 @@ from werkzeug.test import TestResponse import libs.oauth_bearer as oauth_bearer_module import libs.rate_limit as rate_limit_module from app_factory import create_flask_app_with_configs +from constants.oauth_bearer import Scope, TokenType from controllers.common.rbac import PlainApp, RBACCheck, RBACPermission, Workspace from controllers.openapi import bp as openapi_bp from controllers.openapi._catalog import CATALOG_HEADER, catalog_for @@ -92,13 +93,7 @@ from controllers.openapi.auth.spec import EndpointSpec from controllers.openapi.auth.subjects import AccountSubject, ExternalSsoSubject from controllers.openapi.human_input_form import CheckFormSurface from enums import DeploymentEdition, WebAppAccessMode -from libs.oauth_bearer import ( - BearerAuthenticator, - ResolvedRow, - Scope, - TokenType, - sha256_hex, -) +from libs.oauth_bearer import BearerAuthenticator, ResolvedRow, sha256_hex from models.account import Account, AccountStatus, Tenant, TenantAccountJoin, TenantAccountRole from models.enums import EndUserType from models.model import App, EndUser diff --git a/api/tests/unit_tests/controllers/openapi/test_device_approve_deny.py b/api/tests/unit_tests/controllers/openapi/test_device_approve_deny.py index dbe2f7bfae6..6897a2b2e0d 100644 --- a/api/tests/unit_tests/controllers/openapi/test_device_approve_deny.py +++ b/api/tests/unit_tests/controllers/openapi/test_device_approve_deny.py @@ -7,7 +7,8 @@ from flask import Flask from flask.views import MethodView from controllers.openapi import bp as openapi_bp -from controllers.openapi.oauth_device import DeviceApproveApi, DeviceDenyApi +from controllers.openapi.oauth_device import DeviceApproveApi, DeviceDenyApi, _error_response +from services.oauth_device_contracts import ApprovalOutcomeUnknownError if not hasattr(builtins, "MethodView"): builtins.MethodView = MethodView # type: ignore[attr-defined] @@ -50,3 +51,10 @@ def test_approve_and_deny_methods(openapi_app: Flask): deny = _rule(openapi_app, "/openapi/v1/oauth/device/deny") assert "POST" in approve.methods assert "POST" in deny.methods + + +def test_unknown_approval_outcome_is_retryable() -> None: + body, status = _error_response(ApprovalOutcomeUnknownError()) + + assert status == 503 + assert body == {"error": "approval_outcome_unknown"} diff --git a/api/tests/unit_tests/controllers/openapi/test_device_sso.py b/api/tests/unit_tests/controllers/openapi/test_device_sso.py index c59a50fb501..0d5029cd173 100644 --- a/api/tests/unit_tests/controllers/openapi/test_device_sso.py +++ b/api/tests/unit_tests/controllers/openapi/test_device_sso.py @@ -1,19 +1,22 @@ """SSO-branch device-flow endpoints under /openapi/v1/oauth/device/.""" import builtins -from unittest.mock import patch +from dataclasses import dataclass import pytest from flask import Flask from flask.views import MethodView +from werkzeug.exceptions import ServiceUnavailable from controllers.openapi import bp as openapi_bp from controllers.openapi.oauth_device_sso import ( + _raise_http_error, approval_context, approve_external, sso_complete, sso_initiate, ) +from services.oauth_device_contracts import ApprovalOutcomeUnknownError, DeviceSSOCompletion if not hasattr(builtins, "MethodView"): builtins.MethodView = MethodView # type: ignore[attr-defined] @@ -71,13 +74,9 @@ def test_approve_external_dispatches_to_function(openapi_app: Flask): assert openapi_app.view_functions[rule.endpoint] is approve_external -def test_sso_complete_idp_callback_url_uses_canonical_path(): - """sso_initiate hardcodes the IdP callback URL — must point at the - canonical /openapi/v1/ path so IdP-side ACS configuration matches. - """ - from controllers.openapi import oauth_device_sso - - assert oauth_device_sso._SSO_COMPLETE_PATH == "/openapi/v1/oauth/device/sso-complete" +def test_unknown_external_approval_outcome_is_retryable() -> None: + with pytest.raises(ServiceUnavailable, match="approval_outcome_unknown"): + _raise_http_error(ApprovalOutcomeUnknownError()) # --------------------------------------------------------------------------- @@ -142,15 +141,46 @@ def test_device_error_redirect_drops_malformed_user_code(): # --------------------------------------------------------------------------- -def _ee_features(): - from services.entities.feature_entities import LicenseStatus - - return LicenseStatus.ACTIVE +class _CompletionService: + def complete_sso(self, _context, *, inbound_error, inbound_user_code, assertion): + _ = assertion + if inbound_error: + return DeviceSSOCompletion(error_code=inbound_error, user_code=inbound_user_code) + return DeviceSSOCompletion(error_code="sso_failed") -@patch("libs.device_flow_security.SystemFeatureService.get_license_status") -def test_sso_complete_relays_inbound_sso_error(ee_feat, openapi_app): - ee_feat.return_value = _ee_features() +@dataclass(frozen=True, slots=True) +class _FeatureQueries: + valid_enterprise_license: bool + + def has_valid_enterprise_license(self) -> bool: + return self.valid_enterprise_license + + +@dataclass(frozen=True, slots=True) +class _ApplicationServices: + oauth_device: _CompletionService + feature_queries: _FeatureQueries + + +def _install_application_services(monkeypatch: pytest.MonkeyPatch, *, valid_enterprise_license: bool) -> None: + from controllers.openapi import flask_admission, oauth_device_sso + + services = _ApplicationServices( + oauth_device=_CompletionService(), + feature_queries=_FeatureQueries(valid_enterprise_license=valid_enterprise_license), + ) + monkeypatch.setattr(flask_admission, "application_services", lambda: services) + monkeypatch.setattr(oauth_device_sso, "application_services", lambda: services) + + +@pytest.fixture +def admitted_sso(monkeypatch: pytest.MonkeyPatch) -> None: + _install_application_services(monkeypatch, valid_enterprise_license=True) + + +def test_sso_complete_relays_inbound_sso_error(openapi_app, admitted_sso): + _ = admitted_sso client = openapi_app.test_client() resp = client.get( "/openapi/v1/oauth/device/sso-complete?sso_error=sso_failed&user_code=ABCD-1234", @@ -163,10 +193,17 @@ def test_sso_complete_relays_inbound_sso_error(ee_feat, openapi_app): assert "user_code=ABCD-1234" in loc -@patch("libs.device_flow_security.SystemFeatureService.get_license_status") -def test_sso_complete_missing_assertion_redirects_generic(ee_feat, openapi_app): - ee_feat.return_value = _ee_features() +def test_sso_complete_missing_assertion_redirects_generic(openapi_app, admitted_sso): + _ = admitted_sso client = openapi_app.test_client() resp = client.get("/openapi/v1/oauth/device/sso-complete", follow_redirects=False) assert resp.status_code == 302 assert "sso_error=sso_failed" in resp.headers["Location"] + + +def test_sso_admission_rejects_inactive_license(openapi_app, monkeypatch: pytest.MonkeyPatch): + _install_application_services(monkeypatch, valid_enterprise_license=False) + + response = openapi_app.test_client().get("/openapi/v1/oauth/device/sso-complete") + + assert response.status_code == 404 diff --git a/api/tests/unit_tests/controllers/openapi/test_error_contract.py b/api/tests/unit_tests/controllers/openapi/test_error_contract.py index 4d3d2934e05..28e7541733b 100644 --- a/api/tests/unit_tests/controllers/openapi/test_error_contract.py +++ b/api/tests/unit_tests/controllers/openapi/test_error_contract.py @@ -2,7 +2,7 @@ import uuid from types import SimpleNamespace -from unittest.mock import MagicMock, patch +from unittest.mock import patch import pytest from sqlalchemy.orm import Session @@ -18,6 +18,7 @@ from werkzeug.exceptions import ( UnprocessableEntity, ) +from constants.oauth_bearer import TokenType from controllers.common.errors import ( BlockedFileExtensionError, FileTooLargeError, @@ -49,8 +50,9 @@ from controllers.service_api.app.error import ( ProviderQuotaExceededError, ) from controllers.web.error import InvokeRateLimitError as InvokeRateLimitHttpError -from libs.oauth_bearer import AuthContext, TokenType +from libs.oauth_bearer import AuthContext from models.account import Account, AccountStatus +from services.oauth_device_contracts import ExpiredTokenError @pytest.fixture @@ -261,12 +263,9 @@ class TestWireContract: # not intercepted → Flask's default HTML 404, not the canonical JSON body assert "application/json" not in (resp.content_type or "") - @patch("controllers.openapi.oauth_device.DeviceFlowRedis") - def test_oauth_device_token_keeps_rfc8628_shape(self, mock_redis_cls, openapi_app): - store = MagicMock() - mock_redis_cls.return_value = store - store.record_poll.return_value = None # not SlowDownDecision.SLOW_DOWN - store.load_by_device_code.return_value = None # unknown code → expired_token + @patch("controllers.openapi.oauth_device.application_services") + def test_oauth_device_token_keeps_rfc8628_shape(self, mock_application_services, openapi_app): + mock_application_services.return_value.oauth_device.poll.side_effect = ExpiredTokenError client = openapi_app.test_client() diff --git a/api/tests/unit_tests/controllers/openapi/test_oauth_device_http.py b/api/tests/unit_tests/controllers/openapi/test_oauth_device_http.py new file mode 100644 index 00000000000..9ba8eb8faae --- /dev/null +++ b/api/tests/unit_tests/controllers/openapi/test_oauth_device_http.py @@ -0,0 +1,383 @@ +"""Device-flow HTTP contracts with application services isolated at their public boundary.""" + +from collections.abc import Callable +from datetime import UTC, datetime +from http.cookies import SimpleCookie +from types import SimpleNamespace +from unittest.mock import Mock, create_autospec + +import pytest +from flask import Flask + +from controllers.console import flask_admission, wraps +from libs.helper import RateLimiter +from libs.login import AccountWithTenant +from machinery.context import RequestContext +from models.account import Account, AccountStatus +from services.oauth_device_application_service import OAuthDeviceApplicationService +from services.oauth_device_contracts import ( + AccessDeniedError, + AlreadyResolvedError, + ApprovalInProgressError, + ApprovalOutcomeUnknownError, + ApprovalSessionConsumedError, + AuthorizationPendingError, + DeviceApprovalContext, + DeviceAuthorization, + DeviceLookup, + DeviceMutation, + DeviceRequestContext, + DeviceSSOCompletion, + DeviceSSOInitiation, + DeviceStateLostError, + ExpiredOrUnknownError, + ExpiredTokenError, + ExternalApprovalCSRFError, + ExternalApprovalRateLimitError, + ExternalIdentityConflictError, + ExternalUserCodeMismatchError, + ExternalUserCodeNotFoundError, + InvalidApprovalSessionError, + InvalidUserCodeError, + OAuthDeviceError, + OAuthDeviceSSOConfigurationError, + OAuthDeviceSSOInitiationError, + PollTooFastError, + UnsupportedClientError, +) + +_PREFIX = "/openapi/v1/oauth/device" +_COOKIE = "device_approval_grant" + + +@pytest.fixture +def device_service(openapi_app: Flask, monkeypatch: pytest.MonkeyPatch, config_overrides: Callable[..., None]) -> Mock: + service = create_autospec(OAuthDeviceApplicationService, instance=True) + openapi_app.extensions["application_services"] = SimpleNamespace( + oauth_device=service, + feature_queries=SimpleNamespace(has_valid_enterprise_license=lambda: True), + ) + # Supply an authenticated account while retaining the real admission/context wrapper. + account = Account(name="Ada", email="ada@example.com", status=AccountStatus.ACTIVE) + account.id = "account-1" + account_with_tenant = AccountWithTenant(account, "workspace-1") + config_overrides( + LOGIN_DISABLED=True, + ENABLE_OAUTH_BEARER=True, + CONSOLE_API_URL="https://api.example", + CONSOLE_WEB_URL="https://console.example", + ) + monkeypatch.setattr(wraps, "_is_setup_completed", lambda: True) + monkeypatch.setattr(wraps, "current_account_with_tenant", lambda: account_with_tenant) + monkeypatch.setattr(flask_admission, "current_account_with_tenant", lambda: account_with_tenant) + monkeypatch.setattr(RateLimiter, "is_rate_limited", lambda *_args: False) + monkeypatch.setattr(RateLimiter, "increment_rate_limit", lambda *_args: None) + return service + + +def test_device_code_serializes_authorization_and_passes_request_origin( + openapi_app: Flask, device_service: Mock +) -> None: + device_service.start.return_value = DeviceAuthorization( + device_code="opaque-code", + user_code="ABCD-EFGH", + verification_uri="https://console.example/device", + expires_in=900, + interval=5, + ) + + response = openapi_app.test_client().post( + f"{_PREFIX}/code", + json={"client_id": "difyctl", "device_label": "laptop"}, + base_url="https://api.example", + environ_overrides={"REMOTE_ADDR": "192.0.2.10"}, + ) + + assert response.status_code == 200 + assert response.json == { + "device_code": "opaque-code", + "user_code": "ABCD-EFGH", + "verification_uri": "https://console.example/device", + "expires_in": 900, + "interval": 5, + } + device_service.start.assert_called_once_with( + client_id="difyctl", + device_label="laptop", + created_ip="192.0.2.10", + request_origin="https://api.example/", + ) + + +@pytest.mark.parametrize("path", ["code", "token", "approve", "deny", "approve-external"]) +@pytest.mark.parametrize("body", ["{}", "[]", "{invalid"]) +def test_invalid_device_request_never_reaches_service( + openapi_app: Flask, device_service: Mock, path: str, body: str +) -> None: + response = openapi_app.test_client().post(f"{_PREFIX}/{path}", data=body, content_type="application/json") + + assert response.status_code == 400 + assert device_service.mock_calls == [] + + +def test_lookup_requires_user_code(openapi_app: Flask, device_service: Mock) -> None: + response = openapi_app.test_client().get(f"{_PREFIX}/lookup") + + assert response.status_code == 400 + device_service.lookup.assert_not_called() + + +def test_lookup_returns_public_validity_without_token_data(openapi_app: Flask, device_service: Mock) -> None: + device_service.lookup.return_value = DeviceLookup(valid=False, expires_in_remaining=0, client_id=None) + + response = openapi_app.test_client().get(f"{_PREFIX}/lookup?user_code=ABCD-EFGH") + + assert response.status_code == 200 + assert response.json == {"valid": False, "expires_in_remaining": 0, "client_id": None} + device_service.lookup.assert_called_once_with(user_code="ABCD-EFGH") + + +@pytest.mark.parametrize( + ("path", "method", "error", "status", "code"), + [ + ("code", "start", UnsupportedClientError(), 400, "unsupported_client"), + ("token", "poll", PollTooFastError(), 400, "slow_down"), + ("token", "poll", ExpiredTokenError(), 400, "expired_token"), + ("token", "poll", AuthorizationPendingError(), 400, "authorization_pending"), + ("token", "poll", AccessDeniedError(), 400, "access_denied"), + ("approve", "approve", ExpiredOrUnknownError(), 404, "expired_or_unknown"), + ("approve", "approve", AlreadyResolvedError(), 409, "already_resolved"), + ("approve", "approve", ApprovalInProgressError(), 409, "approve_in_progress"), + ("approve", "approve", ApprovalOutcomeUnknownError(), 503, "approval_outcome_unknown"), + ("deny", "deny", DeviceStateLostError(), 409, "state_lost"), + ], +) +def test_device_errors_keep_oauth_response_contract( + openapi_app: Flask, device_service: Mock, path: str, method: str, error: OAuthDeviceError, status: int, code: str +) -> None: + methods = { + "start": device_service.start, + "poll": device_service.poll, + "approve": device_service.approve, + "deny": device_service.deny, + } + methods[method].side_effect = error + + response = openapi_app.test_client().post( + f"{_PREFIX}/{path}", + json={"client_id": "difyctl", "device_label": "CLI", "device_code": "opaque-code", "user_code": "ABCD-EFGH"}, + ) + + assert response.status_code == status + assert response.json == {"error": code} + + +@pytest.mark.parametrize("external", [False, True], ids=["account", "external"]) +def test_poll_serializes_subject_specific_token_payload( + openapi_app: Flask, device_service: Mock, external: bool +) -> None: + payload: dict[str, object] = { + "token": "secret-token", + "expires_at": "2030-01-01T00:00:00Z", + "token_id": "token-1", + "subject_type": "external_sso" if external else "account", + "account": None if external else {"id": "account-1", "email": "ada@example.com", "name": "Ada"}, + "workspaces": [], + "default_workspace_id": None, + } + if external: + payload.update(subject_email="external@example.com", subject_issuer="https://idp.example") + device_service.poll.return_value = payload + + response = openapi_app.test_client().post( + f"{_PREFIX}/token", + json={"device_code": "opaque-code", "client_id": "difyctl"}, + environ_overrides={"REMOTE_ADDR": "192.0.2.10"}, + ) + + assert response.status_code == 200 + assert response.json == payload + device_service.poll.assert_called_once_with(device_code="opaque-code", poll_ip="192.0.2.10") + + +@pytest.mark.parametrize("path", ["approve", "deny"]) +def test_account_decision_uses_admitted_identity(openapi_app: Flask, device_service: Mock, path: str) -> None: + method = device_service.approve if path == "approve" else device_service.deny + method.return_value = DeviceMutation(status="approved" if path == "approve" else "denied") + + response = openapi_app.test_client().post( + f"{_PREFIX}/{path}", + json={"user_code": "ABCD-EFGH", "account_id": "forged", "workspace_id": "foreign"}, + ) + + assert response.status_code == 200 + assert response.json == {"status": method.return_value.status} + assert method.call_args.kwargs == {"user_code": "ABCD-EFGH"} + if path == "approve": + context = method.call_args.args[0] + assert isinstance(context, RequestContext) + assert context.account_id == "account-1" + assert context.active_workspace_id == "workspace-1" + + +@pytest.mark.parametrize("path", ["approve", "deny"]) +def test_disabled_bearer_feature_blocks_account_decisions( + openapi_app: Flask, device_service: Mock, config_overrides: Callable[..., None], path: str +) -> None: + config_overrides(ENABLE_OAUTH_BEARER=False) + + response = openapi_app.test_client().post(f"{_PREFIX}/{path}", json={"user_code": "ABCD-EFGH"}) + + assert response.status_code == 503 + assert device_service.mock_calls == [] + + +def test_sso_initiation_clears_old_approval_cookie(openapi_app: Flask, device_service: Mock) -> None: + device_service.initiate_sso.return_value = DeviceSSOInitiation("https://idp.example/authorize") + + response = openapi_app.test_client().get(f"{_PREFIX}/sso-initiate?user_code=%20ABCD-EFGH%20") + + assert response.status_code == 302 + assert response.headers["Location"] == "https://idp.example/authorize" + cookie = SimpleCookie(response.headers["Set-Cookie"])[_COOKIE] + assert cookie["max-age"] == "0" + assert cookie["path"] == _PREFIX + assert device_service.initiate_sso.call_args.kwargs == {"user_code": "ABCD-EFGH"} + + +def test_sso_initiation_requires_user_code(openapi_app: Flask, device_service: Mock) -> None: + response = openapi_app.test_client().get(f"{_PREFIX}/sso-initiate?user_code=%20") + + assert response.status_code == 400 + device_service.initiate_sso.assert_not_called() + + +def test_sso_completion_sets_scoped_http_only_cookie(openapi_app: Flask, device_service: Mock) -> None: + device_service.complete_sso.return_value = DeviceSSOCompletion(user_code="ABCD-EFGH", approval_grant="signed-grant") + + response = openapi_app.test_client().get(f"{_PREFIX}/sso-complete?sso_assertion=signed-assertion") + + assert response.status_code == 302 + assert response.headers["Location"] == "/device?sso_verified=1" + cookie = SimpleCookie(response.headers["Set-Cookie"])[_COOKIE] + assert cookie.value == "signed-grant" + assert cookie["path"] == _PREFIX + assert cookie["httponly"] + assert cookie["secure"] + assert cookie["samesite"] == "Lax" + assert cookie["max-age"] == "300" + assert device_service.complete_sso.call_args.kwargs == { + "inbound_error": None, + "inbound_user_code": None, + "assertion": "signed-assertion", + } + + +@pytest.mark.parametrize("failure", ["exception", "missing-grant", "untrusted-error"]) +def test_sso_completion_failure_redirects_without_exposing_details( + openapi_app: Flask, device_service: Mock, failure: str +) -> None: + device_service.complete_sso.return_value = DeviceSSOCompletion( + error_code="private-provider-error" if failure == "untrusted-error" else None, + ) + if failure == "exception": + device_service.complete_sso.side_effect = RuntimeError("private-provider-error") + + response = openapi_app.test_client().get(f"{_PREFIX}/sso-complete?sso_assertion=private-assertion") + + assert response.status_code == 302 + assert response.headers["Location"] == "/device?sso_error=sso_failed" + assert "Set-Cookie" not in response.headers + assert b"private-" not in response.data + + +def test_approval_context_reads_cookie_and_serializes_expiry(openapi_app: Flask, device_service: Mock) -> None: + device_service.get_approval_context.return_value = DeviceApprovalContext( + subject_email="external@example.com", + subject_issuer="https://idp.example", + user_code="ABCD-EFGH", + csrf_token="csrf-token", + expires_at=datetime(2030, 1, 1, tzinfo=UTC), + ) + client = openapi_app.test_client() + client.set_cookie(_COOKIE, "signed-grant", path=_PREFIX) + + response = client.get(f"{_PREFIX}/approval-context") + + assert response.status_code == 200 + assert response.json == { + "subject_email": "external@example.com", + "subject_issuer": "https://idp.example", + "user_code": "ABCD-EFGH", + "csrf_token": "csrf-token", + "expires_at": "2030-01-01T00:00:00Z", + } + assert device_service.get_approval_context.call_args.kwargs == {"approval_grant": "signed-grant"} + + +def test_external_approval_uses_cookie_and_csrf_header_then_clears_cookie( + openapi_app: Flask, device_service: Mock +) -> None: + device_service.approve_external.return_value = DeviceMutation("approved") + client = openapi_app.test_client() + client.set_cookie(_COOKIE, "signed-grant", path=_PREFIX) + + response = client.post( + f"{_PREFIX}/approve-external", + headers={"X-CSRF-Token": "csrf-header"}, + json={"user_code": "ABCD-EFGH", "csrf_token": "forged", "approval_grant": "forged"}, + ) + + assert response.status_code == 200 + assert response.json == {"status": "approved"} + assert isinstance(device_service.approve_external.call_args.args[0], DeviceRequestContext) + assert device_service.approve_external.call_args.kwargs == { + "approval_grant": "signed-grant", + "csrf_token": "csrf-header", + "user_code": "ABCD-EFGH", + } + cookie = SimpleCookie(response.headers["Set-Cookie"])[_COOKIE] + assert cookie["max-age"] == "0" + assert cookie["path"] == _PREFIX + + +@pytest.mark.parametrize( + ("path", "method", "error", "status", "message"), + [ + ("sso-initiate", "initiate_sso", InvalidUserCodeError(), 400, "invalid_user_code"), + ("sso-initiate", "initiate_sso", OAuthDeviceSSOConfigurationError(), 502, "console_api_url_unset"), + ("sso-initiate", "initiate_sso", OAuthDeviceSSOInitiationError(), 502, "sso_initiate_failed"), + ("approval-context", "get_approval_context", InvalidApprovalSessionError(), 401, "no_session"), + ("approve-external", "approve_external", InvalidApprovalSessionError(), 401, "invalid_session"), + ("approve-external", "approve_external", ExternalApprovalRateLimitError(), 429, "rate_limited"), + ("approve-external", "approve_external", ExternalApprovalCSRFError(), 403, "csrf_mismatch"), + ("approve-external", "approve_external", ExternalUserCodeMismatchError(), 400, "user_code_mismatch"), + ("approve-external", "approve_external", ExternalUserCodeNotFoundError(), 404, "user_code_not_pending"), + ("approve-external", "approve_external", AlreadyResolvedError(), 409, "user_code_not_pending"), + ("approve-external", "approve_external", ApprovalInProgressError(), 409, "approve_in_progress"), + ("approve-external", "approve_external", ApprovalOutcomeUnknownError(), 503, "approval_outcome_unknown"), + ("approve-external", "approve_external", ExternalIdentityConflictError(), 403, "email_belongs_to_dify_account"), + ("approve-external", "approve_external", ApprovalSessionConsumedError(), 401, "session_already_consumed"), + ("approve-external", "approve_external", DeviceStateLostError(), 409, "state_lost"), + ], +) +def test_sso_errors_preserve_http_status_and_approval_cookie( + openapi_app: Flask, device_service: Mock, path: str, method: str, error: OAuthDeviceError, status: int, message: str +) -> None: + methods = { + "initiate_sso": device_service.initiate_sso, + "get_approval_context": device_service.get_approval_context, + "approve_external": device_service.approve_external, + } + methods[method].side_effect = error + client = openapi_app.test_client() + client.set_cookie(_COOKIE, "signed-grant", path=_PREFIX) + if path == "approve-external": + response = client.post(f"{_PREFIX}/{path}", json={"user_code": "ABCD-EFGH"}) + else: + response = client.get(f"{_PREFIX}/{path}?user_code=ABCD-EFGH") + + assert response.status_code == status + assert message in response.get_data(as_text=True) + assert "Set-Cookie" not in response.headers + assert "signed-grant" not in response.get_data(as_text=True) diff --git a/api/tests/unit_tests/controllers/openapi/test_oauth_sso_claims.py b/api/tests/unit_tests/controllers/openapi/test_oauth_sso_claims.py deleted file mode 100644 index 59fdfc24d67..00000000000 --- a/api/tests/unit_tests/controllers/openapi/test_oauth_sso_claims.py +++ /dev/null @@ -1,73 +0,0 @@ -from __future__ import annotations - -from unittest.mock import patch - -import pytest -from flask import Flask - -from controllers.openapi import bp as openapi_bp - - -@pytest.fixture -def app() -> Flask: - a = Flask(__name__) - a.config["TESTING"] = True - a.register_blueprint(openapi_bp) - return a - - -def _ee_features(): - from services.entities.feature_entities import LicenseStatus - - return LicenseStatus.ACTIVE - - -@patch("controllers.openapi.oauth_device_sso.jws") -@patch("libs.device_flow_security.SystemFeatureService.get_license_status") -def test_sso_complete_rejects_assertion_missing_email(ee_feat, jws_mod, app: Flask): - ee_feat.return_value = _ee_features() - jws_mod.verify.return_value = {"issuer": "https://idp.example", "user_code": "ABCD-EFGH", "nonce": "n"} - jws_mod.AUD_EXT_SUBJECT_ASSERTION = "aud" - jws_mod.KeySet.from_shared_secret.return_value = object() - jws_mod.VerifyError = Exception - - client = app.test_client() - resp = client.get("/openapi/v1/oauth/device/sso-complete?sso_assertion=blob", follow_redirects=False) - assert resp.status_code == 302, resp.data - assert "sso_error=sso_failed" in resp.headers["Location"] - - -@patch("controllers.openapi.oauth_device_sso.jws") -@patch("libs.device_flow_security.SystemFeatureService.get_license_status") -def test_sso_complete_rejects_assertion_empty_issuer(ee_feat, jws_mod, app: Flask): - ee_feat.return_value = _ee_features() - jws_mod.verify.return_value = {"email": "x@y.com", "issuer": "", "user_code": "ABCD-EFGH", "nonce": "n"} - jws_mod.AUD_EXT_SUBJECT_ASSERTION = "aud" - jws_mod.KeySet.from_shared_secret.return_value = object() - jws_mod.VerifyError = Exception - - client = app.test_client() - resp = client.get("/openapi/v1/oauth/device/sso-complete?sso_assertion=blob", follow_redirects=False) - assert resp.status_code == 302 - assert "sso_error=sso_failed" in resp.headers["Location"] - - -def test_verify_approval_grant_raises_on_missing_field(): - from libs import device_flow_security - from libs import jws as jws_mod - - class _FakeKeyset: - active_kid = "k" - - def lookup(self, kid): - return b"secret" - - keyset = _FakeKeyset() - incomplete = jws_mod.sign( - keyset, - payload={"subject_email": "x@y.com", "subject_issuer": "i", "user_code": "ABCD-EFGH", "nonce": "n"}, - aud=jws_mod.AUD_APPROVAL_GRANT, - ttl_seconds=60, - ) - with pytest.raises(jws_mod.VerifyError): - device_flow_security.verify_approval_grant(keyset, incomplete) diff --git a/api/tests/unit_tests/controllers/openapi/test_oauth_sso_csrf.py b/api/tests/unit_tests/controllers/openapi/test_oauth_sso_csrf.py deleted file mode 100644 index e31a308506b..00000000000 --- a/api/tests/unit_tests/controllers/openapi/test_oauth_sso_csrf.py +++ /dev/null @@ -1,22 +0,0 @@ -from __future__ import annotations - -import ast -from pathlib import Path - - -def _repo_root() -> Path: - for parent in Path(__file__).resolve().parents: - if (parent / "api" / "pyproject.toml").exists(): - return parent - raise RuntimeError("repo root not found") - - -def test_approve_external_uses_compare_digest_for_csrf(): - src = (_repo_root() / "api" / "controllers" / "openapi" / "oauth_device_sso.py").read_text() - tree = ast.parse(src) - - fn = next(n for n in ast.walk(tree) if isinstance(n, ast.FunctionDef) and n.name == "approve_external") - fn_src = ast.unparse(fn) - - assert "compare_digest" in fn_src, "approve_external must call secrets.compare_digest for CSRF" - assert "csrf_header != claims.csrf_token" not in fn_src, "approve_external must not use plain != on csrf_token" diff --git a/api/tests/unit_tests/controllers/openapi/test_oauth_sso_host_header.py b/api/tests/unit_tests/controllers/openapi/test_oauth_sso_host_header.py deleted file mode 100644 index 2c9b1778927..00000000000 --- a/api/tests/unit_tests/controllers/openapi/test_oauth_sso_host_header.py +++ /dev/null @@ -1,76 +0,0 @@ -from __future__ import annotations - -from unittest.mock import MagicMock, patch -from urllib.parse import urlparse - -import pytest -from flask import Flask - -from controllers.openapi import bp as openapi_bp - - -@pytest.fixture -def app() -> Flask: - a = Flask(__name__) - a.config["TESTING"] = True - a.register_blueprint(openapi_bp) - return a - - -def _ee_features(): - from services.entities.feature_entities import LicenseStatus - - return LicenseStatus.ACTIVE - - -@patch("controllers.openapi.oauth_device_sso.EnterpriseService") -@patch("controllers.openapi.oauth_device_sso.jws") -@patch("controllers.openapi.oauth_device_sso.DeviceFlowRedis") -@patch("libs.device_flow_security.SystemFeatureService.get_license_status") -@patch("libs.rate_limit.RateLimiter.is_rate_limited", new=MagicMock(return_value=False)) -@patch("libs.rate_limit.RateLimiter.increment_rate_limit", new=MagicMock()) -def test_idp_callback_url_uses_console_api_url_not_host_header( - ee_feat, redis_cls, jws_mod, ent, app: Flask, config_overrides -): - ee_feat.return_value = _ee_features() - config_overrides(CONSOLE_API_URL="https://api.dify.example") - state = MagicMock() - from services.oauth_device_flow import DeviceFlowStatus - - state.status = DeviceFlowStatus.PENDING - redis_cls.return_value.load_by_user_code.return_value = ("dc_x", state) - jws_mod.KeySet.from_shared_secret.return_value = MagicMock() - jws_mod.sign.return_value = "signed-state" - jws_mod.AUD_STATE_ENVELOPE = "aud" - ent.initiate_device_flow_sso.return_value = {"url": "https://idp.example/auth"} - - client = app.test_client() - client.get( - "/openapi/v1/oauth/device/sso-initiate?user_code=ABCD-EFGH", - headers={"Host": "evil.com"}, - ) - - args, kwargs = jws_mod.sign.call_args - signed_payload = args[1] if len(args) > 1 else kwargs["payload"] - callback_url = urlparse(signed_payload["idp_callback_url"]) - assert callback_url.scheme == "https" - assert callback_url.hostname == "api.dify.example" - assert "evil.com" not in signed_payload["idp_callback_url"] - - -@patch("controllers.openapi.oauth_device_sso.DeviceFlowRedis") -@patch("libs.device_flow_security.SystemFeatureService.get_license_status") -@patch("libs.rate_limit.RateLimiter.is_rate_limited", new=MagicMock(return_value=False)) -@patch("libs.rate_limit.RateLimiter.increment_rate_limit", new=MagicMock()) -def test_sso_initiate_fails_closed_when_console_api_url_unset(ee_feat, redis_cls, app: Flask, config_overrides): - ee_feat.return_value = _ee_features() - config_overrides(CONSOLE_API_URL="") - from services.oauth_device_flow import DeviceFlowStatus - - state = MagicMock() - state.status = DeviceFlowStatus.PENDING - redis_cls.return_value.load_by_user_code.return_value = ("dc_x", state) - - client = app.test_client() - resp = client.get("/openapi/v1/oauth/device/sso-initiate?user_code=ABCD-EFGH") - assert resp.status_code == 502 diff --git a/api/tests/unit_tests/controllers/openapi/test_workspaces_members.py b/api/tests/unit_tests/controllers/openapi/test_workspaces_members.py index f1f4293ce42..c31581b8669 100644 --- a/api/tests/unit_tests/controllers/openapi/test_workspaces_members.py +++ b/api/tests/unit_tests/controllers/openapi/test_workspaces_members.py @@ -35,6 +35,7 @@ from sqlalchemy.orm import Session, sessionmaker from werkzeug.exceptions import BadRequest, NotFound from werkzeug.test import TestResponse +from constants.oauth_bearer import TokenType from controllers.openapi import bp as openapi_bp from controllers.openapi import workspaces as workspaces_module from controllers.openapi._catalog import CATALOG_HEADER, catalog_for @@ -50,7 +51,7 @@ from controllers.openapi.auth.loaders import load_caller, load_workspace from controllers.openapi.auth.subjects import AccountSubject from controllers.openapi.workspaces import WorkspaceMemberApi, WorkspaceMembersApi, WorkspaceSwitchApi from enums import DeploymentEdition -from libs.oauth_bearer import AuthContext, TokenType +from libs.oauth_bearer import AuthContext from models import Account, Tenant, TenantAccountJoin from models.account import TenantAccountRole from services.account_service import TenantService as RealTenantService diff --git a/api/tests/unit_tests/extensions/test_redis.py b/api/tests/unit_tests/extensions/test_redis.py index f86420b2b63..4de90d0afb1 100644 --- a/api/tests/unit_tests/extensions/test_redis.py +++ b/api/tests/unit_tests/extensions/test_redis.py @@ -1,4 +1,5 @@ -from unittest.mock import MagicMock +from collections.abc import Mapping +from dataclasses import dataclass, field import pytest from redis import RedisError @@ -15,6 +16,49 @@ from extensions.ext_redis import ( redis_fallback, ) +type _RecordedCall = tuple[str, tuple[object, ...], dict[str, object]] + + +@dataclass +class _RecordingRedisClient: + calls: list[_RecordedCall] = field(default_factory=list) + + def _record(self, method: str, *args: object, **kwargs: object) -> None: + self.calls.append((method, args, kwargs)) + + def register_script(self, script: str): + self._record("register_script", script) + + def execute(*, keys, args, client): + self._record("script", keys, args, client) + return "result" + + return execute + + def get(self, name: str | bytes) -> None: + self._record("get", name) + + def delete(self, *names: str | bytes) -> None: + self._record("delete", *names) + + def lock(self, name: str, **kwargs: object) -> None: + self._record("lock", name, **kwargs) + + def hset(self, name: str | bytes, *args: object, **kwargs: object) -> None: + self._record("hset", name, *args, **kwargs) + + def hgetall(self, name: str | bytes) -> None: + self._record("hgetall", name) + + def hkeys(self, name: str | bytes) -> None: + self._record("hkeys", name) + + def hexists(self, name: str | bytes, key: str | bytes) -> None: + self._record("hexists", name, key) + + def zadd(self, name: str | bytes, mapping: Mapping[object, object], **kwargs: object) -> None: + self._record("zadd", name, mapping, **kwargs) + @pytest.fixture(autouse=True) def _redis_config(config_overrides) -> None: @@ -148,43 +192,57 @@ class TestRedisKeyPrefixHelpers: class TestRedisClientWrapperKeyPrefix: - def test_wrapper_get_prefixes_string_keys(self, config_overrides): - mock_client = MagicMock() + def test_wrapper_registered_script_prefixes_key_arguments(self, config_overrides): + raw_client = _RecordingRedisClient() wrapper = RedisClientWrapper() - wrapper.initialize(mock_client) + wrapper.initialize(raw_client) # type: ignore[arg-type] + + config_overrides(REDIS_KEY_PREFIX="enterprise-a") + script = wrapper.register_script("return redis.call('GET', KEYS[1])") + + assert script(keys=["device_code:abc"], args=["argument"]) == "result" + assert raw_client.calls == [ + ("register_script", ("return redis.call('GET', KEYS[1])",), {}), + ("script", (("enterprise-a:device_code:abc",), ["argument"], raw_client), {}), + ] + + def test_wrapper_get_prefixes_string_keys(self, config_overrides): + client = _RecordingRedisClient() + wrapper = RedisClientWrapper() + wrapper.initialize(client) # type: ignore[arg-type] config_overrides(REDIS_KEY_PREFIX="enterprise-a") wrapper.get("oauth_state:abc") - mock_client.get.assert_called_once_with("enterprise-a:oauth_state:abc") + assert client.calls == [("get", ("enterprise-a:oauth_state:abc",), {})] def test_wrapper_delete_prefixes_multiple_keys(self, config_overrides): - mock_client = MagicMock() + client = _RecordingRedisClient() wrapper = RedisClientWrapper() - wrapper.initialize(mock_client) + wrapper.initialize(client) # type: ignore[arg-type] config_overrides(REDIS_KEY_PREFIX="enterprise-a") wrapper.delete("key:a", "key:b") - mock_client.delete.assert_called_once_with("enterprise-a:key:a", "enterprise-a:key:b") + assert client.calls == [("delete", ("enterprise-a:key:a", "enterprise-a:key:b"), {})] def test_wrapper_lock_prefixes_lock_name(self, config_overrides): - mock_client = MagicMock() + client = _RecordingRedisClient() wrapper = RedisClientWrapper() - wrapper.initialize(mock_client) + wrapper.initialize(client) # type: ignore[arg-type] config_overrides(REDIS_KEY_PREFIX="enterprise-a") wrapper.lock("resource-lock", timeout=10) - mock_client.lock.assert_called_once() - args, kwargs = mock_client.lock.call_args + method, args, kwargs = client.calls[0] + assert method == "lock" assert args == ("enterprise-a:resource-lock",) assert kwargs["timeout"] == 10 def test_wrapper_hash_operations_prefix_key_name(self, config_overrides): - mock_client = MagicMock() + client = _RecordingRedisClient() wrapper = RedisClientWrapper() - wrapper.initialize(mock_client) + wrapper.initialize(client) # type: ignore[arg-type] config_overrides(REDIS_KEY_PREFIX="enterprise-a") wrapper.hset("hash:key", "field", "value") @@ -192,30 +250,32 @@ class TestRedisClientWrapperKeyPrefix: wrapper.hkeys("hash:key") wrapper.hexists("hash:key", "field") - mock_client.hset.assert_called_once_with("enterprise-a:hash:key", "field", "value") - mock_client.hgetall.assert_called_once_with("enterprise-a:hash:key") - mock_client.hkeys.assert_called_once_with("enterprise-a:hash:key") - mock_client.hexists.assert_called_once_with("enterprise-a:hash:key", "field") + assert client.calls == [ + ("hset", ("enterprise-a:hash:key", "field", "value"), {}), + ("hgetall", ("enterprise-a:hash:key",), {}), + ("hkeys", ("enterprise-a:hash:key",), {}), + ("hexists", ("enterprise-a:hash:key", "field"), {}), + ] def test_wrapper_zadd_prefixes_sorted_set_name(self, config_overrides): - mock_client = MagicMock() + client = _RecordingRedisClient() wrapper = RedisClientWrapper() - wrapper.initialize(mock_client) + wrapper.initialize(client) # type: ignore[arg-type] config_overrides(REDIS_KEY_PREFIX="enterprise-a") wrapper.zadd("zset:key", {"member": 1}) - mock_client.zadd.assert_called_once() - args, kwargs = mock_client.zadd.call_args + method, args, kwargs = client.calls[0] + assert method == "zadd" assert args == ("enterprise-a:zset:key", {"member": 1}) assert kwargs["nx"] is False def test_wrapper_preserves_keys_when_prefix_is_empty(self, config_overrides): - mock_client = MagicMock() + client = _RecordingRedisClient() wrapper = RedisClientWrapper() - wrapper.initialize(mock_client) + wrapper.initialize(client) # type: ignore[arg-type] config_overrides(REDIS_KEY_PREFIX=" ") wrapper.get("plain:key") - mock_client.get.assert_called_once_with("plain:key") + assert client.calls == [("get", ("plain:key",), {})] diff --git a/api/tests/unit_tests/libs/test_device_flow_security.py b/api/tests/unit_tests/libs/test_device_flow_security.py new file mode 100644 index 00000000000..5583874f984 --- /dev/null +++ b/api/tests/unit_tests/libs/test_device_flow_security.py @@ -0,0 +1,25 @@ +import pytest + +from libs import device_flow_security +from libs import jws as jws_mod + + +def test_verify_approval_grant_rejects_missing_claim() -> None: + keyset = jws_mod.KeySet.from_entries( + {"key-1": b"test-secret-key-that-is-at-least-32-bytes"}, + active_kid="key-1", + ) + incomplete = jws_mod.sign( + keyset, + payload={ + "subject_email": "external@example.com", + "subject_issuer": "https://idp.example", + "user_code": "ABCD-EFGH", + "nonce": "nonce-1", + }, + aud=jws_mod.AUD_APPROVAL_GRANT, + ttl_seconds=60, + ) + + with pytest.raises(jws_mod.VerifyError): + device_flow_security.verify_approval_grant(keyset, incomplete) diff --git a/api/tests/unit_tests/libs/test_oauth_bearer.py b/api/tests/unit_tests/libs/test_oauth_bearer.py index e6566adcb9a..1e71354dba3 100644 --- a/api/tests/unit_tests/libs/test_oauth_bearer.py +++ b/api/tests/unit_tests/libs/test_oauth_bearer.py @@ -8,7 +8,8 @@ from unittest.mock import MagicMock import pytest -from libs.oauth_bearer import ResolvedRow, Scope, SubjectType, TokenType, _TokenTypeResolver +from constants.oauth_bearer import Scope, SubjectType, TokenType +from libs.oauth_bearer import ResolvedRow, _TokenTypeResolver def _row(account_id: uuid.UUID | None) -> ResolvedRow: diff --git a/api/tests/unit_tests/libs/test_oauth_bearer_rate_limit_ordering.py b/api/tests/unit_tests/libs/test_oauth_bearer_rate_limit_ordering.py index 845e520d946..f345e7f146c 100644 --- a/api/tests/unit_tests/libs/test_oauth_bearer_rate_limit_ordering.py +++ b/api/tests/unit_tests/libs/test_oauth_bearer_rate_limit_ordering.py @@ -4,7 +4,8 @@ from unittest.mock import MagicMock, patch import pytest -from libs.oauth_bearer import BearerAuthenticator, InvalidBearerError, TokenType +from constants.oauth_bearer import TokenType +from libs.oauth_bearer import BearerAuthenticator, InvalidBearerError def _authenticator_with(resolver) -> BearerAuthenticator: diff --git a/api/tests/unit_tests/repositories/test_account_repository.py b/api/tests/unit_tests/repositories/test_account_repository.py index f41c2fadda4..613e39352b5 100644 --- a/api/tests/unit_tests/repositories/test_account_repository.py +++ b/api/tests/unit_tests/repositories/test_account_repository.py @@ -369,3 +369,18 @@ def test_account_repository_updates_email_and_removes_integrations_atomically( assert persisted.email == "new@example.com" assert persisted.normalized_email == "new@example.com" assert sqlite_session.get(AccountIntegrate, integration_id) is None + + +def test_has_active_email_is_case_insensitive_and_excludes_inactive_accounts( + sqlite_session: Session, + sqlite_session_factory: sessionmaker[Session], +) -> None: + active = Account(name="Active", email="Active@Example.com", status=AccountStatus.ACTIVE) + banned = Account(name="Banned", email="banned@example.com", status=AccountStatus.BANNED) + sqlite_session.add_all([active, banned]) + sqlite_session.commit() + repository = SQLAlchemyAccountRepository(sqlite_session_factory) + + assert repository.has_active_email(" active@example.COM ") is True + assert repository.has_active_email("banned@example.com") is False + assert repository.has_active_email("") is False diff --git a/api/tests/unit_tests/repositories/test_oauth_device_token_repository.py b/api/tests/unit_tests/repositories/test_oauth_device_token_repository.py new file mode 100644 index 00000000000..f3787214392 --- /dev/null +++ b/api/tests/unit_tests/repositories/test_oauth_device_token_repository.py @@ -0,0 +1,236 @@ +from dataclasses import dataclass, field +from datetime import UTC, datetime +from functools import partial + +import pytest +from sqlalchemy import select +from sqlalchemy.orm import Session, sessionmaker + +from constants.oauth_bearer import TOKEN_CACHE_KEY_FMT +from extensions.ext_redis import RedisClientWrapper +from libs.oauth_bearer import sha256_hex +from models.oauth import OAuthAccessToken +from repositories.oauth_device_token_repository import SQLAlchemyOAuthDeviceTokenRepository +from services.entities.account_entities import AccountSnapshot +from services.oauth_device_adapters import OAuthDeviceTokenIssuanceGateway + + +@dataclass +class _RawRedis: + deleted: list[str | bytes] = field(default_factory=list) + + def delete(self, *names: str | bytes) -> None: + self.deleted.extend(names) + + +@dataclass +class _TTLPolicy: + days: int = 14 + workspace_ids: list[str | None] = field(default_factory=list) + + def ttl_days(self, workspace_id: str | None) -> int: + self.workspace_ids.append(workspace_id) + return self.days + + +def _account(account_id: str = "account-1", email: str = "ada@example.com") -> AccountSnapshot: + now = datetime.now(UTC) + return AccountSnapshot( + id=account_id, + name="Ada", + email=email, + avatar=None, + is_password_set=True, + interface_language="en-US", + interface_theme="light", + timezone="UTC", + last_login_at=now, + last_login_ip="127.0.0.1", + status="active", + initialized_at=now, + created_at=now, + ) + + +def _repository( + sqlite_session_factory: sessionmaker[Session], +) -> tuple[ + OAuthDeviceTokenIssuanceGateway, + SQLAlchemyOAuthDeviceTokenRepository, + _RawRedis, + _TTLPolicy, +]: + raw_redis = _RawRedis() + redis = RedisClientWrapper() + redis.initialize(raw_redis) # type: ignore[arg-type] + ttl_policy = _TTLPolicy() + repository = SQLAlchemyOAuthDeviceTokenRepository(session_factory=sqlite_session_factory, redis=redis) + return ( + OAuthDeviceTokenIssuanceGateway(tokens=repository, ttl_policy=ttl_policy), + repository, + raw_redis, + ttl_policy, + ) + + +def test_issue_account_token_persists_with_short_lived_session( + sqlite_session: Session, + sqlite_session_factory: sessionmaker[Session], +) -> None: + issuer, _repository_, _redis, ttl_policy = _repository(sqlite_session_factory) + + issued = issuer.issue_account_token( + account=_account(), + workspace_id="workspace-1", + client_id="difyctl", + device_label="Laptop", + ) + + assert issued.token.startswith("dfoa_") + assert ttl_policy.workspace_ids == ["workspace-1"] + sqlite_session.expire_all() + record = sqlite_session.get(OAuthAccessToken, issued.token_id) + assert record is not None + assert record.account_id == "account-1" + assert record.subject_issuer == "dify:account" + assert record.token_hash == sha256_hex(issued.token) + + +def test_issue_rotates_matching_live_token_and_invalidates_old_cache( + sqlite_session: Session, + sqlite_session_factory: sessionmaker[Session], +) -> None: + issuer, _repository_, redis, _ttl_policy = _repository(sqlite_session_factory) + issue_token = partial( + issuer.issue_account_token, + account=_account(), + workspace_id="workspace-1", + client_id="difyctl", + device_label="Laptop", + ) + + first = issue_token() + second = issue_token() + + sqlite_session.expire_all() + records = list(sqlite_session.scalars(select(OAuthAccessToken).order_by(OAuthAccessToken.created_at)).all()) + assert len(records) == 2 + assert sum(record.revoked_at is None for record in records) == 1 + assert records[-1].id == second.token_id + assert redis.deleted == [TOKEN_CACHE_KEY_FMT.format(hash=sha256_hex(first.token))] + + +def test_rollback_rotation_removes_new_token_and_restores_previous_session( + sqlite_session: Session, + sqlite_session_factory: sessionmaker[Session], +) -> None: + issuer, _repository_, redis, _ttl_policy = _repository(sqlite_session_factory) + issue_token = partial( + issuer.issue_account_token, + account=_account(), + workspace_id="workspace-1", + client_id="difyctl", + device_label="Laptop", + ) + first = issue_token() + second = issue_token() + + assert issuer.rollback_token(second) is True + + sqlite_session.expire_all() + first_record = sqlite_session.get(OAuthAccessToken, first.token_id) + assert first_record is not None + assert first_record.revoked_at is None + assert sqlite_session.get(OAuthAccessToken, second.token_id) is None + assert redis.deleted == [ + TOKEN_CACHE_KEY_FMT.format(hash=sha256_hex(first.token)), + TOKEN_CACHE_KEY_FMT.format(hash=sha256_hex(second.token)), + TOKEN_CACHE_KEY_FMT.format(hash=sha256_hex(first.token)), + ] + + +def test_rollback_does_not_restore_predecessor_after_later_rotation( + sqlite_session: Session, + sqlite_session_factory: sessionmaker[Session], +) -> None: + issuer, _repository_, _redis, _ttl_policy = _repository(sqlite_session_factory) + issue_token = partial( + issuer.issue_account_token, + account=_account(), + workspace_id="workspace-1", + client_id="difyctl", + device_label="Laptop", + ) + first = issue_token() + second = issue_token() + third = issue_token() + + assert issuer.rollback_token(second) is False + + sqlite_session.expire_all() + records = list(sqlite_session.scalars(select(OAuthAccessToken)).all()) + assert [record.id for record in records if record.revoked_at is None] == [third.token_id] + first_record = sqlite_session.get(OAuthAccessToken, first.token_id) + assert first_record is not None + assert first_record.revoked_at is not None + + +def test_list_account_sessions_pages_only_live_owned_tokens( + sqlite_session_factory: sessionmaker[Session], +) -> None: + issuer, repository, _redis, _ttl_policy = _repository(sqlite_session_factory) + account = _account() + for device_label in ("Laptop", "Desktop"): + issuer.issue_account_token( + account=account, + workspace_id="workspace-1", + client_id="difyctl", + device_label=device_label, + ) + issuer.issue_account_token( + account=_account("account-2", "grace@example.com"), + workspace_id="workspace-2", + client_id="difyctl", + device_label="Other", + ) + + page = repository.list_account_sessions(account_id="account-1", page=1, limit=1) + + assert page.total == 2 + assert len(page.items) == 1 + assert page.items[0].device_label in {"Laptop", "Desktop"} + + +def test_revoke_account_session_scopes_update_and_invalidates_cache( + sqlite_session: Session, + sqlite_session_factory: sessionmaker[Session], +) -> None: + issuer, repository, redis, _ttl_policy = _repository(sqlite_session_factory) + issued = issuer.issue_account_token( + account=_account(), + workspace_id="workspace-1", + client_id="difyctl", + device_label="Laptop", + ) + + assert repository.revoke_account_session(account_id="account-2", token_id=issued.token_id) is False + assert repository.revoke_account_session(account_id="account-1", token_id=issued.token_id) is True + + sqlite_session.expire_all() + record = sqlite_session.get(OAuthAccessToken, issued.token_id) + assert record is not None + assert record.revoked_at is not None + assert record.token_hash is None + assert redis.deleted == [TOKEN_CACHE_KEY_FMT.format(hash=sha256_hex(issued.token))] + + +def test_issue_external_token_rejects_empty_issuer(sqlite_session_factory: sessionmaker[Session]) -> None: + issuer, _repository_, _redis, _ttl_policy = _repository(sqlite_session_factory) + + with pytest.raises(ValueError, match="non-empty subject_issuer"): + issuer.issue_external_token( + subject_email="external@example.com", + subject_issuer=" ", + client_id="difyctl", + device_label="CLI", + ) diff --git a/api/tests/unit_tests/services/test_account_oauth_service.py b/api/tests/unit_tests/services/test_account_oauth_service.py index d2231064107..12948bc4b55 100644 --- a/api/tests/unit_tests/services/test_account_oauth_service.py +++ b/api/tests/unit_tests/services/test_account_oauth_service.py @@ -117,6 +117,9 @@ class FakeAccounts: def email_exists(self, email: str) -> bool: raise AssertionError(email) + def has_active_email(self, email: str) -> bool: + raise AssertionError(email) + def reset_email(self, account_id: str, *, expected_old_email: str, new_email: str) -> NoReturn: raise AssertionError((account_id, expected_old_email, new_email)) diff --git a/api/tests/unit_tests/services/test_account_service.py b/api/tests/unit_tests/services/test_account_service.py index 18cc02f1959..02165eae94b 100644 --- a/api/tests/unit_tests/services/test_account_service.py +++ b/api/tests/unit_tests/services/test_account_service.py @@ -2886,10 +2886,7 @@ class TestSessionInjectedGetters: assert AccountService.get_account_by_id("missing", session=sqlite_session) is None def test_get_account_by_email_returns_scalar_or_none(self, sqlite_session: Session) -> None: - """Plain getter — case-sensitive equality (callers needing the - case-insensitive existence check use - :meth:`has_active_account_with_email`). - """ + """Plain getter uses case-sensitive equality.""" account = Account(name="Alice", email="alice@example.com") sqlite_session.add(account) sqlite_session.commit() diff --git a/api/tests/unit_tests/services/test_feature_query_service.py b/api/tests/unit_tests/services/test_feature_query_service.py index edbc4a4a463..d66340e9a03 100644 --- a/api/tests/unit_tests/services/test_feature_query_service.py +++ b/api/tests/unit_tests/services/test_feature_query_service.py @@ -1,10 +1,13 @@ from unittest.mock import create_autospec +import pytest + from enums import DeploymentEdition from machinery.context import RequestContext from services.entities.feature_entities import ( FeatureModel, LicenseModel, + LicenseStatus, SystemFeatureModel, VectorSpaceLimitationModel, ) @@ -66,3 +69,23 @@ def test_workspace_id_queries_delegate_without_request_context() -> None: assert service.get_workspace_vector_space("workspace_123") is vector_space gateway.get_workspace_features.assert_called_once_with("workspace_123") gateway.get_vector_space.assert_called_once_with("workspace_123") + + +@pytest.mark.parametrize( + ("status", "expected"), + [ + (LicenseStatus.ACTIVE, True), + (LicenseStatus.EXPIRING, True), + (LicenseStatus.INACTIVE, False), + (LicenseStatus.EXPIRED, False), + ], +) +def test_enterprise_license_admission(status: LicenseStatus, expected: bool) -> None: + gateway = create_autospec(FeatureQueryGateway, instance=True, spec_set=True) + gateway.get_public_system_features.return_value = SystemFeatureModel( + deployment_edition=DeploymentEdition.ENTERPRISE, + license={"status": status}, + ) + service = FeatureQueryService(features=gateway, app_dsl_version="0.7.0") + + assert service.has_valid_enterprise_license() is expected diff --git a/api/tests/unit_tests/services/test_oauth_device_adapters.py b/api/tests/unit_tests/services/test_oauth_device_adapters.py new file mode 100644 index 00000000000..df9981e1b4c --- /dev/null +++ b/api/tests/unit_tests/services/test_oauth_device_adapters.py @@ -0,0 +1,189 @@ +from collections.abc import Callable, Mapping +from dataclasses import dataclass, field +from typing import override + +import pytest + +from extensions.ext_redis import RedisClientWrapper +from libs import jws +from libs.rate_limit import LIMIT_APPROVE_EXT_PER_EMAIL +from services.oauth_device_adapters import ( + DifyConfigOAuthDeviceSettings, + EnterpriseDeviceSSOService, + EnterpriseOAuthDeviceSSOGateway, + EnvironmentOAuthDeviceTokenTTLPolicy, + RedisExternalApprovalLimiter, +) +from services.oauth_device_contracts import OAuthDeviceSSOInitiationError + + +@dataclass +class _EnterpriseService(EnterpriseDeviceSSOService): + response: Mapping[str, object] | None + signed_states: list[str] = field(default_factory=list) + + @override + def initiate_device_flow_sso(self, signed_state: str) -> Mapping[str, object] | None: + self.signed_states.append(signed_state) + return self.response + + +class _FailingEnterpriseService(EnterpriseDeviceSSOService): + @override + def initiate_device_flow_sso(self, signed_state: str) -> Mapping[str, object] | None: + _ = signed_state + raise RuntimeError("enterprise unavailable") + + +@dataclass +class _NonceRedis: + values: dict[str | bytes, object] = field(default_factory=dict) + + def set( + self, + name: str | bytes, + value: object, + *, + nx: bool = False, + ex: int | None = None, + **_kwargs: object, + ) -> bool: + _ = ex + if nx and name in self.values: + return False + self.values[name] = value + return True + + def register_script(self, script: str) -> Callable[..., int]: + def execute( + *, + keys: list[str | bytes], + args: list[object], + client: object | None = None, + ) -> int: + _ = client + key = keys[0] + if "local current" in script: + current = self.values.get(key) + if current == args[0]: + return 1 + if current is not None: + return 0 + self.values[key] = args[0] + return 1 + if self.values.get(key) != args[0]: + return 0 + del self.values[key] + return 1 + + return execute + + +def _gateway( + config_overrides: Callable[..., None], + enterprise_service: EnterpriseDeviceSSOService, +) -> EnterpriseOAuthDeviceSSOGateway: + config_overrides(SECRET_KEY="test-secret-key-that-is-at-least-32-bytes") + return EnterpriseOAuthDeviceSSOGateway( + redis=RedisClientWrapper(), + enterprise_service=enterprise_service, + ) + + +def test_initiate_uses_injected_enterprise_operation(config_overrides: Callable[..., None]) -> None: + enterprise_service = _EnterpriseService(response={"url": "https://idp.example.com/authorize"}) + gateway = _gateway(config_overrides, enterprise_service) + + redirect_url = gateway.initiate( + user_code="ABCD-1234", + callback_url="https://api.example.com/openapi/v1/oauth/device/sso-complete", + ttl_seconds=300, + ) + + assert redirect_url == "https://idp.example.com/authorize" + assert len(enterprise_service.signed_states) == 1 + claims = jws.verify( + jws.KeySet.from_shared_secret(), + enterprise_service.signed_states[0], + expected_aud=jws.AUD_STATE_ENVELOPE, + ) + assert claims["user_code"] == "ABCD-1234" + assert claims["idp_callback_url"] == "https://api.example.com/openapi/v1/oauth/device/sso-complete" + + +@pytest.mark.parametrize("response", [None, {}, {"url": ""}, {"url": 123}]) +def test_initiate_rejects_missing_or_invalid_redirect_url( + config_overrides: Callable[..., None], + response: Mapping[str, object] | None, +) -> None: + gateway = _gateway(config_overrides, _EnterpriseService(response=response)) + + with pytest.raises(OAuthDeviceSSOInitiationError, match="sso_initiate_missing_url"): + gateway.initiate(user_code="ABCD-1234", callback_url="https://api.example.com/callback", ttl_seconds=300) + + +def test_initiate_maps_injected_enterprise_failure(config_overrides: Callable[..., None]) -> None: + gateway = _gateway(config_overrides, _FailingEnterpriseService()) + + with pytest.raises(OAuthDeviceSSOInitiationError, match="sso_initiate_failed"): + gateway.initiate(user_code="ABCD-1234", callback_url="https://api.example.com/callback", ttl_seconds=300) + + +def test_settings_adapter_owns_oauth_device_configuration(config_overrides: Callable[..., None]) -> None: + config_overrides( + inner_OPENAPI_KNOWN_CLIENT_IDS="client-a,client-b", + CONSOLE_WEB_URL="https://console.example.com", + CONSOLE_API_URL="https://api.example.com", + ) + + settings = DifyConfigOAuthDeviceSettings() + + assert settings.known_client_ids == frozenset({"client-a", "client-b"}) + assert settings.verification_base_url == "https://console.example.com" + assert settings.sso_base_url == "https://api.example.com" + + +def test_external_approval_limiter_owns_rate_limit_policy() -> None: + limiter = RedisExternalApprovalLimiter(redis=RedisClientWrapper()) + + assert limiter._rate_limiter.prefix == "rl:subject_email" + assert limiter._rate_limiter.max_attempts == LIMIT_APPROVE_EXT_PER_EMAIL.limit + assert limiter._rate_limiter.time_window == int(LIMIT_APPROVE_EXT_PER_EMAIL.window.total_seconds()) + + +def test_approval_nonce_release_only_removes_matching_reservation(config_overrides: Callable[..., None]) -> None: + config_overrides(REDIS_KEY_PREFIX="") + raw_redis = _NonceRedis() + redis = RedisClientWrapper() + redis.initialize(raw_redis) # type: ignore[arg-type] + gateway = EnterpriseOAuthDeviceSSOGateway( + redis=redis, + enterprise_service=_EnterpriseService(response=None), + ) + + assert gateway.reserve_approval_nonce("nonce-1", "reservation-1") is True + assert gateway.reserve_approval_nonce("nonce-1", "reservation-1") is True + assert gateway.reserve_approval_nonce("nonce-1", "reservation-2") is False + + gateway.release_approval_nonce("nonce-1", "reservation-2") + assert gateway.reserve_approval_nonce("nonce-1", "reservation-2") is False + + gateway.release_approval_nonce("nonce-1", "reservation-1") + assert gateway.reserve_approval_nonce("nonce-1", "reservation-2") is True + + +@pytest.mark.parametrize( + ("configured", "expected"), + [(None, 14), ("invalid", 14), ("0", 1), ("500", 365), ("30", 30)], +) +def test_token_ttl_policy_owns_environment_parsing( + monkeypatch: pytest.MonkeyPatch, + configured: str | None, + expected: int, +) -> None: + if configured is None: + monkeypatch.delenv("OAUTH_TTL_DAYS", raising=False) + else: + monkeypatch.setenv("OAUTH_TTL_DAYS", configured) + + assert EnvironmentOAuthDeviceTokenTTLPolicy().ttl_days("workspace-1") == expected diff --git a/api/tests/unit_tests/services/test_oauth_device_application_service.py b/api/tests/unit_tests/services/test_oauth_device_application_service.py new file mode 100644 index 00000000000..3f6e4991186 --- /dev/null +++ b/api/tests/unit_tests/services/test_oauth_device_application_service.py @@ -0,0 +1,1010 @@ +import logging +from dataclasses import dataclass, field +from datetime import UTC, datetime +from unittest.mock import patch + +import pytest + +from machinery.context import RequestContext +from services.entities.account_entities import AccountSnapshot +from services.oauth_device_application_service import OAuthDeviceApplicationService +from services.oauth_device_contracts import ( + AccessDeniedError, + AlreadyResolvedError, + ApprovalInProgressError, + ApprovalOutcomeUnknownError, + ApprovalSessionConsumedError, + ApprovalTransitionConfirmation, + AuthorizationPendingError, + DeviceFlowStatus, + DeviceMutation, + DeviceRequestContext, + DeviceStateLostError, + DeviceWorkspace, + ExpiredOrUnknownError, + ExpiredTokenError, + ExternalApprovalCSRFError, + ExternalApprovalGrant, + ExternalApprovalRateLimitError, + ExternalIdentityConflictError, + ExternalSubjectAssertion, + ExternalUserCodeMismatchError, + ExternalUserCodeNotFoundError, + InvalidApprovalSessionError, + InvalidSSOAssertionError, + InvalidTransitionError, + InvalidUserCodeError, + IssuedOAuthToken, + OAuthDeviceSession, + OAuthDeviceSessionNotFoundError, + OAuthDeviceSessionPage, + OAuthDeviceSSOConfigurationError, + OAuthDeviceSSOInitiationError, + OAuthDeviceTokenRotation, + PollPayload, + PollTooFastError, + SlowDownDecision, + StateNotFoundError, + UnsupportedClientError, +) + + +@dataclass +class FakeState: + client_id: str = "difyctl" + device_label: str = "CLI" + status: DeviceFlowStatus = DeviceFlowStatus.PENDING + token_id: str | None = None + transition_id: str | None = None + created_ip: str = "127.0.0.1" + poll_payload: PollPayload | None = None + + +@dataclass +class FakeStore: + states: dict[str, tuple[str, FakeState]] = field(default_factory=dict) + poll_decision: SlowDownDecision = SlowDownDecision.OK + approval_acquired: bool = True + approved_payload: PollPayload | None = None + approve_error: Exception | None = None + confirmation: ApprovalTransitionConfirmation | None = None + confirm_error: Exception | None = None + acquired: list[tuple[str, str, int]] = field(default_factory=list) + released: list[tuple[str, str]] = field(default_factory=list) + + def start(self, client_id: str, device_label: str, created_ip: str) -> tuple[str, str, int]: + _ = (client_id, device_label, created_ip) + return "device-1", "ABCD-EFGH", 900 + + def load_by_user_code(self, user_code: str) -> tuple[str, FakeState] | None: + return self.states.get(user_code) + + def load_by_device_code(self, device_code: str) -> FakeState | None: + found = next((state for code, state in self.states.values() if code == device_code), None) + return found + + def approve( + self, + device_code: str, + transition_id: str, + token_id: str, + poll_payload: PollPayload, + ) -> None: + if self.approve_error is not None: + raise self.approve_error + self.approved_payload = poll_payload + state = self.load_by_device_code(device_code) + if state is not None: + state.status = DeviceFlowStatus.APPROVED + state.transition_id = transition_id + state.token_id = token_id + state.poll_payload = poll_payload + + def deny(self, device_code: str, transition_id: str) -> None: + state = self.load_by_device_code(device_code) + if state is not None: + state.status = DeviceFlowStatus.DENIED + state.transition_id = transition_id + + def confirm_approval( + self, + device_code: str, + transition_id: str, + token_id: str, + ) -> ApprovalTransitionConfirmation: + if self.confirm_error is not None: + raise self.confirm_error + if self.confirmation is not None: + return self.confirmation + state = self.load_by_device_code(device_code) + if state is None: + return ApprovalTransitionConfirmation.UNKNOWN + if ( + state.status is DeviceFlowStatus.APPROVED + and state.transition_id == transition_id + and state.token_id == token_id + ): + return ApprovalTransitionConfirmation.PUBLISHED + return ApprovalTransitionConfirmation.NOT_PUBLISHED + + def consume_on_poll(self, device_code: str) -> FakeState | None: + for user_code, (code, state) in self.states.items(): + if code == device_code: + del self.states[user_code] + return state + return None + + def record_poll(self, device_code: str, interval_seconds: int) -> SlowDownDecision: + _ = (device_code, interval_seconds) + return self.poll_decision + + def try_acquire_approval(self, guard_id: str, owner_id: str, ttl_seconds: int) -> bool: + self.acquired.append((guard_id, owner_id, ttl_seconds)) + return self.approval_acquired + + def release_approval(self, guard_id: str, owner_id: str) -> None: + self.released.append((guard_id, owner_id)) + + +@dataclass +class FakeAccounts: + account: AccountSnapshot + active_external_emails: set[str] = field(default_factory=set) + + def get(self, account_id: str) -> AccountSnapshot | None: + return self.account if account_id == self.account.id else None + + def has_active_email(self, email: str) -> bool: + return email in self.active_external_emails + + +@dataclass +class FakeWorkspaces: + rows: tuple[DeviceWorkspace, ...] + error: Exception | None = None + + def list_for_device_flow(self, account_id: str) -> tuple[DeviceWorkspace, ...]: + _ = account_id + if self.error is not None: + raise self.error + return self.rows + + +@dataclass +class FakeTokens: + account_token: IssuedOAuthToken + external_token: IssuedOAuthToken + external_issues: int = 0 + account_issues: int = 0 + rollbacks: list[str] = field(default_factory=list) + sessions: tuple[OAuthDeviceSession, ...] = () + revoke_allowed: bool = True + + def issue_account_token( + self, + *, + account: AccountSnapshot, + workspace_id: str, + client_id: str, + device_label: str, + ) -> IssuedOAuthToken: + _ = (account, workspace_id, client_id, device_label) + self.account_issues += 1 + return self.account_token + + def issue_external_token( + self, + *, + subject_email: str, + subject_issuer: str, + client_id: str, + device_label: str, + ) -> IssuedOAuthToken: + _ = (subject_email, subject_issuer, client_id, device_label) + self.external_issues += 1 + return self.external_token + + def rollback_token(self, token: IssuedOAuthToken) -> bool: + self.rollbacks.append(token.token_id) + return True + + def list_account_sessions(self, *, account_id: str, page: int, limit: int) -> OAuthDeviceSessionPage: + _ = account_id + return OAuthDeviceSessionPage(page=page, limit=limit, total=len(self.sessions), items=self.sessions) + + def revoke_account_session(self, *, account_id: str, token_id: str) -> bool: + _ = (account_id, token_id) + return self.revoke_allowed + + +@dataclass +class FakeSSO: + assertion: ExternalSubjectAssertion + grant: ExternalApprovalGrant + assertion_nonce_available: bool = True + approval_nonce_available: bool = True + callback_url: str | None = None + reserved_nonces: list[tuple[str, str]] = field(default_factory=list) + released_nonces: list[tuple[str, str]] = field(default_factory=list) + + def initiate(self, *, user_code: str, callback_url: str, ttl_seconds: int) -> str: + _ = (user_code, ttl_seconds) + self.callback_url = callback_url + return "https://idp.example/authorize" + + def verify_assertion(self, assertion: str) -> ExternalSubjectAssertion: + _ = assertion + return self.assertion + + def mint_approval_grant( + self, + *, + issuer: str, + subject_email: str, + subject_issuer: str, + user_code: str, + ) -> str: + _ = (issuer, subject_email, subject_issuer, user_code) + return "approval-grant" + + def verify_approval_grant(self, token: str) -> ExternalApprovalGrant: + _ = token + return self.grant + + def consume_assertion_nonce(self, nonce: str) -> bool: + _ = nonce + return self.assertion_nonce_available + + def reserve_approval_nonce(self, nonce: str, reservation_id: str) -> bool: + self.reserved_nonces.append((nonce, reservation_id)) + return self.approval_nonce_available + + def release_approval_nonce(self, nonce: str, reservation_id: str) -> None: + self.released_nonces.append((nonce, reservation_id)) + + +@dataclass +class FakeLimiter: + limited: bool = False + recorded: list[str] = field(default_factory=list) + + def is_rate_limited(self, subject_email: str) -> bool: + _ = subject_email + return self.limited + + def record(self, subject_email: str) -> None: + self.recorded.append(subject_email) + + +@dataclass(frozen=True, slots=True) +class FakeSettings: + known_client_ids: frozenset[str] = frozenset({"difyctl"}) + verification_base_url: str | None = "https://console.example" + sso_base_url: str | None = "https://api.example" + + +@dataclass +class Harness: + service: OAuthDeviceApplicationService + store: FakeStore + accounts: FakeAccounts + workspaces: FakeWorkspaces + tokens: FakeTokens + sso: FakeSSO + limiter: FakeLimiter + + +def _account() -> AccountSnapshot: + now = datetime.now(UTC) + return AccountSnapshot( + id="account-1", + name="Ada", + email="ada@example.com", + avatar=None, + is_password_set=True, + interface_language="en-US", + interface_theme="light", + timezone="UTC", + last_login_at=now, + last_login_ip="127.0.0.1", + status="active", + initialized_at=now, + created_at=now, + ) + + +def _harness(*, settings: FakeSettings | None = None) -> Harness: + account = _account() + store = FakeStore(states={"ABCD-EFGH": ("device-1", FakeState())}) + accounts = FakeAccounts(account) + token = IssuedOAuthToken( + token="secret-token", + expires_at="2030-01-01T00:00:00+00:00", + rotation=OAuthDeviceTokenRotation( + token_id="token-1", + replaced_token_id=None, + replaced_token_hash=None, + ), + ) + tokens = FakeTokens(account_token=token, external_token=token) + external_assertion = ExternalSubjectAssertion( + subject_email="external@example.com", + subject_issuer="https://idp.example", + user_code="ABCD-EFGH", + nonce="assertion-nonce", + ) + approval_grant = ExternalApprovalGrant( + subject_email=external_assertion.subject_email, + subject_issuer=external_assertion.subject_issuer, + user_code=external_assertion.user_code, + nonce="approval-nonce", + csrf_token="csrf-token", + expires_at=datetime(2030, 1, 1, tzinfo=UTC), + ) + sso = FakeSSO(assertion=external_assertion, grant=approval_grant) + limiter = FakeLimiter() + workspaces = FakeWorkspaces((DeviceWorkspace(id="workspace-1", name="Primary", role="owner", current=True),)) + service = OAuthDeviceApplicationService( + store=store, + accounts=accounts, + workspaces=workspaces, + tokens=tokens, + sessions=tokens, + sso=sso, + external_approval_limiter=limiter, + settings=settings or FakeSettings(), + ) + return Harness( + service=service, + store=store, + accounts=accounts, + workspaces=workspaces, + tokens=tokens, + sso=sso, + limiter=limiter, + ) + + +def test_start_rejects_unknown_client_before_writing_state() -> None: + harness = _harness() + + with pytest.raises(UnsupportedClientError): + harness.service.start( + client_id="unknown", + device_label="CLI", + created_ip="127.0.0.1", + request_origin="https://api.example", + ) + + +def test_poll_preserves_slow_down_precedence() -> None: + harness = _harness() + harness.store.poll_decision = SlowDownDecision.SLOW_DOWN + + with pytest.raises(PollTooFastError): + harness.service.poll(device_code="device-1", poll_ip="127.0.0.1") + + +def test_cross_ip_audit_uses_real_token_id_from_poll_payload(caplog: pytest.LogCaptureFixture) -> None: + harness = _harness() + _device_code, state = harness.store.states["ABCD-EFGH"] + state.status = DeviceFlowStatus.APPROVED + state.token_id = "transition:approve-1:token-1" + state.poll_payload = { + "token": "secret-token", + "expires_at": "2030-01-01T00:00:00+00:00", + "subject_type": "account", + "account": {"id": "account-1"}, + "workspaces": [], + "default_workspace_id": None, + "token_id": "token-1", + } + + with caplog.at_level(logging.WARNING): + harness.service.poll(device_code="device-1", poll_ip="192.0.2.10") + + assert "token_id=token-1" in caplog.text + assert "transition:approve-1" not in caplog.text + + +def test_account_sessions_delegate_to_token_repository() -> None: + harness = _harness() + session = OAuthDeviceSession( + id="token-1", + prefix="dfoa_", + client_id="difyctl", + device_label="CLI", + created_at=None, + last_used_at=None, + expires_at=None, + ) + harness.tokens.sessions = (session,) + + result = harness.service.list_account_sessions(account_id="account-1", page=2, limit=10) + + assert result.items == (session,) + assert result.page == 2 + + +def test_revoke_account_session_hides_missing_or_foreign_token() -> None: + harness = _harness() + harness.tokens.revoke_allowed = False + + with pytest.raises(OAuthDeviceSessionNotFoundError): + harness.service.revoke_account_session(account_id="account-1", token_id="foreign-token") + + +def test_approve_mints_and_publishes_account_payload() -> None: + harness = _harness() + + result = harness.service.approve( + RequestContext("request-1", None, "account-1", "workspace-1"), + user_code="abcd-efgh", + ) + + assert result.status == "approved" + assert harness.store.approved_payload is not None + assert harness.store.approved_payload["default_workspace_id"] == "workspace-1" + assert harness.store.released == [harness.store.acquired[0][:2]] + + +def test_approve_does_not_release_guard_owned_by_another_request() -> None: + harness = _harness() + harness.store.approval_acquired = False + + with pytest.raises(ApprovalInProgressError): + harness.service.approve( + RequestContext("request-1", None, "account-1", "workspace-1"), + user_code="ABCD-EFGH", + ) + + assert harness.store.released == [] + + +def test_account_approve_rolls_back_rotation_when_state_publish_fails() -> None: + harness = _harness() + harness.store.approve_error = StateNotFoundError("device-1") + + with pytest.raises(DeviceStateLostError): + harness.service.approve( + RequestContext("request-1", None, "account-1", "workspace-1"), + user_code="ABCD-EFGH", + ) + + assert harness.tokens.rollbacks == ["token-1"] + assert harness.store.released == [harness.store.acquired[0][:2]] + + +def test_account_approve_keeps_token_when_ambiguous_publish_is_confirmed() -> None: + harness = _harness() + harness.store.approve_error = ConnectionError("response lost") + harness.store.confirmation = ApprovalTransitionConfirmation.PUBLISHED + + result = harness.service.approve( + RequestContext("request-1", None, "account-1", "workspace-1"), + user_code="ABCD-EFGH", + ) + + assert result.status == "approved" + assert harness.tokens.rollbacks == [] + + +def test_account_approve_does_not_compensate_unknown_publish_outcome() -> None: + harness = _harness() + harness.store.approve_error = ConnectionError("response lost") + harness.store.confirmation = ApprovalTransitionConfirmation.UNKNOWN + + with pytest.raises(ApprovalOutcomeUnknownError): + harness.service.approve( + RequestContext("request-1", None, "account-1", "workspace-1"), + user_code="ABCD-EFGH", + ) + + assert harness.tokens.rollbacks == [] + + +def test_account_approve_completes_workspace_query_before_rotating_token() -> None: + harness = _harness() + harness.workspaces.error = RuntimeError("workspace unavailable") + + with pytest.raises(RuntimeError, match="workspace unavailable"): + harness.service.approve( + RequestContext("request-1", None, "account-1", "workspace-1"), + user_code="ABCD-EFGH", + ) + + assert harness.tokens.account_issues == 0 + + +def test_sso_initiation_uses_configured_api_origin() -> None: + harness = _harness() + + result = harness.service.initiate_sso( + DeviceRequestContext("request-1", None), + user_code="abcd-efgh", + ) + + assert result.redirect_url == "https://idp.example/authorize" + assert harness.sso.callback_url == "https://api.example/openapi/v1/oauth/device/sso-complete" + + +def test_sso_completion_exchanges_valid_assertion_for_approval_grant() -> None: + harness = _harness() + + result = harness.service.complete_sso( + DeviceRequestContext("request-1", None), + inbound_error=None, + inbound_user_code=None, + assertion="signed-assertion", + ) + + assert result.error_code is None + assert result.user_code == "ABCD-EFGH" + assert result.approval_grant == "approval-grant" + + +@pytest.mark.parametrize("csrf_token", ["", "wrong"]) +def test_external_approval_rejects_bad_csrf_before_minting(csrf_token: str) -> None: + harness = _harness() + + with pytest.raises(ExternalApprovalCSRFError): + harness.service.approve_external( + DeviceRequestContext("request-1", None), + approval_grant="grant", + csrf_token=csrf_token, + user_code="ABCD-EFGH", + ) + + assert harness.tokens.external_issues == 0 + assert harness.sso.reserved_nonces == [] + assert harness.store.acquired == [] + + +def test_external_approval_publishes_external_payload() -> None: + harness = _harness() + + result = harness.service.approve_external( + DeviceRequestContext("request-1", None), + approval_grant="grant", + csrf_token="csrf-token", + user_code="abcd-efgh", + ) + + assert result.status == "approved" + assert harness.store.approved_payload is not None + assert harness.store.approved_payload["subject_type"] == "external_sso" + assert harness.store.approved_payload["account"] is None + assert harness.tokens.external_issues == 1 + assert len(harness.sso.reserved_nonces) == 1 + assert harness.sso.released_nonces == [] + + +def test_external_approval_releases_nonce_and_rolls_back_when_publish_fails() -> None: + harness = _harness() + harness.store.approve_error = StateNotFoundError("device-1") + + with pytest.raises(DeviceStateLostError): + harness.service.approve_external( + DeviceRequestContext("request-1", None), + approval_grant="grant", + csrf_token="csrf-token", + user_code="ABCD-EFGH", + ) + + assert harness.tokens.rollbacks == ["token-1"] + assert len(harness.sso.reserved_nonces) == 1 + assert harness.sso.released_nonces == harness.sso.reserved_nonces + assert harness.store.released == [harness.store.acquired[0][:2]] + + +def test_external_approval_keeps_recoverable_nonce_when_publish_outcome_is_unknown() -> None: + harness = _harness() + harness.store.approve_error = ConnectionError("response lost") + harness.store.confirmation = ApprovalTransitionConfirmation.UNKNOWN + + with pytest.raises(ApprovalOutcomeUnknownError): + harness.service.approve_external( + DeviceRequestContext("request-1", None), + approval_grant="grant", + csrf_token="csrf-token", + user_code="ABCD-EFGH", + ) + + assert harness.tokens.rollbacks == [] + assert len(harness.sso.reserved_nonces) == 1 + assert harness.sso.released_nonces == [] + + +def test_rotation_guard_is_shared_across_device_codes_for_the_same_token_identity() -> None: + harness = _harness() + context = RequestContext("request-1", None, "account-1", "workspace-1") + harness.service.approve(context, user_code="ABCD-EFGH") + harness.store.states["ABCD-EFGH"] = ("device-2", FakeState()) + + harness.service.approve(context, user_code="ABCD-EFGH") + + assert len(harness.store.acquired) == 2 + assert harness.store.acquired[0][0] == harness.store.acquired[1][0] + assert harness.store.acquired[0][1] != harness.store.acquired[1][1] + assert all(ttl == 60 for _, _, ttl in harness.store.acquired) + + +@pytest.mark.parametrize("base_url", [None, "https://console.example/"]) +def test_start_returns_device_codes_with_configured_or_request_origin(base_url: str | None) -> None: + harness = _harness(settings=FakeSettings(verification_base_url=base_url)) + + result = harness.service.start( + client_id="difyctl", device_label="CLI", created_ip="127.0.0.1", request_origin="https://request.example/" + ) + + assert result.device_code == "device-1" + assert result.user_code == "ABCD-EFGH" + assert result.verification_uri == (base_url or "https://request.example/") + "device" + assert result.expires_in == 900 + assert result.interval == 5 + + +@pytest.mark.parametrize( + ("state", "expected_error"), + [ + (None, ExpiredTokenError), + (FakeState(), AuthorizationPendingError), + (FakeState(status=DeviceFlowStatus.DENIED), AccessDeniedError), + (FakeState(status=DeviceFlowStatus.APPROVED), ExpiredTokenError), + ], + ids=["expired", "pending", "denied", "approved-without-token"], +) +def test_poll_rejects_unavailable_tokens(state: FakeState | None, expected_error: type[Exception]) -> None: + harness = _harness() + harness.store.states = {} if state is None else {"ABCD-EFGH": ("device-1", state)} + + with pytest.raises(expected_error): + harness.service.poll(device_code="device-1", poll_ip="127.0.0.1") + + if state is not None and state.status is DeviceFlowStatus.PENDING: + assert harness.store.states["ABCD-EFGH"][1] is state + assert harness.tokens.account_issues == harness.tokens.external_issues == 0 + + +def test_poll_does_not_return_a_token_consumed_by_another_request() -> None: + harness = _harness() + harness.service.approve(RequestContext("request-1", None, "account-1", "workspace-1"), user_code="ABCD-EFGH") + + with patch.object(harness.store, "consume_on_poll", return_value=None), pytest.raises(ExpiredTokenError): + harness.service.poll(device_code="device-1", poll_ip="127.0.0.1") + + +def test_approved_token_is_only_returned_once() -> None: + harness = _harness() + harness.service.approve(RequestContext("request-1", None, "account-1", "workspace-1"), user_code="ABCD-EFGH") + + result = harness.service.poll(device_code="device-1", poll_ip="127.0.0.1") + + assert result["token"] == "secret-token" + assert result["account"] == {"id": "account-1", "email": "ada@example.com", "name": "Ada"} + with pytest.raises(ExpiredTokenError): + harness.service.poll(device_code="device-1", poll_ip="127.0.0.1") + + +@pytest.mark.parametrize("status", [None, *DeviceFlowStatus]) +def test_lookup_only_accepts_pending_codes(status: DeviceFlowStatus | None) -> None: + harness = _harness() + harness.store.states = {} if status is None else {"ABCD-EFGH": ("device-1", FakeState(status=status))} + + result = harness.service.lookup(user_code=" abcd-efgh ") + + assert result.valid is (status is DeviceFlowStatus.PENDING) + assert result.expires_in_remaining == (900 if result.valid else 0) + assert result.client_id == (None if status is None else "difyctl") + + +@pytest.mark.parametrize("operation", ["approve", "deny"]) +@pytest.mark.parametrize("status", [None, DeviceFlowStatus.APPROVED, DeviceFlowStatus.DENIED]) +def test_account_decisions_reject_expired_or_resolved_codes(operation: str, status: DeviceFlowStatus | None) -> None: + harness = _harness() + harness.store.states = {} if status is None else {"ABCD-EFGH": ("device-1", FakeState(status=status))} + + expected_error = ExpiredOrUnknownError if status is None else AlreadyResolvedError + if operation == "approve": + with pytest.raises(expected_error): + harness.service.approve( + RequestContext("request-1", None, "account-1", "workspace-1"), user_code="ABCD-EFGH" + ) + else: + with pytest.raises(expected_error): + harness.service.deny(user_code="ABCD-EFGH") + + assert harness.tokens.account_issues == 0 + assert harness.store.acquired == [] + + +def test_denial_changes_poll_result_without_issuing_token() -> None: + harness = _harness() + + result = harness.service.deny(user_code=" abcd-efgh ") + + assert result.status == "denied" + with pytest.raises(AccessDeniedError): + harness.service.poll(device_code="device-1", poll_ip="127.0.0.1") + assert harness.tokens.account_issues == harness.tokens.external_issues == 0 + + +@pytest.mark.parametrize("error", [StateNotFoundError("expired"), InvalidTransitionError("already approved")]) +def test_deny_reports_state_lost_if_transition_races(error: Exception) -> None: + harness = _harness() + + with patch.object(harness.store, "deny", side_effect=error), pytest.raises(DeviceStateLostError): + harness.service.deny(user_code="ABCD-EFGH") + + assert harness.store.states["ABCD-EFGH"][1].status is DeviceFlowStatus.PENDING + + +def test_account_approval_requires_existing_account() -> None: + harness = _harness() + + with pytest.raises(DeviceStateLostError): + harness.service.approve( + RequestContext("request-1", None, "missing-account", "workspace-1"), user_code="ABCD-EFGH" + ) + + assert harness.tokens.account_issues == 0 + assert harness.store.acquired == [] + + +def _approve(harness: Harness, *, external: bool) -> DeviceMutation: + if external: + return harness.service.approve_external( + DeviceRequestContext("request-1", None), + approval_grant="grant", + csrf_token="csrf-token", + user_code="ABCD-EFGH", + ) + return harness.service.approve(RequestContext("request-1", None, "account-1", "workspace-1"), user_code="ABCD-EFGH") + + +@pytest.mark.parametrize("external", [False, True], ids=["account", "external"]) +@pytest.mark.parametrize("change", ["expired", "replaced", "approved", "denied"]) +def test_approval_revalidates_state_after_acquiring_guard(external: bool, change: str) -> None: + harness = _harness() + initial = harness.store.states["ABCD-EFGH"] + current: tuple[str, FakeState] | None + if change == "expired": + current = None + expected_error = ExternalUserCodeNotFoundError if external else DeviceStateLostError + elif change == "replaced": + current = ("another-device", FakeState()) + expected_error = DeviceStateLostError + else: + current = ("device-1", FakeState(status=DeviceFlowStatus(change))) + expected_error = AlreadyResolvedError + + with ( + patch.object(harness.store, "load_by_user_code", side_effect=[initial, current]), + pytest.raises(expected_error), + ): + _approve(harness, external=external) + + assert harness.tokens.account_issues == harness.tokens.external_issues == 0 + assert harness.sso.reserved_nonces == [] + assert harness.store.released == [harness.store.acquired[0][:2]] + + +@pytest.mark.parametrize("external", [False, True], ids=["account", "external"]) +def test_unpublished_approval_rolls_back_token_and_allows_retry(external: bool) -> None: + harness = _harness() + harness.store.approve_error = ConnectionError("publish failed") + harness.store.confirmation = ApprovalTransitionConfirmation.NOT_PUBLISHED + + with pytest.raises(ConnectionError, match="publish failed"): + _approve(harness, external=external) + + assert harness.tokens.rollbacks == ["token-1"] + assert harness.store.released == [harness.store.acquired[0][:2]] + if external: + assert len(harness.sso.released_nonces) == 1 + assert harness.sso.released_nonces == harness.sso.reserved_nonces + + +@pytest.mark.parametrize("external", [False, True], ids=["account", "external"]) +def test_confirmation_failure_preserves_potentially_published_token(external: bool) -> None: + harness = _harness() + harness.store.approve_error = ConnectionError("response lost") + harness.store.confirm_error = ConnectionError("confirmation unavailable") + + with pytest.raises(ApprovalOutcomeUnknownError): + _approve(harness, external=external) + + assert harness.tokens.rollbacks == [] + assert harness.sso.released_nonces == [] + assert harness.store.released == [harness.store.acquired[0][:2]] + + +def test_external_token_issuance_failure_releases_reserved_nonce() -> None: + harness = _harness() + + with ( + patch.object(harness.tokens, "issue_external_token", side_effect=RuntimeError("issuer unavailable")), + pytest.raises(RuntimeError, match="issuer unavailable"), + ): + _approve(harness, external=True) + + assert harness.tokens.rollbacks == [] + assert len(harness.sso.released_nonces) == 1 + assert harness.sso.released_nonces == harness.sso.reserved_nonces + assert harness.store.released == [harness.store.acquired[0][:2]] + + +def test_failed_token_compensation_does_not_prevent_nonce_and_guard_cleanup() -> None: + harness = _harness() + harness.store.approve_error = StateNotFoundError("expired") + + with ( + patch.object(harness.tokens, "rollback_token", side_effect=ConnectionError("database unavailable")), + pytest.raises(DeviceStateLostError), + ): + _approve(harness, external=True) + + assert len(harness.sso.released_nonces) == 1 + assert harness.sso.released_nonces == harness.sso.reserved_nonces + assert harness.store.released == [harness.store.acquired[0][:2]] + + +@pytest.mark.parametrize( + "rejection", ["missing-grant", "invalid-grant", "rate-limit", "user-code", "identity", "replay", "busy"] +) +def test_external_approval_rejections_do_not_issue_tokens(rejection: str) -> None: + harness = _harness() + errors = { + "missing-grant": InvalidApprovalSessionError, + "invalid-grant": InvalidApprovalSessionError, + "rate-limit": ExternalApprovalRateLimitError, + "user-code": ExternalUserCodeMismatchError, + "identity": ExternalIdentityConflictError, + "replay": ApprovalSessionConsumedError, + "busy": ApprovalInProgressError, + } + harness.limiter.limited = rejection == "rate-limit" + harness.sso.approval_nonce_available = rejection != "replay" + harness.store.approval_acquired = rejection != "busy" + if rejection == "identity": + harness.accounts.active_external_emails.add(harness.sso.grant.subject_email) + + with ( + patch.object( + harness.sso, + "verify_approval_grant", + return_value=harness.sso.grant, + side_effect=InvalidApprovalSessionError if rejection == "invalid-grant" else None, + ), + pytest.raises(errors[rejection]), + ): + harness.service.approve_external( + DeviceRequestContext("request-1", None), + approval_grant="" if rejection == "missing-grant" else "grant", + csrf_token="csrf-token", + user_code="OTHER-CODE" if rejection == "user-code" else "ABCD-EFGH", + ) + + assert harness.tokens.external_issues == 0 + assert harness.store.approved_payload is None + assert harness.sso.released_nonces == [] + if rejection in {"identity", "replay"}: + assert harness.store.released == [harness.store.acquired[0][:2]] + else: + assert harness.store.released == [] + assert harness.sso.reserved_nonces == [] + if rejection == "rate-limit": + assert harness.limiter.recorded == [] + + +@pytest.mark.parametrize( + "rejection", + ["missing-assertion", "invalid-assertion", "replayed-assertion", "expired-code", "resolved-code", "identity"], +) +def test_sso_completion_rejections_never_mint_approval_grant(rejection: str) -> None: + harness = _harness() + harness.sso.assertion_nonce_available = rejection != "replayed-assertion" + if rejection == "expired-code": + harness.store.states.clear() + elif rejection == "resolved-code": + harness.store.states["ABCD-EFGH"][1].status = DeviceFlowStatus.APPROVED + elif rejection == "identity": + harness.accounts.active_external_emails.add(harness.sso.assertion.subject_email) + + with ( + patch.object( + harness.sso, + "verify_assertion", + return_value=harness.sso.assertion, + side_effect=InvalidSSOAssertionError("invalid signature") if rejection == "invalid-assertion" else None, + ), + patch.object(harness.sso, "mint_approval_grant") as mint, + ): + result = harness.service.complete_sso( + DeviceRequestContext("request-1", None), + inbound_error=None, + inbound_user_code=None, + assertion=None if rejection == "missing-assertion" else "signed-assertion", + ) + + assert result.error_code == ("email_belongs_to_dify_account" if rejection == "identity" else "sso_failed") + assert result.approval_grant is None + mint.assert_not_called() + + +def test_sso_completion_preserves_upstream_error_without_consuming_assertion() -> None: + harness = _harness() + with patch.object(harness.sso, "verify_assertion") as verify: + result = harness.service.complete_sso( + DeviceRequestContext("request-1", None), + inbound_error="sso_failed", + inbound_user_code="ABCD-EFGH", + assertion="signed-assertion", + ) + + assert result.error_code == "sso_failed" + assert result.user_code == "ABCD-EFGH" + assert result.approval_grant is None + verify.assert_not_called() + + +@pytest.mark.parametrize("status", [None, DeviceFlowStatus.APPROVED, DeviceFlowStatus.DENIED]) +def test_sso_initiation_requires_pending_user_code(status: DeviceFlowStatus | None) -> None: + harness = _harness() + harness.store.states = {} if status is None else {"ABCD-EFGH": ("device-1", FakeState(status=status))} + + with patch.object(harness.sso, "initiate") as initiate, pytest.raises(InvalidUserCodeError): + harness.service.initiate_sso(DeviceRequestContext("request-1", None), user_code="ABCD-EFGH") + + initiate.assert_not_called() + + +def test_sso_initiation_requires_configured_callback_origin() -> None: + harness = _harness(settings=FakeSettings(sso_base_url=None)) + + with patch.object(harness.sso, "initiate") as initiate, pytest.raises(OAuthDeviceSSOConfigurationError): + harness.service.initiate_sso(DeviceRequestContext("request-1", None), user_code="ABCD-EFGH") + + initiate.assert_not_called() + + +def test_sso_initiation_rejects_empty_provider_redirect() -> None: + harness = _harness() + + with patch.object(harness.sso, "initiate", return_value=""), pytest.raises(OAuthDeviceSSOInitiationError): + harness.service.initiate_sso(DeviceRequestContext("request-1", None), user_code="ABCD-EFGH") + + +def test_sso_completion_reports_configuration_error_without_grant() -> None: + harness = _harness(settings=FakeSettings(sso_base_url=None)) + + result = harness.service.complete_sso( + DeviceRequestContext("request-1", None), + inbound_error=None, + inbound_user_code=None, + assertion="signed-assertion", + ) + + assert result.error_code == "sso_failed" + assert result.user_code == "ABCD-EFGH" + assert result.approval_grant is None + + +def test_approval_context_requires_verified_grant() -> None: + harness = _harness() + context = DeviceRequestContext("request-1", None) + + with pytest.raises(InvalidApprovalSessionError): + harness.service.get_approval_context(context, approval_grant="") + with ( + patch.object(harness.sso, "verify_approval_grant", side_effect=InvalidApprovalSessionError), + pytest.raises(InvalidApprovalSessionError), + ): + harness.service.get_approval_context(context, approval_grant="invalid") + + result = harness.service.get_approval_context(context, approval_grant="valid") + assert result.subject_email == harness.sso.grant.subject_email + assert result.subject_issuer == harness.sso.grant.subject_issuer + assert result.user_code == harness.sso.grant.user_code + assert result.csrf_token == harness.sso.grant.csrf_token + assert result.expires_at == harness.sso.grant.expires_at diff --git a/api/tests/unit_tests/services/test_oauth_device_flow.py b/api/tests/unit_tests/services/test_oauth_device_flow.py new file mode 100644 index 00000000000..0d511d9c5b4 --- /dev/null +++ b/api/tests/unit_tests/services/test_oauth_device_flow.py @@ -0,0 +1,191 @@ +import json +from dataclasses import dataclass, field + +from services.oauth_device_contracts import ApprovalTransitionConfirmation, DeviceFlowStatus, PollPayload +from services.oauth_device_flow import DeviceFlowRedis, DeviceFlowState + + +@dataclass +class _Redis: + response: str + delete_error: Exception | None = None + deleted: list[str] = field(default_factory=list) + script_sources: list[str] = field(default_factory=list) + script_calls: list[tuple[str, tuple[str, ...], tuple[object, ...]]] = field(default_factory=list) + + def set(self, name: str, value: str, *, nx: bool, ex: int) -> bool: + _ = (name, value, nx, ex) + return True + + def delete(self, name: str) -> None: + if self.delete_error is not None: + raise self.delete_error + self.deleted.append(name) + + def get(self, _name: str) -> str: + return self.response + + def register_script(self, script: str): + self.script_sources.append(script) + + def execute(*, keys, args=()): + self.script_calls.append((script, tuple(keys), tuple(args or ()))) + if "return raw" in script: + return self.response + return 1 + + return execute + + +def test_consume_on_poll_uses_one_cluster_safe_script_key() -> None: + state = DeviceFlowState( + user_code="ABCD-EFGH", + client_id="difyctl", + device_label="CLI", + status=DeviceFlowStatus.APPROVED, + ) + redis = _Redis(response=state.to_json()) + store = DeviceFlowRedis(redis) + + result = store.consume_on_poll("device-1") + + assert result is not None + assert result.user_code == "ABCD-EFGH" + consume_script = next(source for source in redis.script_sources if "return raw" in source) + assert "KEYS[2]" not in consume_script + assert [(keys, args) for source, keys, args in redis.script_calls if source == consume_script] == [ + (("device_code:device-1",), ()) + ] + assert redis.deleted == ["user_code:ABCD-EFGH"] + + +def test_user_code_cleanup_failure_does_not_lose_consumed_token() -> None: + state = DeviceFlowState( + user_code="ABCD-EFGH", + client_id="difyctl", + device_label="CLI", + status=DeviceFlowStatus.APPROVED, + ) + redis = _Redis(response=state.to_json(), delete_error=ConnectionError("redis unavailable")) + + result = DeviceFlowRedis(redis).consume_on_poll("device-1") + + assert result is not None + assert result.user_code == "ABCD-EFGH" + + +def test_approve_and_deny_share_atomic_pending_transition_script() -> None: + redis = _Redis(response="") + store = DeviceFlowRedis(redis) + poll_payload: PollPayload = { + "token": "secret-token", + "expires_at": "2030-01-01T00:00:00+00:00", + "subject_type": "account", + "account": {"id": "account-1"}, + "workspaces": [], + "default_workspace_id": None, + "token_id": "token-1", + } + + store.approve("device-1", "approve-1", "token-1", poll_payload) + store.deny("device-2", "deny-1") + + transition_script = next(source for source in redis.script_sources if "decoded.status ~= 'pending'" in source) + transition_calls = [(keys, args) for source, keys, args in redis.script_calls if source == transition_script] + assert "redis.call('SETEX'" in transition_script + assert "decoded.transition_id" not in transition_script + assert [args[0] for _, args in transition_calls] == ["approved", "denied"] + assert transition_calls[0][1][1] == "transition:approve-1:token-1" + assert [keys for keys, _ in transition_calls] == [("device_code:device-1",), ("device_code:device-2",)] + + +def test_approval_confirmation_reads_marker_from_existing_token_id_field() -> None: + state = DeviceFlowState( + user_code="ABCD-EFGH", + client_id="difyctl", + device_label="CLI", + status=DeviceFlowStatus.APPROVED, + token_id="transition:approve-1:token-1", + ) + store = DeviceFlowRedis(_Redis(response=state.to_json())) + + assert store.confirm_approval("device-1", "approve-1", "token-1") is ApprovalTransitionConfirmation.PUBLISHED + assert store.confirm_approval("device-1", "approve-2", "token-1") is ApprovalTransitionConfirmation.NOT_PUBLISHED + + +def test_approval_guard_release_compares_owner_before_delete() -> None: + redis = _Redis(response="") + store = DeviceFlowRedis(redis) + + assert store.try_acquire_approval("rotation-1", "owner-1", 900) is True + store.release_approval("rotation-1", "owner-1") + + release_script = next(source for source in redis.script_sources if "ARGV[1]" in source and "DEL" in source) + assert "redis.call('GET', KEYS[1]) == ARGV[1]" in release_script + assert redis.script_calls[-1][1:] == (("oauth_device:approval_guard:rotation-1",), ("owner-1",)) + + +def test_approved_state_stores_plaintext_token_only_in_poll_payload() -> None: + poll_payload: PollPayload = { + "token": "secret-token", + "expires_at": "2030-01-01T00:00:00+00:00", + "subject_type": "account", + "account": {"id": "account-1"}, + "workspaces": [], + "default_workspace_id": None, + "token_id": "token-1", + } + serialized = DeviceFlowState( + user_code="ABCD-EFGH", + client_id="difyctl", + device_label="CLI", + status=DeviceFlowStatus.APPROVED, + token_id="transition:approve-1:token-1", + poll_payload=poll_payload, + ).to_json() + data = json.loads(serialized) + + assert data["poll_payload"]["token"] == "secret-token" + assert serialized.count("secret-token") == 1 + assert {"subject_email", "account_id", "subject_issuer", "minted_token", "transition_id"}.isdisjoint(data) + + +def test_transition_marker_does_not_add_a_top_level_state_field() -> None: + serialized = DeviceFlowState( + user_code="ABCD-EFGH", + client_id="difyctl", + device_label="CLI", + status=DeviceFlowStatus.APPROVED, + token_id="transition:approve-1:token-1", + poll_payload=None, + ).to_json() + old_schema_fields = { + "user_code", + "client_id", + "device_label", + "status", + "token_id", + "created_at", + "created_ip", + "last_poll_at", + "poll_payload", + } + + assert set(json.loads(serialized)) <= old_schema_fields + + +def test_state_reader_tolerates_legacy_approval_fields() -> None: + legacy = { + "user_code": "ABCD-EFGH", + "client_id": "difyctl", + "device_label": "CLI", + "status": "pending", + "subject_email": "legacy@example.com", + "account_id": "account-1", + "subject_issuer": "dify:account", + "minted_token": "legacy-token", + } + + state = DeviceFlowState.from_json(json.dumps(legacy)) + + assert state.status is DeviceFlowStatus.PENDING diff --git a/api/tests/unit_tests/services/test_oauth_device_flow_mint.py b/api/tests/unit_tests/services/test_oauth_device_flow_mint.py index f16493b7774..2de88b2bbff 100644 --- a/api/tests/unit_tests/services/test_oauth_device_flow_mint.py +++ b/api/tests/unit_tests/services/test_oauth_device_flow_mint.py @@ -2,43 +2,79 @@ from __future__ import annotations -import uuid -from unittest.mock import MagicMock, patch +from datetime import UTC, datetime +from unittest.mock import create_autospec import pytest -from libs.oauth_bearer import TokenType -from services.oauth_device_flow import ACCOUNT_ISSUER_SENTINEL, UpsertOutcome, mint_oauth_token - -UPSERT = "services.oauth_device_flow._upsert" - - -def _mint(token_type: TokenType, subject_issuer: str | None) -> MagicMock: - outcome = UpsertOutcome(token_id=uuid.uuid4(), rotated=False, old_hash=None) - with patch(UPSERT, return_value=outcome) as upsert: - mint_oauth_token( - MagicMock(), - subject_email="who@example.com", - subject_issuer=subject_issuer, - account_id=str(uuid.uuid4()) if token_type.subject.bound_to_account else None, - client_id="client", - device_label="device", - token_type=token_type, - ttl_days=1, - session=MagicMock(), - ) - return upsert +from constants.oauth_bearer import TokenType +from libs.oauth_bearer import sha256_hex +from services.entities.account_entities import AccountSnapshot +from services.oauth_device_adapters import OAuthDeviceTokenIssuanceGateway +from services.oauth_device_application_service import OAuthDeviceTokenPersistence, OAuthDeviceTokenTTLPolicy +from services.oauth_device_contracts import ACCOUNT_ISSUER_SENTINEL, OAuthDeviceTokenRotation @pytest.mark.parametrize("token_type", list(TokenType), ids=lambda t: t.value) def test_the_issuer_follows_the_subject_binding(token_type: TokenType) -> None: + tokens = create_autospec(OAuthDeviceTokenPersistence, instance=True) + tokens.rotate_token.return_value = OAuthDeviceTokenRotation( + token_id="token-1", replaced_token_id=None, replaced_token_hash=None + ) + ttl_policy = create_autospec(OAuthDeviceTokenTTLPolicy, instance=True) + ttl_policy.ttl_days.return_value = 1 + gateway = OAuthDeviceTokenIssuanceGateway(tokens=tokens, ttl_policy=ttl_policy) + if token_type.subject.bound_to_account: - upsert = _mint(token_type, None) - assert upsert.call_args.kwargs["subject_issuer"] == ACCOUNT_ISSUER_SENTINEL - with pytest.raises(ValueError, match="ACCOUNT_ISSUER_SENTINEL"): - _mint(token_type, "https://idp.example.com") + account = AccountSnapshot( + id="account-1", + name="Caller", + email="who@example.com", + avatar=None, + is_password_set=True, + interface_language=None, + interface_theme=None, + timezone=None, + last_login_at=None, + last_login_ip=None, + status="active", + initialized_at=None, + created_at=datetime.now(UTC), + ) + issued = gateway.issue_account_token( + account=account, workspace_id="workspace-1", client_id="client", device_label="device" + ) + expected_issuer = ACCOUNT_ISSUER_SENTINEL + expected_account = account.id else: - upsert = _mint(token_type, "https://idp.example.com") - assert upsert.call_args.kwargs["subject_issuer"] == "https://idp.example.com" - with pytest.raises(ValueError, match="subject_issuer"): - _mint(token_type, " ") + issued = gateway.issue_external_token( + subject_email="who@example.com", + subject_issuer="https://idp.example.com", + client_id="client", + device_label="device", + ) + expected_issuer = "https://idp.example.com" + expected_account = None + + tokens.rotate_token.assert_called_once() + written = tokens.rotate_token.call_args.args[0] + assert written.subject_issuer == expected_issuer + assert written.account_id == expected_account + assert written.prefix == token_type.prefix + assert TokenType.for_token(issued.token) is token_type + assert written.token_hash == sha256_hex(issued.token) + + +@pytest.mark.parametrize("issuer", ["", " "]) +def test_external_issuer_is_required_before_persistence(issuer: str) -> None: + tokens = create_autospec(OAuthDeviceTokenPersistence, instance=True) + gateway = OAuthDeviceTokenIssuanceGateway( + tokens=tokens, ttl_policy=create_autospec(OAuthDeviceTokenTTLPolicy, instance=True) + ) + + with pytest.raises(ValueError, match="subject_issuer"): + gateway.issue_external_token( + subject_email="who@example.com", subject_issuer=issuer, client_id="client", device_label="device" + ) + + tokens.rotate_token.assert_not_called()