Merge 81d4dfd0d6 into db2e94348f

packages/dbgpt-serve/src/dbgpt_serve/agent/hub/controller.py

@@ -1,8 +1,10 @@
 import logging
 from abc import ABC
-from typing import List
+from functools import cache
+from typing import List, Optional
 
-from fastapi import APIRouter, Body, File, UploadFile
+from fastapi import APIRouter, Body, Depends, File, HTTPException, UploadFile
+from fastapi.security.http import HTTPAuthorizationCredentials, HTTPBearer
 
 from dbgpt.agent.resource.tool.autogpt.plugins_util import scan_plugins
 from dbgpt.agent.resource.tool.pack import AutoGPTPluginToolPack
@@ -24,6 +26,57 @@ from .model.model import MyPluginVO, PluginHubVO
 router = APIRouter()
 logger = logging.getLogger(__name__)
 
+get_bearer_token = HTTPBearer(auto_error=False)
+_api_keys_config: Optional[str] = None
+
+
+def set_api_keys_config(api_keys: Optional[str]):
+    global _api_keys_config
+    _api_keys_config = api_keys
+
+
+@cache
+def _parse_api_keys(api_keys: str) -> List[str]:
+    if not api_keys:
+        return []
+    return [key.strip() for key in api_keys.split(",")]
+
+
+async def check_api_key(
+    auth: Optional[HTTPAuthorizationCredentials] = Depends(get_bearer_token),
+) -> Optional[str]:
+    global _api_keys_config
+    if _api_keys_config:
+        api_keys = _parse_api_keys(_api_keys_config)
+        if auth is None or (token := auth.credentials) not in api_keys:
+            raise HTTPException(
+                status_code=401,
+                detail={
+                    "error": {
+                        "message": "",
+                        "type": "invalid_request_error",
+                        "param": None,
+                        "code": "invalid_api_key",
+                    }
+                },
+            )
+        return token
+    return None
+
+
+def _validate_user_access(
+    requested_user: Optional[str], authenticated_user: Optional[str]
+) -> str:
+    global _api_keys_config
+    if _api_keys_config and authenticated_user:
+        if requested_user and requested_user != authenticated_user:
+            raise HTTPException(
+                status_code=403,
+                detail="Access denied: Cannot access another user's resources",
+            )
+        return requested_user or authenticated_user
+    return requested_user or "default"
+
 
 class ModulePlugin(BaseComponent, ABC):
     name = ComponentType.PLUGIN_HUB
@@ -34,6 +87,16 @@ class ModulePlugin(BaseComponent, ABC):
 
     def init_app(self, system_app: SystemApp):
         system_app.app.include_router(router, prefix="/api", tags=["Agent"])
+        try:
+            global_api_keys = system_app.config.get("dbgpt.app.global.api_keys")
+            if global_api_keys:
+                set_api_keys_config(global_api_keys)
+            elif hasattr(system_app, "config") and hasattr(
+                system_app.config, "API_KEYS"
+            ):
+                set_api_keys_config(system_app.config.API_KEYS)
+        except Exception:
+            pass
 
     def refresh_plugins(self):
         self.plugins = scan_plugins(PLUGINS_DIR)
@@ -89,18 +152,26 @@ async def get_agent_list(filter: PagenationFilter[PluginHubFilter] = Body()):
 
 
 @router.post("/v1/agent/my", response_model=Result[List[MyPluginVO]])
-async def my_agents(user: str = None):
+async def my_agents(
+    user: str = None, authenticated_user: Optional[str] = Depends(check_api_key)
+):
     logger.info(f"my_agents:{user}")
-    agents = plugin_hub.get_my_plugin(user)
+    effective_user = _validate_user_access(user, authenticated_user)
+    agents = plugin_hub.get_my_plugin(effective_user)
     agent_dicts = MyPluginEntity.to_vo(agents)
     return Result.succ(agent_dicts)
 
 
 @router.post("/v1/agent/install", response_model=Result[str])
-async def agent_install(plugin_name: str, user: str = None):
+async def agent_install(
+    plugin_name: str,
+    user: str = None,
+    authenticated_user: Optional[str] = Depends(check_api_key),
+):
     logger.info(f"agent_install:{plugin_name},{user}")
+    effective_user = _validate_user_access(user, authenticated_user)
     try:
-        plugin_hub.install_plugin(plugin_name, user)
+        plugin_hub.install_plugin(plugin_name, effective_user)
 
         module_plugin.refresh_plugins()
 
@@ -111,10 +182,15 @@ async def agent_install(plugin_name: str, user: str = None):
 
 
 @router.post("/v1/agent/uninstall", response_model=Result[str])
-async def agent_uninstall(plugin_name: str, user: str = None):
+async def agent_uninstall(
+    plugin_name: str,
+    user: str = None,
+    authenticated_user: Optional[str] = Depends(check_api_key),
+):
     logger.info(f"agent_uninstall:{plugin_name},{user}")
+    effective_user = _validate_user_access(user, authenticated_user)
     try:
-        plugin_hub.uninstall_plugin(plugin_name, user)
+        plugin_hub.uninstall_plugin(plugin_name, effective_user)
 
         module_plugin.refresh_plugins()
         return Result.succ(None)
@@ -124,10 +200,15 @@ async def agent_uninstall(plugin_name: str, user: str = None):
 
 
 @router.post("/v1/personal/agent/upload", response_model=Result[str])
-async def personal_agent_upload(doc_file: UploadFile = File(...), user: str = None):
+async def personal_agent_upload(
+    doc_file: UploadFile = File(...),
+    user: str = None,
+    authenticated_user: Optional[str] = Depends(check_api_key),
+):
     logger.info(f"personal_agent_upload:{doc_file.filename},{user}")
+    effective_user = _validate_user_access(user, authenticated_user)
     try:
-        await plugin_hub.upload_my_plugin(doc_file, user)
+        await plugin_hub.upload_my_plugin(doc_file, effective_user)
         module_plugin.refresh_plugins()
         return Result.succ(None)
     except Exception as e: