feat(security): 对于图谱文件上传，添加从MinIO URL下载文件的临时文件上下文管理器
重构文件上传服务以使用新的MinIO上下文管理器 移除本地文件路径支持，仅允许MinIO URL
This commit is contained in:
Wenjie Zhang
parent
2524b60d9a
commit
dd5dc53a64
@ -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"}
|
||||
|
||||
@ -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
|
||||
|
||||
@ -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
|
||||
|
||||
Loading…
Reference in New Issue
Block a user