enh/sec: verify folder data integrity

#17182
This commit is contained in:
Timothy Jaeryang Baek 2025-09-07 01:04:56 +04:00
parent d9f97c832c
commit 40e40d1ddd
4 changed files with 53 additions and 3 deletions

View File

@ -147,6 +147,15 @@ class FilesTable:
with get_db() as db:
return [FileModel.model_validate(file) for file in db.query(File).all()]
def check_access_by_user_id(self, id, user_id, permission="write") -> bool:
file = self.get_file_by_id(id)
if not file:
return False
if file.user_id == user_id:
return True
# Implement additional access control logic here as needed
return False
def get_files_by_ids(self, ids: list[str]) -> list[FileModel]:
with get_db() as db:
return [

View File

@ -129,7 +129,9 @@ class KnowledgeTable:
def get_knowledge_bases(self) -> list[KnowledgeUserModel]:
with get_db() as db:
all_knowledge = db.query(Knowledge).order_by(Knowledge.updated_at.desc()).all()
all_knowledge = (
db.query(Knowledge).order_by(Knowledge.updated_at.desc()).all()
)
user_ids = list(set(knowledge.user_id for knowledge in all_knowledge))
@ -149,6 +151,15 @@ class KnowledgeTable:
)
return knowledge_bases
def check_access_by_user_id(self, id, user_id, permission="write") -> bool:
knowledge = self.get_knowledge_by_id(id)
if not knowledge:
return False
if knowledge.user_id == user_id:
return True
user_group_ids = {group.id for group in Groups.get_groups_by_member_id(user_id)}
return has_access(user_id, permission, knowledge.access_control, user_group_ids)
def get_knowledge_bases_by_user_id(
self, user_id: str, permission: str = "write"
) -> list[KnowledgeUserModel]:
@ -158,7 +169,9 @@ class KnowledgeTable:
knowledge_base
for knowledge_base in knowledge_bases
if knowledge_base.user_id == user_id
or has_access(user_id, permission, knowledge_base.access_control, user_group_ids)
or has_access(
user_id, permission, knowledge_base.access_control, user_group_ids
)
]
def get_knowledge_by_id(self, id: str) -> Optional[KnowledgeModel]:

View File

@ -15,6 +15,9 @@ from open_webui.models.folders import (
Folders,
)
from open_webui.models.chats import Chats
from open_webui.models.files import Files
from open_webui.models.knowledge import Knowledges
from open_webui.config import UPLOAD_DIR
from open_webui.env import SRC_LOG_LEVELS
@ -45,6 +48,31 @@ router = APIRouter()
async def get_folders(user=Depends(get_verified_user)):
folders = Folders.get_folders_by_user_id(user.id)
# Verify folder data integrity
for folder in folders:
if folder.data:
if "files" in folder.data:
valid_files = []
for file in folder.data["files"]:
if file.get("type") == "file":
if Files.check_access_by_user_id(
file.get("id"), user.id, "read"
):
valid_files.append(file)
elif file.get("type") == "collection":
if Knowledges.check_access_by_user_id(
file.get("id"), user.id, "read"
):
valid_files.append(file)
else:
valid_files.append(file)
folder.data["files"] = valid_files
Folders.update_folder_by_id_and_user_id(
folder.id, user.id, FolderUpdateForm(data=folder.data)
)
return [
{
**folder.model_dump(),

File diff suppressed because one or more lines are too long