Просмотр исходного кода

refac: user permissions validation

Timothy Jaeryang Baek 5 месяцев назад
Родитель
Сommit
37f19f68eb

+ 11 - 3
backend/open_webui/apps/webui/routers/knowledge.py

@@ -1,7 +1,7 @@
 import json
 from typing import Optional, Union
 from pydantic import BaseModel
-from fastapi import APIRouter, Depends, HTTPException, status
+from fastapi import APIRouter, Depends, HTTPException, status, Request
 import logging
 
 from open_webui.apps.webui.models.knowledge import (
@@ -16,7 +16,7 @@ from open_webui.apps.retrieval.main import process_file, ProcessFileForm
 
 from open_webui.constants import ERROR_MESSAGES
 from open_webui.utils.utils import get_admin_user, get_verified_user
-from open_webui.utils.access_control import has_access
+from open_webui.utils.access_control import has_access, has_permission
 
 
 from open_webui.env import SRC_LOG_LEVELS
@@ -129,8 +129,16 @@ async def get_knowledge_list(user=Depends(get_verified_user)):
 
 @router.post("/create", response_model=Optional[KnowledgeResponse])
 async def create_new_knowledge(
-    form_data: KnowledgeForm, user=Depends(get_verified_user)
+    request: Request, form_data: KnowledgeForm, user=Depends(get_verified_user)
 ):
+    if user.role != "admin" and not has_permission(
+        user.id, "workspace.knowledge", request.app.state.config.USER_PERMISSIONS
+    ):
+        raise HTTPException(
+            status_code=status.HTTP_401_UNAUTHORIZED,
+            detail=ERROR_MESSAGES.UNAUTHORIZED,
+        )
+
     knowledge = Knowledges.insert_new_knowledge(user.id, form_data)
 
     if knowledge:

+ 9 - 1
backend/open_webui/apps/webui/routers/models.py

@@ -11,7 +11,7 @@ from fastapi import APIRouter, Depends, HTTPException, Request, status
 
 
 from open_webui.utils.utils import get_admin_user, get_verified_user
-from open_webui.utils.access_control import has_access
+from open_webui.utils.access_control import has_access, has_permission
 
 
 router = APIRouter()
@@ -47,9 +47,17 @@ async def get_base_models(user=Depends(get_admin_user)):
 
 @router.post("/create", response_model=Optional[ModelModel])
 async def create_new_model(
+    request: Request,
     form_data: ModelForm,
     user=Depends(get_verified_user),
 ):
+    if user.role != "admin" and not has_permission(
+        user.id, "workspace.models", request.app.state.config.USER_PERMISSIONS
+    ):
+        raise HTTPException(
+            status_code=status.HTTP_401_UNAUTHORIZED,
+            detail=ERROR_MESSAGES.UNAUTHORIZED,
+        )
 
     model = Models.get_model_by_id(form_data.id)
     if model:

+ 13 - 3
backend/open_webui/apps/webui/routers/prompts.py

@@ -2,9 +2,9 @@ from typing import Optional
 
 from open_webui.apps.webui.models.prompts import PromptForm, PromptModel, Prompts
 from open_webui.constants import ERROR_MESSAGES
-from fastapi import APIRouter, Depends, HTTPException, status
+from fastapi import APIRouter, Depends, HTTPException, status, Request
 from open_webui.utils.utils import get_admin_user, get_verified_user
-from open_webui.utils.access_control import has_access
+from open_webui.utils.access_control import has_access, has_permission
 
 router = APIRouter()
 
@@ -39,7 +39,17 @@ async def get_prompt_list(user=Depends(get_verified_user)):
 
 
 @router.post("/create", response_model=Optional[PromptModel])
-async def create_new_prompt(form_data: PromptForm, user=Depends(get_verified_user)):
+async def create_new_prompt(
+    request: Request, form_data: PromptForm, user=Depends(get_verified_user)
+):
+    if user.role != "admin" and not has_permission(
+        user.id, "workspace.prompts", request.app.state.config.USER_PERMISSIONS
+    ):
+        raise HTTPException(
+            status_code=status.HTTP_401_UNAUTHORIZED,
+            detail=ERROR_MESSAGES.UNAUTHORIZED,
+        )
+
     prompt = Prompts.get_prompt_by_command(form_data.command)
     if prompt is None:
         prompt = Prompts.insert_new_prompt(user.id, form_data)

+ 9 - 1
backend/open_webui/apps/webui/routers/tools.py

@@ -9,7 +9,7 @@ from open_webui.constants import ERROR_MESSAGES
 from fastapi import APIRouter, Depends, HTTPException, Request, status
 from open_webui.utils.tools import get_tools_specs
 from open_webui.utils.utils import get_admin_user, get_verified_user
-from open_webui.utils.access_control import has_access
+from open_webui.utils.access_control import has_access, has_permission
 
 
 router = APIRouter()
@@ -64,6 +64,14 @@ async def create_new_tools(
     form_data: ToolForm,
     user=Depends(get_verified_user),
 ):
+    if user.role != "admin" and not has_permission(
+        user.id, "workspace.knowledge", request.app.state.config.USER_PERMISSIONS
+    ):
+        raise HTTPException(
+            status_code=status.HTTP_401_UNAUTHORIZED,
+            detail=ERROR_MESSAGES.UNAUTHORIZED,
+        )
+
     if not form_data.id.isidentifier():
         raise HTTPException(
             status_code=status.HTTP_400_BAD_REQUEST,