From 7446bb35a3e3abc1aa856413ee1d2a9bca067d5a Mon Sep 17 00:00:00 2001 From: wangxiaolei Date: Mon, 9 Nov 2026 02:26:30 +0000 Subject: [PATCH] chore: remove unnessary rbac call (#45857) --- api/services/enterprise/base.py | 57 +++++++++--- api/services/knowledge/dataset_service.py | 74 ++++++---------- .../services/enterprise/test_base.py | 87 +++++++++++++++++++ .../knowledge/test_dataset_service.py | 6 +- 4 files changed, 164 insertions(+), 60 deletions(-) diff --git a/api/services/enterprise/base.py b/api/services/enterprise/base.py index 5ddb51b0696..52ed3fd3b40 100644 --- a/api/services/enterprise/base.py +++ b/api/services/enterprise/base.py @@ -1,7 +1,10 @@ +import atexit import logging import os +import threading from collections.abc import Mapping -from typing import Any +from http.cookiejar import CookieJar, DefaultCookiePolicy +from typing import Any, ClassVar import httpx @@ -159,6 +162,34 @@ class EnterpriseRequest(BaseRequest): secret_key = os.environ.get("ENTERPRISE_API_SECRET_KEY", "ENTERPRISE_API_SECRET_KEY") secret_key_header = "Enterprise-Api-Secret-Key" + _rbac_client: ClassVar[httpx.Client | None] = None + _rbac_client_lock: ClassVar[threading.Lock] = threading.Lock() + + @classmethod + def _close_rbac_clients(cls) -> None: + """Close this process's RBAC connections during shutdown.""" + with cls._rbac_client_lock: + if cls._rbac_client is not None: + cls._rbac_client.close() + cls._rbac_client = None + + @classmethod + def _get_rbac_client(cls) -> httpx.Client: + """Return a process-local client for fixed endpoint/proxy settings. + + Initialize only in workers after process startup, never before forking. + Connections are shared; credentials and tenant headers remain request-local. + Response cookies are rejected so one account cannot affect another call. + The pool owns shutdown; callers must not close the returned client. + """ + with cls._rbac_client_lock: + if cls._rbac_client is None: + cls._rbac_client = httpx.Client( + mounts=cls._build_mounts(), + cookies=CookieJar(policy=DefaultCookiePolicy(allowed_domains=())), + ) + return cls._rbac_client + @classmethod def send_inner_rbac_request( cls, @@ -194,7 +225,6 @@ class EnterpriseRequest(BaseRequest): raise ValueError("ENTERPRISE_RBAC_API_URL is required when RBAC_ENABLED=true") url = f"{cls.rbac_base_url}{endpoint}" - mounts = cls._build_mounts() try: traceparent = generate_traceparent_header() @@ -204,19 +234,22 @@ class EnterpriseRequest(BaseRequest): except Exception: logger.debug("Failed to generate traceparent header", exc_info=True) - with httpx.Client(mounts=mounts) as client: - request_kwargs: dict[str, Any] = { - "json": json, - "params": params, - "headers": {"Content-Type": "application/json", cls.secret_key_header: cls.secret_key, **inner_headers}, - "timeout": timeout if timeout is not None else dify_config.ENTERPRISE_RBAC_REQUEST_TIMEOUT, - } - response = client.request(method, url, **request_kwargs) - if not response.is_success: - cls._handle_error_response(response) + client = cls._get_rbac_client() + request_kwargs: dict[str, Any] = { + "json": json, + "params": params, + "headers": {"Content-Type": "application/json", cls.secret_key_header: cls.secret_key, **inner_headers}, + "timeout": timeout if timeout is not None else dify_config.ENTERPRISE_RBAC_REQUEST_TIMEOUT, + } + response = client.request(method, url, **request_kwargs) + if not response.is_success: + cls._handle_error_response(response) return response.json() +atexit.register(EnterpriseRequest._close_rbac_clients) + + class EnterprisePluginManagerRequest(BaseRequest): base_url = os.environ.get("ENTERPRISE_PLUGIN_MANAGER_API_URL", "ENTERPRISE_PLUGIN_MANAGER_API_URL") secret_key = os.environ.get("ENTERPRISE_PLUGIN_MANAGER_API_SECRET_KEY", "ENTERPRISE_PLUGIN_MANAGER_API_SECRET_KEY") diff --git a/api/services/knowledge/dataset_service.py b/api/services/knowledge/dataset_service.py index 98878b8dfd5..10765dfc80d 100644 --- a/api/services/knowledge/dataset_service.py +++ b/api/services/knowledge/dataset_service.py @@ -133,15 +133,6 @@ class AutoDisableLogsDict(TypedDict): class DatasetService: - @staticmethod - def _can_manage_all_datasets(tenant_id: str, account_id: str) -> bool: - if not dify_config.RBAC_ENABLED: - return False - - permissions = enterprise_rbac_service.RBACService.MyPermissions.get(tenant_id, account_id) - workspace_permission_keys = getattr(getattr(permissions, "workspace", None), "permission_keys", []) or [] - return "dataset.create_and_management" in workspace_permission_keys - @staticmethod def get_datasets( page: int, @@ -164,6 +155,9 @@ class DatasetService: Service API callers skip account ACLs while retaining the tenant boundary. A resource-scoped grant further restricts the query to its authorized IDs. + Account callers must supply the resolved RBAC resource scope when restricted. + This method preserves that scope and pagination without fetching another + permission snapshot; include_all does not bypass the RBAC resource filter. """ query = select(Dataset).where(Dataset.tenant_id == tenant_id).order_by(Dataset.created_at.desc(), Dataset.id) @@ -193,46 +187,34 @@ class DatasetService: query = query.where(Dataset.id.in_(permitted_dataset_ids)) else: return [], 0 - else: - if dify_config.RBAC_ENABLED: - can_manage_all_datasets = include_all and DatasetService._can_manage_all_datasets( - str(tenant_id), str(user.id) - ) - should_show_all_datasets = include_all and can_manage_all_datasets - else: - should_show_all_datasets = user.current_role == TenantAccountRole.OWNER and include_all - + elif not dify_config.RBAC_ENABLED: + should_show_all_datasets = user.current_role == TenantAccountRole.OWNER and include_all if not should_show_all_datasets: - if dify_config.RBAC_ENABLED: - # RBAC mode: show all datasets. Permission control is enforced - # via permission_keys on each item and @rbac_permission_required decorators. - pass + # Keep legacy visibility rules when RBAC is disabled. + if permitted_dataset_ids and len(permitted_dataset_ids) > 0: + query = query.where( + sa.or_( + Dataset.permission == DatasetPermissionEnum.ALL_TEAM, + sa.and_( + Dataset.permission == DatasetPermissionEnum.ONLY_ME, + Dataset.maintainer == user.id, + ), + sa.and_( + Dataset.permission == DatasetPermissionEnum.PARTIAL_TEAM, + Dataset.id.in_(permitted_dataset_ids), + ), + ) + ) else: - # Keep legacy visibility rules when RBAC is disabled. - if permitted_dataset_ids and len(permitted_dataset_ids) > 0: - query = query.where( - sa.or_( - Dataset.permission == DatasetPermissionEnum.ALL_TEAM, - sa.and_( - Dataset.permission == DatasetPermissionEnum.ONLY_ME, - Dataset.maintainer == user.id, - ), - sa.and_( - Dataset.permission == DatasetPermissionEnum.PARTIAL_TEAM, - Dataset.id.in_(permitted_dataset_ids), - ), - ) - ) - else: - query = query.where( - sa.or_( - Dataset.permission == DatasetPermissionEnum.ALL_TEAM, - sa.and_( - Dataset.permission == DatasetPermissionEnum.ONLY_ME, - Dataset.maintainer == user.id, - ), - ) + query = query.where( + sa.or_( + Dataset.permission == DatasetPermissionEnum.ALL_TEAM, + sa.and_( + Dataset.permission == DatasetPermissionEnum.ONLY_ME, + Dataset.maintainer == user.id, + ), ) + ) elif service_api_grant is None: if dify_config.RBAC_ENABLED: # Without an account we cannot resolve RBAC resource visibility. diff --git a/api/tests/unit_tests/services/enterprise/test_base.py b/api/tests/unit_tests/services/enterprise/test_base.py index 76348705578..b303f274d79 100644 --- a/api/tests/unit_tests/services/enterprise/test_base.py +++ b/api/tests/unit_tests/services/enterprise/test_base.py @@ -5,8 +5,10 @@ generated and included in HTTP requests made by EnterpriseRequest. """ from collections.abc import Iterator +from http.cookiejar import CookieJar from unittest.mock import MagicMock, patch +import httpx import pytest from services.enterprise.base import EnterpriseRequest @@ -64,3 +66,88 @@ class TestTraceparentPropagation: assert headers["traceparent"] == expected_traceparent assert headers["Content-Type"] == "application/json" assert headers["Enterprise-Api-Secret-Key"] == "test-secret-key" + + +type PooledTransport = tuple[list[httpx.Request], list[httpx.Client]] + + +class TestRBACConnectionPooling: + @pytest.fixture + def pooled_transport(self, monkeypatch: pytest.MonkeyPatch) -> Iterator[PooledTransport]: + from services.enterprise import base + + EnterpriseRequest._close_rbac_clients() + requests: list[httpx.Request] = [] + + def handle(request: httpx.Request) -> httpx.Response: + requests.append(request) + return httpx.Response(200, json={"ok": True}, headers={"set-cookie": "session=first; Path=/"}) + + real_client = httpx.Client + clients: list[httpx.Client] = [] + + def build(*, mounts: dict[str, httpx.BaseTransport] | None, cookies: CookieJar) -> httpx.Client: + client = real_client(transport=httpx.MockTransport(handle), mounts=mounts, cookies=cookies) + clients.append(client) + return client + + monkeypatch.setattr(base.httpx, "Client", build) + monkeypatch.setattr(EnterpriseRequest, "rbac_base_url", "https://rbac.example.com") + monkeypatch.setattr(base, "generate_traceparent_header", lambda: None) + yield requests, clients + EnterpriseRequest._close_rbac_clients() + for client in clients: + client.close() + + def test_reuses_client_without_sharing_identity_or_cookies(self, pooled_transport: PooledTransport) -> None: + requests, clients = pooled_transport + assert EnterpriseRequest.send_inner_rbac_request( + "GET", "/permissions", tenant_id="tenant-a", account_id="account-a", timeout=3 + ) == {"ok": True} + EnterpriseRequest.send_inner_rbac_request("GET", "/permissions", tenant_id="tenant-b", timeout=7) + + assert len(clients) == 1 + assert not clients[0].is_closed + assert requests[0].headers["X-Inner-Tenant-Id"] == "tenant-a" + assert requests[0].headers["X-Inner-Account-Id"] == "account-a" + assert requests[1].headers["X-Inner-Tenant-Id"] == "tenant-b" + assert "X-Inner-Account-Id" not in requests[1].headers + assert "cookie" not in requests[1].headers + assert requests[0].extensions["timeout"]["read"] == 3 + assert requests[1].extensions["timeout"]["read"] == 7 + + def test_close_releases_shared_client(self, pooled_transport: PooledTransport) -> None: + client = EnterpriseRequest._get_rbac_client() + EnterpriseRequest._close_rbac_clients() + assert client.is_closed + assert EnterpriseRequest._rbac_client is None + + def test_concurrent_calls_share_one_client(self, pooled_transport: PooledTransport) -> None: + from concurrent.futures import ThreadPoolExecutor + + requests, clients = pooled_transport + + def send(index: int) -> object: + return EnterpriseRequest.send_inner_rbac_request( + "GET", "/permissions", tenant_id=f"tenant-{index}", account_id=f"account-{index}" + ) + + with ThreadPoolExecutor(max_workers=4) as executor: + assert list(executor.map(send, range(12))) == [{"ok": True}] * 12 + assert len(clients) == 1 + assert len(requests) == 12 + for request in requests: + assert request.headers["X-Inner-Tenant-Id"].removeprefix("tenant-") == request.headers[ + "X-Inner-Account-Id" + ].removeprefix("account-") + assert "cookie" not in request.headers + + def test_error_does_not_close_shared_client(self, pooled_transport: PooledTransport) -> None: + from services.errors.enterprise import EnterpriseAPIForbiddenError + + client = EnterpriseRequest._get_rbac_client() + with patch.object(client, "request", return_value=httpx.Response(403, json={"message": "denied"})): + with pytest.raises(EnterpriseAPIForbiddenError): + EnterpriseRequest.send_inner_rbac_request("GET", "/permissions", tenant_id="tenant-a") + assert not client.is_closed + assert EnterpriseRequest.send_inner_rbac_request("GET", "/permissions", tenant_id="tenant-a") == {"ok": True} diff --git a/api/tests/unit_tests/services/knowledge/test_dataset_service.py b/api/tests/unit_tests/services/knowledge/test_dataset_service.py index 214dbc4db4a..290ee854918 100644 --- a/api/tests/unit_tests/services/knowledge/test_dataset_service.py +++ b/api/tests/unit_tests/services/knowledge/test_dataset_service.py @@ -649,8 +649,9 @@ class TestDatasetServiceDatasetCases: assert total == 2 assert {dataset.id for dataset in datasets} == {creator_one.id, creator_two.id} + @pytest.mark.parametrize("include_all", [False, True]) def test_get_datasets_applies_rbac_resource_scope_and_maintainer_override( - self, application_tags, config_overrides: Callable[..., None], sqlite_session: Session + self, application_tags, config_overrides: Callable[..., None], sqlite_session: Session, include_all: bool ) -> None: user = _dataset_service_dataset_account(role=TenantAccountRole.NORMAL) accessible = _dataset_service_dataset_dataset( @@ -667,7 +668,7 @@ class TestDatasetServiceDatasetCases: with ( _dataset_service_dataset_patch( "services.knowledge.dataset_service.enterprise_rbac_service.RBACService.MyPermissions.get", - return_value=SimpleNamespace(workspace=SimpleNamespace(permission_keys=[])), + side_effect=AssertionError("List queries must reuse the caller-resolved RBAC scope"), ), ): datasets, total = _dataset_service_dataset_DatasetService.get_datasets( @@ -677,6 +678,7 @@ class TestDatasetServiceDatasetCases: tenant_id="tenant-1", user=user, accessible_dataset_ids=[accessible.id], + include_all=include_all, tags=application_tags, )