refactor(api): isolate OAuth device flow orchestration (#41361)

This commit is contained in:
Byron.wang authored and GitHub committed 2026-09-23 06:36:48 +00:00
1 parent 792d154608
commit 16d3302bff
62 files changed
+4616 -1395

No files matched your search

+18 -3
View File
@@ -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
+65
View File
@@ -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}"
+6 -2
View File
@@ -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:
+10 -19
View File
@@ -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.
+1 -1
View File
@@ -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
+1 -1
View File
@@ -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
+1 -1
View File
@@ -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 (
+1 -1
View File
@@ -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
@@ -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
+1 -1
View File
@@ -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
+3 -2
View File
@@ -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
+1 -1
View File
@@ -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/<string:app_id>/files")
@@ -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
+1 -1
View File
@@ -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
+86 -312
View File
@@ -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
+125 -323
View File
@@ -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
+1 -1
View File
@@ -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
+1 -1
View File
@@ -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
@@ -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)),
+28 -1
View File
@@ -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)
+14 -50
View File
@@ -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
+1 -59
View File
@@ -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"
+17 -1
View File
@@ -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,
@@ -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
@@ -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 = (
+2
View File
@@ -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: ...
+1 -19
View File
@@ -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."""
+7
View File
@@ -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
+336
View File
@@ -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}")
@@ -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,
},
)
+276
View File
@@ -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
+156 -304
View File
@@ -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}"
@@ -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
@@ -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
@@ -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:
@@ -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))
@@ -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
@@ -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 (
@@ -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
@@ -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
@@ -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
@@ -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"}
@@ -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
@@ -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()
@@ -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)
@@ -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)
@@ -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"
@@ -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
@@ -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
+85 -25
View File
@@ -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",), {})]
@@ -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)
@@ -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:
@@ -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:
@@ -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
@@ -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",
)
@@ -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))
@@ -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()
@@ -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
@@ -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
File diff suppressed because it is too large. Load diff
@@ -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
@@ -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()