feat(rbac): agent RBAC/ACL scenes and resource-locator check bundles (#41768)

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

api/controllers/common/rbac/__init__.py

@@ -0,0 +1,29 @@
+from core.rbac import RBACPermission, RBACResourceScope
+
+from .checks import RBAC_CHECKS_ATTR, RBACCheck, enforce_rbac_checks
+from .locators import (
+    AgentBehindApp,
+    AgentId,
+    DatasetByPipeline,
+    DatasetId,
+    PlainApp,
+    ResourceIdentity,
+    ResourceLocator,
+    Workspace,
+)
+
+__all__ = [
+    "RBAC_CHECKS_ATTR",
+    "AgentBehindApp",
+    "AgentId",
+    "DatasetByPipeline",
+    "DatasetId",
+    "PlainApp",
+    "RBACCheck",
+    "RBACPermission",
+    "RBACResourceScope",
+    "ResourceIdentity",
+    "ResourceLocator",
+    "Workspace",
+    "enforce_rbac_checks",
+]

api/controllers/common/rbac/checks.py

@@ -0,0 +1,65 @@
+from collections.abc import Mapping, Sequence
+from dataclasses import dataclass
+
+from flask import request
+from werkzeug.exceptions import Forbidden, NotFound
+
+from configs import dify_config
+from core.rbac import RBACPermission, RBACResourceScope
+from services.enterprise.rbac_service import RBACService
+
+from .locators import ResourceLocator, Workspace
+
+__all__ = ["RBAC_CHECKS_ATTR", "RBACCheck", "enforce_rbac_checks"]
+
+RBAC_CHECKS_ATTR = "rbac_checks"
+
+
+@dataclass(frozen=True)
+class RBACCheck:
+    scene: RBACPermission
+    locator: ResourceLocator
+
+    def __post_init__(self) -> None:
+        if isinstance(self.locator, Workspace) or self.scene.scope is RBACResourceScope.WORKSPACE:
+            return
+        if self.locator.scope is not self.scene.scope:
+            raise ValueError(
+                f"{self.scene} is a {self.scene.scope} scene but {self.locator!r} locates {self.locator.scope}"
+            )
+
+
+def enforce_rbac_checks(
+    *,
+    tenant_id: str,
+    account_id: str,
+    checks: Sequence[RBACCheck],
+    path_args: Mapping[str, object] | None = None,
+) -> None:
+    if not dify_config.RBAC_ENABLED:
+        return
+    try:
+        view_args = request.view_args or {}
+    except RuntimeError:
+        view_args = {}
+    merged: dict[str, object] = {**view_args, **(path_args or {})}
+    applicable = [
+        (check, identity) for check in checks if (identity := check.locator.locate(tenant_id, merged)) is not None
+    ]
+    if not applicable:
+        raise NotFound()
+    for check, identity in applicable:
+        is_workspace = identity.scope is RBACResourceScope.WORKSPACE
+        owner = None if is_workspace else check.locator.owner_id(tenant_id, identity)
+        if owner is not None and owner == account_id:
+            return
+        allowed = RBACService.CheckAccess.check(
+            tenant_id,
+            account_id,
+            scene=check.scene,
+            resource_type=None if is_workspace else identity.scope,
+            resource_id=None if is_workspace else identity.id,
+        )
+        if allowed:
+            return
+    raise Forbidden()

api/controllers/common/rbac/locators.py

@@ -0,0 +1,144 @@
+from collections.abc import Mapping
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Protocol, override, runtime_checkable
+
+from flask import g
+from werkzeug.exceptions import NotFound
+
+from core.rbac import RBACResourceScope
+from services.rbac_resource_service import RBACResourceService
+
+if TYPE_CHECKING:
+    from models.agent import Agent
+
+__all__ = [
+    "AgentBehindApp",
+    "AgentId",
+    "DatasetByPipeline",
+    "DatasetId",
+    "PlainApp",
+    "ResourceIdentity",
+    "ResourceLocator",
+    "Workspace",
+    "agent_binding",
+]
+
+_AGENT_BINDING_CACHE_KEY = "_rbac_agent_bindings"
+
+
+@dataclass(frozen=True)
+class ResourceIdentity:
+    scope: RBACResourceScope
+    id: str
+
+
+@runtime_checkable
+class ResourceLocator(Protocol):
+    scope: RBACResourceScope
+
+    def locate(self, tenant_id: str, path_args: Mapping[str, object]) -> ResourceIdentity | None: ...
+
+    def owner_id(self, tenant_id: str, identity: ResourceIdentity) -> str | None: ...
+
+
+def _required(path_args: Mapping[str, object], param: str) -> str:
+    value = path_args.get(param)
+    if not value:
+        raise ValueError(f"Missing {param} in request path")
+    return str(value)
+
+
+def agent_binding(tenant_id: str, app_id: str) -> "Agent | None":
+    cache: dict[tuple[str, str], Agent | None] = g.setdefault(_AGENT_BINDING_CACHE_KEY, {})
+    key = (tenant_id, app_id)
+    if key not in cache:
+        cache[key] = RBACResourceService.get_app_agent_binding(tenant_id, app_id)
+    return cache[key]
+
+
+class Workspace:
+    scope = RBACResourceScope.WORKSPACE
+
+    def locate(self, tenant_id: str, path_args: Mapping[str, object]) -> ResourceIdentity | None:
+        return ResourceIdentity(self.scope, "")
+
+    def owner_id(self, tenant_id: str, identity: ResourceIdentity) -> str | None:
+        return None
+
+    def __repr__(self) -> str:
+        return "Workspace()"
+
+
+class _ParamLocator:
+    scope: RBACResourceScope
+    default_param: str
+
+    def __init__(self, param: str | None = None):
+        self.param = param or self.default_param
+
+    def __repr__(self) -> str:
+        return f"{type(self).__name__}({self.param!r})"
+
+
+class AgentId(_ParamLocator):
+    scope = RBACResourceScope.AGENT
+    default_param = "agent_id"
+
+    def locate(self, tenant_id: str, path_args: Mapping[str, object]) -> ResourceIdentity | None:
+        return ResourceIdentity(self.scope, _required(path_args, self.param))
+
+    def owner_id(self, tenant_id: str, identity: ResourceIdentity) -> str | None:
+        return None
+
+
+class PlainApp(_ParamLocator):
+    scope = RBACResourceScope.APP
+    default_param = "app_id"
+
+    def locate(self, tenant_id: str, path_args: Mapping[str, object]) -> ResourceIdentity | None:
+        app_id = _required(path_args, self.param)
+        if agent_binding(tenant_id, app_id) is not None:
+            return None
+        return ResourceIdentity(self.scope, app_id)
+
+    def owner_id(self, tenant_id: str, identity: ResourceIdentity) -> str | None:
+        return RBACResourceService.get_app_maintainer(tenant_id, identity.id)
+
+
+class AgentBehindApp(_ParamLocator):
+    scope = RBACResourceScope.AGENT
+    default_param = "app_id"
+
+    def locate(self, tenant_id: str, path_args: Mapping[str, object]) -> ResourceIdentity | None:
+        from models.agent import AgentScope
+
+        binding = agent_binding(tenant_id, _required(path_args, self.param))
+        if binding is None or binding.scope == AgentScope.WORKFLOW_ONLY:
+            return None
+        return ResourceIdentity(self.scope, str(binding.id))
+
+    def owner_id(self, tenant_id: str, identity: ResourceIdentity) -> str | None:
+        return None
+
+
+class DatasetId(_ParamLocator):
+    scope = RBACResourceScope.DATASET
+    default_param = "dataset_id"
+
+    def locate(self, tenant_id: str, path_args: Mapping[str, object]) -> ResourceIdentity | None:
+        return ResourceIdentity(self.scope, _required(path_args, self.param))
+
+    def owner_id(self, tenant_id: str, identity: ResourceIdentity) -> str | None:
+        return RBACResourceService.get_dataset_maintainer(tenant_id, identity.id)
+
+
+class DatasetByPipeline(DatasetId):
+    default_param = "pipeline_id"
+
+    @override
+    def locate(self, tenant_id: str, path_args: Mapping[str, object]) -> ResourceIdentity | None:
+        pipeline_id = _required(path_args, self.param)
+        dataset_id = RBACResourceService.get_dataset_id_by_pipeline(tenant_id, pipeline_id)
+        if dataset_id is None:
+            raise NotFound("Dataset not found for pipeline")
+        return ResourceIdentity(self.scope, dataset_id)

api/controllers/common/wraps.py

@@ -1,181 +1,32 @@
 from collections.abc import Callable
 from functools import wraps
 
-from sqlalchemy import select
-from werkzeug.exceptions import Forbidden, NotFound
-
 from configs import dify_config
-from core.rbac import RBACPermission, RBACResourceScope
-from extensions.ext_database import db
+from controllers.common.rbac import RBAC_CHECKS_ATTR, RBACCheck, enforce_rbac_checks
 from libs.login import current_account_with_tenant
-from models.dataset import Dataset
-from models.model import App
-from services.agent.roster_service import AgentRosterService
-from services.enterprise.rbac_service import RBACService
 
-__all__ = ["RBACPermission", "RBACResourceScope", "enforce_rbac_access", "rbac_permission_required"]
+__all__ = ["rbac_permission_required"]
 
 
-def enforce_rbac_access(
-    *,
-    tenant_id: str,
-    account_id: str,
-    resource_type: RBACResourceScope,
-    scene: RBACPermission,
-    resource_required: bool = True,
-    path_args: dict[str, object] | None = None,
-) -> None:
-    """Enforce enterprise RBAC for an explicit account/tenant pair.
-
-    This is the flask-login-independent core of the RBAC gate so it can run
-    inside request-handling layers that resolve the caller themselves (e.g. the
-    openapi auth pipeline, which has the account on ``AuthData`` before
-    flask-login is mounted).
-
-    No-op when ``RBAC_ENABLED`` is ``False``. For resource-scoped checks the
-    resource ID is taken from ``path_args`` merged with ``request.view_args``;
-    resource ownership short-circuits the check. Raises ``Forbidden`` when
-    access is denied. For workspace-level checks pass ``resource_required=False``
-    so the RBAC request omits ``resource_id``.
-
-    Args:
-        tenant_id: The tenant the access is evaluated against.
-        account_id: The account requesting access.
-        resource_type: The :class:`RBACResourceScope` member (app/dataset/workspace).
-        scene: The :class:`RBACPermission` permission point, e.g. ``RBACPermission.APP_DELETE``.
-        resource_required: Whether a concrete resource ID is required.
-        path_args: Extra path arguments to merge with ``request.view_args``.
-    """
-    if not dify_config.RBAC_ENABLED:
-        return
-
-    check_resource_type = None if resource_type == RBACResourceScope.WORKSPACE else resource_type
-    resource_id = None
-    if resource_required and check_resource_type:
-        resource_id = _extract_resource_id(resource_type, tenant_id, path_args)
-        if _is_resource_owned_by_current_user(tenant_id, account_id, resource_type, resource_id):
-            return
-    allowed = RBACService.CheckAccess.check(
-        tenant_id,
-        account_id,
-        scene=scene,
-        resource_type=check_resource_type,
-        resource_id=resource_id,
-    )
-    if not allowed:
-        raise Forbidden()
-
-
-def rbac_permission_required[**P, R](
-    resource_type: RBACResourceScope,
-    scene: RBACPermission,
-    *,
-    resource_required: bool = True,
-) -> Callable[[Callable[P, R]], Callable[P, R]]:
-    """Check enterprise RBAC permissions for the current flask-login user.
-
-    When ``RBAC_ENABLED`` is ``False`` the decorator is a no-op and the
-    request passes through unchanged. When enabled it resolves the current
-    account/tenant and delegates to :func:`enforce_rbac_access`, raising
-    ``Forbidden`` if access is denied.
-
-    Args:
-        resource_type: The :class:`RBACResourceScope` member (app/dataset/workspace).
-        scene: The :class:`RBACPermission` permission point, e.g. ``RBACPermission.APP_DELETE``.
-        resource_required: Whether a concrete resource ID is required.
-    """
+def rbac_permission_required[**P, R](*checks: RBACCheck) -> Callable[[Callable[P, R]], Callable[P, R]]:
+    if not checks:
+        raise ValueError("rbac_permission_required requires at least one RBACCheck")
 
     def decorator(view: Callable[P, R]) -> Callable[P, R]:
-        @wraps(view)
+        @wraps(view, updated=())
         def decorated(*args: P.args, **kwargs: P.kwargs) -> R:
             if not dify_config.RBAC_ENABLED:
                 return view(*args, **kwargs)
-
             current_user, current_tenant_id = current_account_with_tenant()
-            enforce_rbac_access(
+            enforce_rbac_checks(
                 tenant_id=current_tenant_id,
                 account_id=current_user.id,
-                resource_type=resource_type,
-                scene=scene,
-                resource_required=resource_required,
+                checks=checks,
                 path_args=kwargs,
             )
             return view(*args, **kwargs)
 
+        setattr(decorated, RBAC_CHECKS_ATTR, checks)
         return decorated
 
     return decorator
-
-
-def _is_resource_owned_by_current_user(
-    tenant_id: str, account_id: str, resource_type: RBACResourceScope, resource_id: str
-) -> bool:
-    if resource_type == RBACResourceScope.APP:
-        maintainer = db.session.scalar(
-            select(App.maintainer).where(
-                App.id == resource_id,
-                App.tenant_id == tenant_id,
-                App.status == "normal",
-            )
-        )
-        return maintainer == account_id
-
-    if resource_type == RBACResourceScope.DATASET:
-        maintainer = db.session.scalar(
-            select(Dataset.maintainer).where(
-                Dataset.id == resource_id,
-                Dataset.tenant_id == tenant_id,
-            )
-        )
-        return maintainer == account_id
-
-    return False
-
-
-def _extract_resource_id(
-    resource_type: RBACResourceScope, tenant_id: str, path_args: dict[str, object] | None = None
-) -> str:
-    """Extract the resource ID from matched path arguments.
-
-    Some legacy route classes use neutral names such as ``resource_id`` for
-    app/dataset resources, and Agent routes carry ``agent_id``, which is
-    resolved to the App backing that Agent.
-    Dataset endpoints behind a rag-pipeline route contain ``pipeline_id``
-    instead of ``dataset_id``. In that case we look up the associated
-    ``Dataset`` row via ``Dataset.pipeline_id``.
-    """
-    from flask import request
-
-    view_args = request.view_args or {}
-    matched_args = {**view_args, **(path_args or {})}
-
-    if resource_type == RBACResourceScope.APP:
-        app_id = matched_args.get("app_id")
-        if app_id:
-            return str(app_id)  # pyrefly: ignore[unnecessary-type-conversion]
-
-        agent_id = matched_args.get("agent_id")
-        if agent_id:
-            authz_app_id = AgentRosterService(db.session).peek_authz_app_id(tenant_id=tenant_id, agent_id=str(agent_id))
-            return authz_app_id or str(agent_id)
-
-        resource_id = matched_args.get("resource_id")
-        if resource_id:
-            return str(resource_id)  # pyrefly: ignore[unnecessary-type-conversion]
-        raise ValueError("Missing app_id in request path")
-
-    if resource_type == RBACResourceScope.DATASET:
-        dataset_id = matched_args.get("dataset_id") or matched_args.get("resource_id")
-        if dataset_id:
-            return str(dataset_id)
-
-        pipeline_id = matched_args.get("pipeline_id")
-        if pipeline_id:
-            dataset = db.session.scalar(
-                select(Dataset).where(Dataset.pipeline_id == str(pipeline_id), Dataset.tenant_id == tenant_id)
-            )
-            if not dataset:
-                raise NotFound("Dataset not found for pipeline")
-            return str(dataset.id)  # pyrefly: ignore[unnecessary-type-conversion]
-        raise ValueError("Missing dataset_id or pipeline_id in request path")
-    raise ValueError(f"Unknown resource_type: {resource_type}")