diff --git a/api/apps/restful_apis/chat_api.py b/api/apps/restful_apis/chat_api.py index 19fe442de0..9a4d5b1418 100644 --- a/api/apps/restful_apis/chat_api.py +++ b/api/apps/restful_apis/chat_api.py @@ -353,21 +353,32 @@ async def list_chats(): page_number = int(request.args.get("page", 0)) items_per_page = int(request.args.get("page_size", 0)) + tenants = TenantService.get_joined_tenants_by_user_id(current_user.id) + authorized_owner_ids = {member["tenant_id"] for member in tenants} + authorized_owner_ids.add(current_user.id) + if owner_ids: - chats, total = await thread_pool_exec( - DialogService.get_by_tenant_ids, - owner_ids, current_user.id, 0, 0, orderby, desc, keywords, **exact_filters, - ) - chats = [chat for chat in chats if chat["tenant_id"] in owner_ids] - total = len(chats) - if page_number and items_per_page: - start = (page_number - 1) * items_per_page - chats = chats[start : start + items_per_page] + requested_owner_ids = set(owner_ids) + unauthorized_owner_ids = requested_owner_ids - authorized_owner_ids + if unauthorized_owner_ids: + logging.warning( + "Rejected list_chats request: user=%s attempted unauthorized owner_ids=%s", + current_user.id, + sorted(unauthorized_owner_ids), + ) + return get_json_result( + data=False, + message="Only authorized owner_ids can be queried.", + code=RetCode.OPERATING_ERROR, + ) + effective_owner_ids = list(requested_owner_ids) else: - chats, total = await thread_pool_exec( - DialogService.get_by_tenant_ids, - [], current_user.id, page_number, items_per_page, orderby, desc, keywords, **exact_filters, - ) + effective_owner_ids = list(authorized_owner_ids) + + chats, total = await thread_pool_exec( + DialogService.get_by_tenant_ids, + effective_owner_ids, current_user.id, page_number, items_per_page, orderby, desc, keywords, **exact_filters, + ) return get_json_result( data={"chats": [_build_chat_response(chat) for chat in chats], "total": total} diff --git a/api/apps/restful_apis/search_api.py b/api/apps/restful_apis/search_api.py index c56d0ff834..7755704e4d 100644 --- a/api/apps/restful_apis/search_api.py +++ b/api/apps/restful_apis/search_api.py @@ -15,6 +15,7 @@ # import json +import logging from quart import Response, request from api.db.services.dialog_service import async_ask @@ -75,15 +76,31 @@ def list_searches(): owner_ids = request.args.getlist("owner_ids") try: - if not owner_ids: - tenants = [] - search_apps, total = SearchService.get_by_tenant_ids(tenants, current_user.id, page_number, items_per_page, orderby, desc, keywords) + tenants = TenantService.get_joined_tenants_by_user_id(current_user.id) + authorized_owner_ids = {member["tenant_id"] for member in tenants} + authorized_owner_ids.add(current_user.id) + + if owner_ids: + requested_owner_ids = set(owner_ids) + unauthorized_owner_ids = requested_owner_ids - authorized_owner_ids + if unauthorized_owner_ids: + logging.warning( + "Rejected list_searches request: user=%s attempted unauthorized owner_ids=%s", + current_user.id, + sorted(unauthorized_owner_ids), + ) + return get_json_result( + data=False, + message="Only authorized owner_ids can be queried.", + code=RetCode.OPERATING_ERROR, + ) + effective_owner_ids = list(requested_owner_ids) else: - search_apps, total = SearchService.get_by_tenant_ids(owner_ids, current_user.id, 0, 0, orderby, desc, keywords) - search_apps = [s for s in search_apps if s["tenant_id"] in owner_ids] - total = len(search_apps) - if page_number and items_per_page: - search_apps = search_apps[(page_number - 1) * items_per_page: page_number * items_per_page] + effective_owner_ids = list(authorized_owner_ids) + + search_apps, total = SearchService.get_by_tenant_ids( + effective_owner_ids, current_user.id, page_number, items_per_page, orderby, desc, keywords + ) return get_json_result(data={"search_apps": search_apps, "total": total}) except Exception as e: return server_error_response(e) diff --git a/test/testcases/test_http_api/test_chat_assistant_management/conftest.py b/test/testcases/test_http_api/test_chat_assistant_management/conftest.py index 330732db6d..60d5e43210 100644 --- a/test/testcases/test_http_api/test_chat_assistant_management/conftest.py +++ b/test/testcases/test_http_api/test_chat_assistant_management/conftest.py @@ -18,7 +18,7 @@ from common import batch_create_chat_assistants, delete_all_chat_assistants, get from utils import wait_for -@wait_for(30, 1, "Document parsing timeout") +@wait_for(200, 1, "Document parsing timeout") def condition(_auth, _dataset_id): res = list_documents(_auth, _dataset_id) for doc in res["data"]["docs"]: diff --git a/test/testcases/test_http_api/test_chat_assistant_management/test_chat_sdk_routes_unit.py b/test/testcases/test_http_api/test_chat_assistant_management/test_chat_sdk_routes_unit.py index 1094ae4292..fa0894f142 100644 --- a/test/testcases/test_http_api/test_chat_assistant_management/test_chat_sdk_routes_unit.py +++ b/test/testcases/test_http_api/test_chat_assistant_management/test_chat_sdk_routes_unit.py @@ -201,6 +201,7 @@ def _load_chat_module(monkeypatch): class _StubRetCode(int, Enum): SUCCESS = 0 DATA_ERROR = 102 + OPERATING_ERROR = 103 AUTHENTICATION_ERROR = 109 class _StubStatusEnum(str, Enum): @@ -376,6 +377,10 @@ def _load_chat_module(monkeypatch): def get_by_id(_tenant_id): return True, SimpleNamespace(llm_id="glm-4") + @staticmethod + def get_joined_tenants_by_user_id(_user_id): + return [{"tenant_id": "tenant-1"}, {"tenant_id": "team-tenant-2"}] + class _StubUserTenantService: @staticmethod def query(**_kwargs): @@ -886,6 +891,112 @@ def test_list_chats_keeps_zero_pagination_semantics(monkeypatch): assert len(res["data"]["chats"]) == 1 +@pytest.mark.p2 +def test_list_chats_rejects_unauthorized_owner_ids(monkeypatch): + module = _load_chat_module(monkeypatch) + monkeypatch.setattr( + module, + "request", + SimpleNamespace( + args=SimpleNamespace( + get=lambda key, default=None: { + "keywords": "", + "page": "0", + "page_size": "0", + "orderby": "create_time", + "desc": "true", + "id": None, + "name": None, + }.get(key, default), + getlist=lambda key: ["foreign-tenant-id"] if key == "owner_ids" else [], + ) + ), + ) + res = _run(module.list_chats.__wrapped__()) + assert res["code"] == module.RetCode.OPERATING_ERROR + assert "authorized owner_ids" in res["message"] + + +@pytest.mark.p2 +def test_list_chats_authorized_multi_tenant(monkeypatch): + module = _load_chat_module(monkeypatch) + captured = {} + monkeypatch.setattr( + module, + "request", + SimpleNamespace( + args=SimpleNamespace( + get=lambda key, default=None: { + "keywords": "", + "page": "1", + "page_size": "10", + "orderby": "create_time", + "desc": "true", + "id": None, + "name": None, + }.get(key, default), + getlist=lambda key: ["tenant-1", "team-tenant-2"] if key == "owner_ids" else [], + ) + ), + ) + + def _get_by_tenant_ids(owner_ids, user_id, *args, **kwargs): + captured["owner_ids"] = owner_ids + captured["user_id"] = user_id + return ( + [ + {**_DummyDialogRecord().to_dict(), "tenant_id": "tenant-1", "id": "c1"}, + {**_DummyDialogRecord().to_dict(), "tenant_id": "team-tenant-2", "id": "c2"}, + ], + 2, + ) + + monkeypatch.setattr(module.DialogService, "get_by_tenant_ids", _get_by_tenant_ids) + monkeypatch.setattr(module.KnowledgebaseService, "get_by_id", lambda _id: (True, _DummyKB())) + + res = _run(module.list_chats.__wrapped__()) + assert res["code"] == 0 + assert res["data"]["total"] == 2 + assert {c["id"] for c in res["data"]["chats"]} == {"c1", "c2"} + assert set(captured["owner_ids"]) == {"tenant-1", "team-tenant-2"} + assert captured["user_id"] == "tenant-1" + + +@pytest.mark.p2 +def test_list_chats_defaults_to_authorized_owner_ids_when_omitted(monkeypatch): + module = _load_chat_module(monkeypatch) + captured = {} + + monkeypatch.setattr( + module, + "request", + SimpleNamespace( + args=SimpleNamespace( + get=lambda key, default=None: { + "keywords": "", + "page": "1", + "page_size": "10", + "orderby": "create_time", + "desc": "true", + "id": None, + "name": None, + }.get(key, default), + getlist=lambda _key: [], + ) + ), + ) + + def _get_by_tenant_ids(owner_ids, *_args, **_kwargs): + captured["owner_ids"] = owner_ids + return ([], 0) + + monkeypatch.setattr(module.DialogService, "get_by_tenant_ids", _get_by_tenant_ids) + res = _run(module.list_chats.__wrapped__()) + + assert res["code"] == 0 + assert set(captured["owner_ids"]) == {"tenant-1", "team-tenant-2"} + + @pytest.mark.p2 def test_chat_session_create_and_update_guard_matrix_unit(monkeypatch): module = _load_chat_module(monkeypatch) diff --git a/test/testcases/test_web_api/test_search_app/test_search_routes_unit.py b/test/testcases/test_web_api/test_search_app/test_search_routes_unit.py index 3de9f3c156..9ea8f0f348 100644 --- a/test/testcases/test_web_api/test_search_app/test_search_routes_unit.py +++ b/test/testcases/test_web_api/test_search_app/test_search_routes_unit.py @@ -225,6 +225,10 @@ def _load_search_api(monkeypatch): def get_by_id(_tenant_id): return True, SimpleNamespace(id=_tenant_id) + @staticmethod + def get_joined_tenants_by_user_id(_user_id): + return [{"tenant_id": "tenant-1"}, {"tenant_id": "team-tenant-2"}] + class _UserTenantService: @staticmethod def query(**_kwargs): @@ -491,19 +495,30 @@ def test_list_and_delete_route_matrix_unit(monkeypatch): module, {"keywords": "k", "page": "1", "page_size": "1", "orderby": "create_time", "desc": "true", "owner_ids": ["tenant-1"]}, ) - monkeypatch.setattr( - module.SearchService, - "get_by_tenant_ids", - lambda _tenants, _uid, _page, _size, _orderby, _desc, _keywords: ( - [{"id": "x", "tenant_id": "tenant-1"}, {"id": "y", "tenant_id": "tenant-2"}], - 2, - ), - ) + + def _get_by_tenant_ids_filtered(tenants, _uid, page, size, _orderby, _desc, _keywords): + all_items = [{"id": "x", "tenant_id": "tenant-1"}, {"id": "y", "tenant_id": "tenant-1"}] + filtered = [item for item in all_items if item["tenant_id"] in set(tenants)] + total = len(filtered) + if page and size: + filtered = filtered[(page - 1) * size : page * size] + return filtered, total + + monkeypatch.setattr(module.SearchService, "get_by_tenant_ids", _get_by_tenant_ids_filtered) res = module.list_searches() assert res["code"] == 0 - assert res["data"]["total"] == 1 + assert res["data"]["total"] == 2 assert len(res["data"]["search_apps"]) == 1 - assert res["data"]["search_apps"][0]["tenant_id"] == "tenant-1" + + # list: unauthorized owner_ids + _set_request_args( + monkeypatch, + module, + {"keywords": "", "page": "0", "page_size": "10", "orderby": "create_time", "desc": "true", "owner_ids": ["other-tenant"]}, + ) + res = module.list_searches() + assert res["code"] == module.RetCode.OPERATING_ERROR + assert "authorized owner_ids" in res["message"] # list: exception def _raise_list(*_args, **_kwargs): @@ -542,3 +557,63 @@ def test_list_and_delete_route_matrix_unit(monkeypatch): res = module.delete_search(search_id="search-1") assert res["code"] == module.RetCode.EXCEPTION_ERROR assert "rm boom" in res["message"] + + +@pytest.mark.p2 +def test_list_searches_authorized_multi_tenant(monkeypatch): + module = _load_search_api(monkeypatch) + captured = {} + + _set_request_args( + monkeypatch, + module, + { + "keywords": "", + "page": "1", + "page_size": "10", + "orderby": "create_time", + "desc": "true", + "owner_ids": ["tenant-1", "team-tenant-2"], + }, + ) + + def _get_by_tenant_ids(owner_ids, user_id, *args, **kwargs): + captured["owner_ids"] = owner_ids + captured["user_id"] = user_id + return ( + [ + {"id": "s1", "tenant_id": "tenant-1"}, + {"id": "s2", "tenant_id": "team-tenant-2"}, + ], + 2, + ) + + monkeypatch.setattr(module.SearchService, "get_by_tenant_ids", _get_by_tenant_ids) + res = module.list_searches() + assert res["code"] == 0 + assert res["data"]["total"] == 2 + assert {s["id"] for s in res["data"]["search_apps"]} == {"s1", "s2"} + assert set(captured["owner_ids"]) == {"tenant-1", "team-tenant-2"} + assert captured["user_id"] == "tenant-1" + + +@pytest.mark.p2 +def test_list_searches_defaults_to_authorized_owner_ids_when_omitted(monkeypatch): + module = _load_search_api(monkeypatch) + captured = {} + + _set_request_args( + monkeypatch, + module, + {"keywords": "", "page": "1", "page_size": "10", "orderby": "create_time", "desc": "true"}, + ) + + def _get_by_tenant_ids(owner_ids, *_args, **_kwargs): + captured["owner_ids"] = owner_ids + return ([], 0) + + monkeypatch.setattr(module.SearchService, "get_by_tenant_ids", _get_by_tenant_ids) + res = module.list_searches() + + assert res["code"] == 0 + assert set(captured["owner_ids"]) == {"tenant-1", "team-tenant-2"}