fix(api): enforce workspace membership + role checks in auth pipeline (#36931)

Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>

api/controllers/openapi/auth/composition.py

@@ -1,11 +1,13 @@
 from __future__ import annotations
 
 from controllers.openapi.auth.conditions import (
-    EDITION_CE,
     EDITION_EE,
+    HAS_ALLOWED_ROLES,
     LOADED_APP_IS_PRIVATE,
     PATH_HAS_APP_ID,
     WEBAPP_AUTH_ENABLED,
+    WORKSPACE_MEMBERSHIP_REQUIRED,
+    WORKSPACE_SCOPED,
 )
 from controllers.openapi.auth.data import Edition
 from controllers.openapi.auth.flow import When
@@ -15,14 +17,18 @@ from controllers.openapi.auth.prepare import (
     load_app,
     load_app_access_mode,
     load_tenant,
+    load_tenant_from_request,
+    load_workspace_role,
     resolve_external_user,
 )
 from controllers.openapi.auth.verify import (
     check_acl,
-    check_app_access,
-    check_membership,
+    check_app_api_enabled,
     check_private_app_permission,
     check_scope,
+    check_workspace_member,
+    check_workspace_mismatch,
+    check_workspace_role,
 )
 from libs.oauth_bearer import TokenType
 
@@ -30,13 +36,17 @@ account_pipeline = AuthPipeline(
     prepare=[
         When(PATH_HAS_APP_ID, then=load_app),
         When(PATH_HAS_APP_ID, then=load_tenant),
-        load_account,  # all tokens here are account tokens
+        When(WORKSPACE_MEMBERSHIP_REQUIRED, then=load_tenant_from_request),
+        load_account,
+        When(WORKSPACE_SCOPED, then=load_workspace_role),
         When(PATH_HAS_APP_ID & EDITION_EE, then=load_app_access_mode),
     ],
     auth=[
+        When(PATH_HAS_APP_ID, then=check_app_api_enabled),
         check_scope,
-        When(EDITION_CE & PATH_HAS_APP_ID, then=check_membership),
-        When(EDITION_EE & PATH_HAS_APP_ID & ~WEBAPP_AUTH_ENABLED, then=check_app_access),
+        When(WORKSPACE_SCOPED, then=check_workspace_member),
+        When(PATH_HAS_APP_ID, then=check_workspace_mismatch),
+        When(HAS_ALLOWED_ROLES, then=check_workspace_role),
         When(PATH_HAS_APP_ID & EDITION_EE & WEBAPP_AUTH_ENABLED, then=check_acl),
         When(EDITION_EE & LOADED_APP_IS_PRIVATE, then=check_private_app_permission),
     ],
@@ -50,6 +60,7 @@ external_sso_pipeline = AuthPipeline(
         When(PATH_HAS_APP_ID, then=load_app_access_mode),
     ],
     auth=[
+        When(PATH_HAS_APP_ID, then=check_app_api_enabled),
         check_scope,
         When(PATH_HAS_APP_ID & WEBAPP_AUTH_ENABLED, then=check_acl),
         When(LOADED_APP_IS_PRIVATE, then=check_private_app_permission),

api/controllers/openapi/auth/conditions.py

@@ -50,4 +50,11 @@ EDITION_SAAS = config_cond(lambda: current_edition() == Edition.SAAS)
 
 WEBAPP_AUTH_ENABLED = config_cond(lambda: FeatureService.get_system_features().webapp_auth.enabled)
 
+WORKSPACE_MEMBERSHIP_REQUIRED = request_cond(lambda ctx: ctx.workspace_membership)
+HAS_ALLOWED_ROLES = request_cond(lambda ctx: ctx.allowed_roles is not None)
+
+# Caller must belong to the resolved tenant: either an app-scoped path (tenant
+# from the app) or an explicit workspace-membership path (tenant from request).
+WORKSPACE_SCOPED = PATH_HAS_APP_ID | WORKSPACE_MEMBERSHIP_REQUIRED
+
 LOADED_APP_IS_PRIVATE = data_cond(lambda data: data.app_access_mode == WebAppAccessMode.PRIVATE)

api/controllers/openapi/auth/data.py

@@ -9,7 +9,7 @@ from werkzeug.exceptions import InternalServerError
 
 from configs import dify_config
 from libs.oauth_bearer import Scope, TokenType
-from models.account import Account, Tenant
+from models.account import Account, Tenant, TenantAccountRole
 from models.model import App, EndUser
 from services.enterprise.enterprise_service import WebAppAccessMode
 
@@ -41,6 +41,8 @@ class RequestContext(BaseModel):
     token_type: TokenType
     scope: Scope | None = None
     path_params: dict[str, str]
+    workspace_membership: bool = False
+    allowed_roles: frozenset[TenantAccountRole] | None = None
 
 
 class AuthData(BaseModel):
@@ -56,10 +58,14 @@ class AuthData(BaseModel):
     external_identity: ExternalIdentity | None = None
     path_params: dict[str, str] = Field(default_factory=dict)
 
+    allowed_roles: frozenset[TenantAccountRole] | None = None
+
     app: App | None = None
     tenant: Tenant | None = None
     app_access_mode: WebAppAccessMode | None = None
 
+    tenant_role: TenantAccountRole | None = None
+
     caller: Account | EndUser | None = None
     caller_kind: Literal["account", "end_user"] | None = None
 

api/controllers/openapi/auth/pipeline.py

@@ -34,6 +34,7 @@ from libs.oauth_bearer import (
     reset_auth_ctx,
     set_auth_ctx,
 )
+from models.account import TenantAccountRole
 from services.feature_service import FeatureService, LicenseStatus
 
 
@@ -56,11 +57,15 @@ class AuthPipeline:
         view: Callable,
         *,
         scope: Scope | None,
+        workspace_membership: bool = False,
+        allowed_roles: frozenset[TenantAccountRole] | None = None,
     ) -> Any:
         req_ctx = RequestContext(
             token_type=identity.token_type,
             scope=scope,
             path_params=dict(request.view_args or {}),
+            workspace_membership=workspace_membership,
+            allowed_roles=allowed_roles,
         )
 
         data = AuthData(
@@ -71,6 +76,7 @@ class AuthPipeline:
             scopes=frozenset(identity.scopes),
             tenants=dict(identity.verified_tenants),
             required_scope=scope,
+            allowed_roles=allowed_roles,
             path_params=dict(req_ctx.path_params),
             external_identity=(
                 ExternalIdentity(email=identity.subject_email, issuer=identity.subject_issuer)
@@ -121,6 +127,41 @@ class PipelineRouter:
         scope: Scope | None = None,
         allowed_token_types: frozenset[TokenType] | None = None,
         edition: frozenset[Edition] | None = None,
+        workspace_membership: bool = False,
+        allowed_roles: frozenset[TenantAccountRole] | None = None,
+    ) -> Callable:
+        return self._make_decorator(
+            scope=scope,
+            allowed_token_types=allowed_token_types,
+            edition=edition,
+            workspace_membership=workspace_membership,
+            allowed_roles=allowed_roles,
+        )
+
+    def guard_workspace(
+        self,
+        *,
+        scope: Scope | None = None,
+        allowed_token_types: frozenset[TokenType] | None = None,
+        edition: frozenset[Edition] | None = None,
+        allowed_roles: frozenset[TenantAccountRole] | None = None,
+    ) -> Callable:
+        return self._make_decorator(
+            scope=scope,
+            allowed_token_types=allowed_token_types,
+            edition=edition,
+            workspace_membership=True,
+            allowed_roles=allowed_roles,
+        )
+
+    def _make_decorator(
+        self,
+        *,
+        scope: Scope | None,
+        allowed_token_types: frozenset[TokenType] | None,
+        edition: frozenset[Edition] | None,
+        workspace_membership: bool,
+        allowed_roles: frozenset[TenantAccountRole] | None,
     ) -> Callable:
         def decorator(view: Callable) -> Callable:
             @wraps(view)
@@ -132,6 +173,8 @@ class PipelineRouter:
                     scope=scope,
                     allowed_token_types=allowed_token_types,
                     edition=edition,
+                    workspace_membership=workspace_membership,
+                    allowed_roles=allowed_roles,
                 )
 
             return decorated
@@ -147,6 +190,8 @@ class PipelineRouter:
         scope: Scope | None,
         allowed_token_types: frozenset[TokenType] | None,
         edition: frozenset[Edition] | None,
+        workspace_membership: bool = False,
+        allowed_roles: frozenset[TenantAccountRole] | None = None,
     ) -> Any:
         # 404 not 403 — this edition doesn't expose the feature at all
         if edition is not None and current_edition() not in edition:
@@ -182,7 +227,15 @@ class PipelineRouter:
             if not license_checked and Edition.EE in route.required_edition:
                 _check_license()
 
-        return route.pipeline._run(identity, args, kwargs, view, scope=scope)
+        return route.pipeline._run(
+            identity,
+            args,
+            kwargs,
+            view,
+            scope=scope,
+            workspace_membership=workspace_membership,
+            allowed_roles=allowed_roles,
+        )
 
 
 def _should_run(step: Any, req_ctx: RequestContext, data: AuthData | None) -> bool:

api/controllers/openapi/auth/prepare.py

@@ -1,5 +1,8 @@
 from __future__ import annotations
 
+import uuid
+
+from flask import request
 from werkzeug.exceptions import Forbidden, InternalServerError, NotFound, Unauthorized
 
 from controllers.openapi.auth.data import AuthData
@@ -13,16 +16,18 @@ from services.enterprise.enterprise_service import EnterpriseService, WebAppAcce
 
 
 def load_app(data: AuthData) -> None:
+    if data.app is not None:
+        return
     app_id = data.path_params["app_id"]
     app = AppService.get_app_by_id(db.session, app_id)
     if not app or app.status != "normal":
         raise NotFound("app not found")
-    if not app.enable_api:
-        raise Forbidden("service_api_disabled")
     data.app = app
 
 
 def load_tenant(data: AuthData) -> None:
+    if data.tenant is not None:
+        return
     if data.app is None:
         raise InternalServerError("pipeline_invariant_violated: app not loaded before load_tenant")
     tenant = TenantService.get_tenant_by_id(db.session, str(data.app.tenant_id))
@@ -31,7 +36,25 @@ def load_tenant(data: AuthData) -> None:
     data.tenant = tenant
 
 
+def load_tenant_from_request(data: AuthData) -> None:
+    if data.tenant is not None:
+        return
+    workspace_id = data.path_params.get("workspace_id") or request.args.get("workspace_id")
+    if not workspace_id:
+        raise NotFound("workspace not found")
+    try:
+        uuid.UUID(workspace_id)
+    except ValueError:
+        raise NotFound("workspace not found")
+    tenant = TenantService.get_tenant_by_id(db.session, workspace_id)
+    if tenant is None or tenant.status == TenantStatus.ARCHIVE:
+        raise NotFound("workspace not found")
+    data.tenant = tenant
+
+
 def load_account(data: AuthData) -> None:
+    if data.caller is not None:
+        return
     account = AccountService.get_account_by_id(db.session, str(data.account_id))
     if account is None:
         raise Unauthorized("account not found")
@@ -41,6 +64,19 @@ def load_account(data: AuthData) -> None:
     data.caller_kind = "account"
 
 
+def load_workspace_role(data: AuthData) -> None:
+    if data.tenant_role is not None:
+        return
+    if data.tenant is None or data.account_id is None:
+        return
+    if data.caller is not None and getattr(data.caller, "status", None) != "active":
+        return
+    role = TenantService.get_account_role_in_tenant(db.session, str(data.account_id), str(data.tenant.id))
+    if role is None:
+        return
+    data.tenant_role = role
+
+
 def resolve_external_user(data: AuthData) -> None:
     if data.tenant is None or data.app is None or data.external_identity is None:
         raise Unauthorized("missing context for external user resolution")

api/controllers/openapi/auth/role_gate.py

@@ -1,77 +0,0 @@
-"""Workspace role gate.
-
-Layered on top of `validate_bearer` + `accept_subjects(SubjectType.ACCOUNT)`
-for routes whose access depends on the caller's `TenantAccountJoin.role`
-in the workspace named by the `workspace_id` path parameter.
-
-Usage::
-
-    @openapi_ns.route("/workspaces/<string:workspace_id>/members")
-    class Members(Resource):
-        @validate_bearer(accept=ACCEPT_USER_ANY)
-        @accept_subjects(SubjectType.ACCOUNT)
-        @require_workspace_role()  # any member
-        def get(self, workspace_id: str): ...
-
-        @validate_bearer(accept=ACCEPT_USER_ANY)
-        @accept_subjects(SubjectType.ACCOUNT)
-        @require_workspace_role(TenantAccountRole.OWNER, TenantAccountRole.ADMIN)
-        def post(self, workspace_id: str): ...
-
-Non-member callers get 404 (matching `GET /openapi/v1/workspaces/<id>`)
-so workspace IDs do not leak across tenants. A member without one of the
-allowed roles gets 403.
-"""
-
-from __future__ import annotations
-
-from collections.abc import Callable
-from functools import wraps
-from typing import TypeVar
-
-from werkzeug.exceptions import Forbidden, NotFound
-
-from extensions.ext_database import db
-from libs.oauth_bearer import try_get_auth_ctx
-from models.account import TenantAccountRole
-from services.account_service import TenantService
-
-F = TypeVar("F", bound=Callable[..., object])
-
-
-def require_workspace_role(*allowed_roles: TenantAccountRole) -> Callable[[F], F]:
-    """Gate a route on the caller's role in ``workspace_id``.
-
-    Pass no roles to require only membership. Pass one or more roles to
-    require the caller's role be in that set.
-    """
-
-    allowed = frozenset(allowed_roles)
-
-    def deco(fn: F) -> F:
-        @wraps(fn)
-        def wrapper(*args: object, **kwargs: object) -> object:
-            ctx = try_get_auth_ctx()
-            if ctx is None or ctx.account_id is None:
-                raise RuntimeError(
-                    "require_workspace_role called without account-bearer context; "
-                    "stack validate_bearer + accept_subjects(SubjectType.ACCOUNT) above it"
-                )
-
-            workspace_id = kwargs.get("workspace_id")
-            if not workspace_id:
-                raise RuntimeError("require_workspace_role expects a 'workspace_id' route parameter")
-
-            role = TenantService.get_account_role_in_tenant(db.session, str(ctx.account_id), str(workspace_id))
-
-            if role is None:
-                raise NotFound("workspace not found")
-
-            if allowed and role not in allowed:
-                raise Forbidden("insufficient workspace role")
-
-            return fn(*args, **kwargs)
-
-        return wrapper  # type: ignore[return-value]
-
-    return deco

api/controllers/openapi/auth/verify.py

@@ -1,10 +1,11 @@
 from __future__ import annotations
 
-from werkzeug.exceptions import Forbidden, Unauthorized
+from flask import request
+from werkzeug.exceptions import Forbidden, NotFound, UnprocessableEntity
 
 from controllers.openapi.auth.data import AuthData
 from extensions.ext_database import db
-from libs.oauth_bearer import Scope, TokenType, check_workspace_membership
+from libs.oauth_bearer import Scope, TokenType
 from services.account_service import AccountService, TenantService
 from services.enterprise.enterprise_service import EnterpriseService, WebAppAccessMode
 
@@ -17,17 +18,39 @@ def check_scope(data: AuthData) -> None:
     raise Forbidden("insufficient_scope")
 
 
-def check_membership(data: AuthData) -> None:
+def check_workspace_member(data: AuthData) -> None:
+    """Assert the caller belongs to the resolved tenant.
+
+    `load_workspace_role` stashes the membership role (None when the caller is
+    not a member or is inactive). A missing membership surfaces as 404, not
+    403, so workspace IDs don't leak across tenants.
+    """
+    if data.tenant_role is None:
+        raise NotFound("workspace not found")
+
+
+def check_workspace_mismatch(data: AuthData) -> None:
     if data.tenant is None:
-        raise Unauthorized("tenant unset")
-    if data.account_id is None:
-        raise Unauthorized("account_id unset")
-    check_workspace_membership(
-        account_id=data.account_id,
-        tenant_id=data.tenant.id,
-        token_hash=data.token_hash,
-        membership_cache=data.tenants,
-    )
+        return
+    request_workspace_id = data.path_params.get("workspace_id") or request.args.get("workspace_id")
+    if request_workspace_id and request_workspace_id != str(data.tenant.id):
+        raise UnprocessableEntity("workspace_id does not match app's workspace")
+
+
+def check_workspace_role(data: AuthData) -> None:
+    if data.allowed_roles is None:
+        return
+    if data.tenant_role is None:
+        raise NotFound("workspace not found")
+    if data.tenant_role not in data.allowed_roles:
+        raise Forbidden("insufficient workspace role")
+
+
+def check_app_api_enabled(data: AuthData) -> None:
+    if data.app is None:
+        return
+    if not data.app.enable_api:
+        raise Forbidden("service_api_disabled")
 
 
 def check_app_access(data: AuthData) -> None: