From dd5dc53a64ef14ab5fe25b444cef4a87fb6bb0fe Mon Sep 17 00:00:00 2001 From: Wenjie Zhang Date: Fri, 2 Jan 2026 22:52:33 +0800 Subject: [PATCH] =?UTF-8?q?feat(security):=20=E5=AF=B9=E4=BA=8E=E5=9B=BE?= =?UTF-8?q?=E8=B0=B1=E6=96=87=E4=BB=B6=E4=B8=8A=E4=BC=A0=EF=BC=8C=E6=B7=BB?= =?UTF-8?q?=E5=8A=A0=E4=BB=8EMinIO=20URL=E4=B8=8B=E8=BD=BD=E6=96=87?= =?UTF-8?q?=E4=BB=B6=E7=9A=84=E4=B8=B4=E6=97=B6=E6=96=87=E4=BB=B6=E4=B8=8A?= =?UTF-8?q?=E4=B8=8B=E6=96=87=E7=AE=A1=E7=90=86=E5=99=A8?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 重构文件上传服务以使用新的MinIO上下文管理器 移除本地文件路径支持，仅允许MinIO URL --- server/routers/graph_router.py | 25 +++--- .../services/upload_graph_service.py | 58 ++++--------- src/storage/minio/client.py | 83 +++++++++++++++++++ 3 files changed, 111 insertions(+), 55 deletions(-) diff --git a/server/routers/graph_router.py b/server/routers/graph_router.py index 80c3373e..7cca6ea4 100644 --- a/server/routers/graph_router.py +++ b/server/routers/graph_router.py @@ -2,11 +2,12 @@ import traceback from fastapi import APIRouter, Body, Depends, HTTPException, Query -from src.storage.db.models import User from server.utils.auth_middleware import get_admin_user from src import graph_base, knowledge_base -from src.knowledge.adapters.factory import GraphAdapterFactory from src.knowledge.adapters.base import GraphAdapter +from src.knowledge.adapters.factory import GraphAdapterFactory +from src.storage.db.models import User +from src.storage.minio.client import StorageError from src.utils.logging_config import logger graph = APIRouter(prefix="/graph", tags=["graph"]) @@ -272,21 +273,17 @@ async def add_neo4j_entities( batch_size: int | None = Body(None), current_user: User = Depends(get_admin_user), ): - """通过JSONL文件添加图谱实体到Neo4j""" + """通过JSONL文件添加图谱实体到Neo4j（只接受 MinIO URL）""" try: - # 验证文件路径 - if not file_path or not isinstance(file_path, str): - return {"success": False, "message": "文件路径不能为空", "status": "failed"} - - file_path = file_path.strip() - if not file_path: - return {"success": False, "message": "文件路径不能为空", "status": "failed"} - - if not file_path.endswith(".jsonl"): - return {"success": False, "message": "文件格式错误，请上传jsonl文件", "status": "failed"} - + # 服务层会验证 URL 并从 MinIO 下载文件 await graph_base.jsonl_file_add_entity(file_path, kgdb_name, embed_model_name, batch_size) return {"success": True, "message": "实体添加成功", "status": "success"} + except StorageError as e: + # MinIO 验证或下载错误 + raise HTTPException(status_code=400, detail=str(e)) + except ValueError as e: + # 本地路径拒绝错误 + raise HTTPException(status_code=400, detail=str(e)) except Exception as e: logger.error(f"添加实体失败: {e}, {traceback.format_exc()}") return {"success": False, "message": f"添加实体失败: {e}", "status": "failed"} diff --git a/src/knowledge/services/upload_graph_service.py b/src/knowledge/services/upload_graph_service.py index 94b5fc1e..361a7b69 100644 --- a/src/knowledge/services/upload_graph_service.py +++ b/src/knowledge/services/upload_graph_service.py @@ -1,6 +1,5 @@ import json import os -import tempfile import traceback import warnings from urllib.parse import urlparse @@ -93,61 +92,38 @@ class UploadGraphService: # 检测 file_path 是否是 URL parsed_url = urlparse(file_path) - temp_file_path = None try: if parsed_url.scheme in ("http", "https"): # 如果是 URL logger.info(f"检测到 URL，正在从 MinIO 下载文件: {file_path}") - # 从 URL 解析 bucket_name 和 object_name - # URL 格式: http://host:port/bucket_name/object_name - path_parts = parsed_url.path.lstrip("/").split("/", 1) - if len(path_parts) < 2: - raise ValueError(f"无法解析 MinIO URL: {file_path}") - - bucket_name = path_parts[0] - object_name = path_parts[1] - - # 从 MinIO 下载文件 + # 使用 MinIO 客户端下载到临时文件（使用上下文管理器） minio_client = get_minio_client() - file_data = await minio_client.adownload_file(bucket_name, object_name) + async with minio_client.temp_file_from_url( + file_path, allowed_extensions=[".jsonl"] + ) as actual_file_path: + def read_triples(file_path): + with open(file_path, encoding="utf-8") as file: + for line in file: + if line.strip(): + yield json.loads(line.strip()) - # 创建临时文件保存下载的内容 - with tempfile.NamedTemporaryFile( - mode="w", suffix=".jsonl", delete=False, encoding="utf-8" - ) as temp_file: - temp_file.write(file_data.decode("utf-8")) - temp_file_path = temp_file.name + triples = list(read_triples(actual_file_path)) + await self.txt_add_vector_entity(triples, kgdb_name, embed_model_name, batch_size) + # 退出 with 块后，临时文件自动清理 - logger.info(f"文件已下载到临时路径: {temp_file_path}") - actual_file_path = temp_file_path else: - # 本地文件路径 - actual_file_path = file_path - - def read_triples(file_path): - with open(file_path, encoding="utf-8") as file: - for line in file: - if line.strip(): - yield json.loads(line.strip()) - - triples = list(read_triples(actual_file_path)) - - await self.txt_add_vector_entity(triples, kgdb_name, embed_model_name, batch_size) + # 本地文件路径 - 拒绝不安全的本地路径 + raise ValueError( + "不支持本地文件路径，只允许 MinIO URL。请先通过文件上传接口上传文件。" + ) except Exception as e: logger.error(f"处理文件失败: {e}") raise finally: - # 清理临时文件 - if temp_file_path and os.path.exists(temp_file_path): - try: - os.unlink(temp_file_path) - logger.info(f"已删除临时文件: {temp_file_path}") - except Exception as e: - logger.warning(f"删除临时文件失败: {e}") + self.connection.status = "open" - self.connection.status = "open" # 更新并保存图数据库信息 self.save_graph_info() return kgdb_name diff --git a/src/storage/minio/client.py b/src/storage/minio/client.py index afecf6d0..97d3a32c 100644 --- a/src/storage/minio/client.py +++ b/src/storage/minio/client.py @@ -6,6 +6,7 @@ MinIO 存储客户端 import asyncio import json import os +from contextlib import asynccontextmanager from datetime import timedelta from io import BytesIO @@ -275,6 +276,88 @@ class MinIOClient: logger.warning(f"设置存储桶 '{bucket_name}' 公共读取策略失败: {e}") raise StorageError(f"无法设置存储桶公共访问策略: {e}") + @asynccontextmanager + async def temp_file_from_url( + self, + url: str, + allowed_extensions: list[str] | None = None, + ): + """ + 异步上下文管理器：从 MinIO URL 下载文件到临时文件，使用后自动清理 + + Args: + url: MinIO 文件 URL + allowed_extensions: 允许的文件扩展名列表（可选） + + Yields: + str: 临时文件路径 + + Raises: + StorageError: 如果 URL 无效或下载失败 + """ + import tempfile + from urllib.parse import urlparse + + # 验证 URL + if not url or not isinstance(url, str): + raise StorageError("URL 不能为空") + + url = url.strip() + + if not url.startswith(("http://", "https://")): + raise StorageError("无效的 MinIO URL，只允许 http/https") + + parsed = urlparse(url) + + # 验证主机 + endpoint_host = self.endpoint.split("://")[-1].split(":")[0] + url_host = parsed.netloc.split(":")[0] + + if endpoint_host != url_host and url_host != os.environ.get("HOST_IP", "localhost"): + raise StorageError(f"不允许的外部 URL: {url_host}") + + # 检查路径遍历 + if ".." in url or "\\" in url: + raise StorageError("URL 包含路径遍历字符") + + # 验证扩展名 + if allowed_extensions and not any(url.endswith(ext) for ext in allowed_extensions): + raise StorageError(f"文件扩展名不符合要求，允许: {', '.join(allowed_extensions)}") + + # 解析 bucket 和 object name + path_parts = parsed.path.lstrip("/").split("/", 1) + if len(path_parts) != 2: + raise StorageError("无法解析 MinIO URL") + + bucket_name, object_name = path_parts + + # 下载文件 + file_data = await self.adownload_file(bucket_name, object_name) + logger.info(f"成功从 MinIO 下载文件: {object_name} ({len(file_data)} bytes)") + + # 创建临时文件 + if allowed_extensions: + suffix = next((ext for ext in allowed_extensions if url.endswith(ext)), ".tmp") + else: + suffix = f".{object_name.split('.')[-1]}" + + temp_path = None + try: + with tempfile.NamedTemporaryFile(mode="wb", suffix=suffix, delete=False) as temp_file: + temp_file.write(file_data) + temp_path = temp_file.name + + logger.info(f"文件已下载到临时路径: {temp_path}") + yield temp_path + + finally: + if temp_path and os.path.exists(temp_path): + try: + os.unlink(temp_path) + logger.info(f"已删除临时文件: {temp_path}") + except Exception as e: + logger.warning(f"删除临时文件失败: {e}") + # 全局客户端实例 _default_client = None