github/DB-GPT
1
0
Fork 0
mirror of https://github.com/csunny/DB-GPT.git synced 2025-07-28 06:17:14 +00:00
Code Issues Packages Projects Releases Wiki Activity

Fix Critical Directory Traversal Vulnerability (#2098)

Browse Source
This commit is contained in:
Raphael 2024-11-17 17:46:33 -08:00 committed by GitHub
parent b392d51adf
commit 780ce803e3
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 65 additions and 48 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

113
dbgpt/app/knowledge/api.py
View File

@ -3,8 +3,9 @@ import os
import shutil import shutil
import tempfile import tempfile
from typing import List from typing import List
from pathlib import Path
from fastapi import APIRouter, Depends, File, Form, UploadFile from fastapi import APIRouter, Depends, File, Form, UploadFile, HTTPException
from dbgpt._private.config import Config from dbgpt._private.config import Config
from dbgpt.app.knowledge.request.request import ( from dbgpt.app.knowledge.request.request import (
@ -332,54 +333,70 @@ def document_delete(space_name: str, query_request: DocumentQueryRequest):
@router.post("/knowledge/{space_name}/document/upload") @router.post("/knowledge/{space_name}/document/upload")
async def document_upload( async def document_upload(
space_name: str, space_name: str,
doc_name: str = Form(...), doc_name: str = Form(...),
doc_type: str = Form(...), doc_type: str = Form(...),
doc_file: UploadFile = File(...), doc_file: UploadFile = File(...),
): ):
print(f"/document/upload params: {space_name}") print(f"/document/upload params: {space_name}")
try: try:
if doc_file: if doc_file:
if not os.path.exists(os.path.join(KNOWLEDGE_UPLOAD_ROOT_PATH, space_name)): # Sanitize inputs to prevent path traversal
os.makedirs(os.path.join(KNOWLEDGE_UPLOAD_ROOT_PATH, space_name)) safe_space_name = os.path.basename(space_name)
# We can not move temp file in windows system when we open file in context of `with` safe_filename = os.path.basename(doc_file.filename)
tmp_fd, tmp_path = tempfile.mkstemp(
dir=os.path.join(KNOWLEDGE_UPLOAD_ROOT_PATH, space_name) # Create absolute paths and verify they are within allowed directory
) upload_dir = os.path.abspath(os.path.join(KNOWLEDGE_UPLOAD_ROOT_PATH, safe_space_name))
with os.fdopen(tmp_fd, "wb") as tmp: target_path = os.path.abspath(os.path.join(upload_dir, safe_filename))
tmp.write(await doc_file.read())
shutil.move( if not os.path.abspath(KNOWLEDGE_UPLOAD_ROOT_PATH) in target_path:
tmp_path, raise HTTPException(status_code=400, detail="Invalid path detected")
os.path.join(KNOWLEDGE_UPLOAD_ROOT_PATH, space_name, doc_file.filename),
) if not os.path.exists(upload_dir):
request = KnowledgeDocumentRequest() os.makedirs(upload_dir)
request.doc_name = doc_name
request.doc_type = doc_type # Create temp file
request.content = os.path.join( tmp_fd, tmp_path = tempfile.mkstemp(dir=upload_dir)
KNOWLEDGE_UPLOAD_ROOT_PATH, space_name, doc_file.filename
) try:
space_res = knowledge_space_service.get_knowledge_space( with os.fdopen(tmp_fd, "wb") as tmp:
KnowledgeSpaceRequest(name=space_name) tmp.write(await doc_file.read())
)
if len(space_res) == 0: shutil.move(tmp_path, target_path)
# create default space
if "default" != space_name: request = KnowledgeDocumentRequest()
raise Exception(f"you have not create your knowledge space.") request.doc_name = doc_name
knowledge_space_service.create_knowledge_space( request.doc_type = doc_type
KnowledgeSpaceRequest( request.content = target_path
name=space_name,
desc="first db-gpt rag application", space_res = knowledge_space_service.get_knowledge_space(
owner="dbgpt", KnowledgeSpaceRequest(name=safe_space_name)
) )
) if len(space_res) == 0:
return Result.succ( # create default space
knowledge_space_service.create_knowledge_document( if "default" != safe_space_name:
space=space_name, request=request raise Exception(f"you have not create your knowledge space.")
) knowledge_space_service.create_knowledge_space(
) KnowledgeSpaceRequest(
return Result.failed(code="E000X", msg=f"doc_file is None") name=safe_space_name,
except Exception as e: desc="first db-gpt rag application",
return Result.failed(code="E000X", msg=f"document add error {e}") owner="dbgpt",
)
)
return Result.succ(
knowledge_space_service.create_knowledge_document(
space=safe_space_name, request=request
)
)
except Exception as e:
# Clean up temp file if anything goes wrong
if os.path.exists(tmp_path):
os.unlink(tmp_path)
raise e
return Result.failed(code="E000X", msg=f"doc_file is None")
except Exception as e:
return Result.failed(code="E000X", msg=f"document add error {e}")
@router.post("/knowledge/{space_name}/document/sync") @router.post("/knowledge/{space_name}/document/sync")