fix(security): 代码审查安全加固 — 三批次修复(6H/13M/12L)

第一批(快速修复):
- H6: main.py --debug 默认值 True→False,防止 Werkzeug RCE
- M2+M3: /search 增加 validate_query + top_k 范围限制(1-50)
- M4: context_count 范围限制(0-10) + 异常捕获
- L3: assert → raise RuntimeError(生产环境 API Key 检查)
- H1: SSE 错误事件移除 traceback 字段

第二批(安全加固):
- H2+H3: 文档接口路径遍历 realpath 校验 + 文件类型/大小限制
- H4+H5: 批量上传文件大小检查
- M6: LIKE 查询通配符转义
- M1: 37 处 str(e) 异常信息统一脱敏(6 文件)
- M5: CORS 生产环境限制来源
- M7: SESSION_MANAGER None 保护(503)
- M11: subprocess 参数注入防护(白名单 + -- 分隔符)

第三批(架构改进):
- M8+M9: 提取 JSON 解析共享工具(extract_json_object/list)
- M10: Prompt 注入检测防御(prompt_guard.py)
- M12: 解析器文件大小限制(Excel 50MB/TXT 20MB/PDF 100MB)
- M13: 全局单例竞态条件双重检查锁定(engine/bm25/intent_analyzer)
This commit is contained in:
lacerate551
2026-06-05 15:26:32 +08:00
parent a7206377cc
commit 90b915232a
19 changed files with 436 additions and 102 deletions

View File

@@ -67,7 +67,17 @@ def create_app() -> 'Flask':
static_folder = os.path.join(PROJECT_ROOT, 'chat-ui')
app = Flask(__name__, static_folder=static_folder, static_url_path='')
CORS(app)
# CORS 配置:生产环境限制来源,开发环境允许全部
if IS_PROD:
cors_origins = os.environ.get('CORS_ORIGINS', '').split(',') if os.environ.get('CORS_ORIGINS') else []
if cors_origins:
CORS(app, origins=cors_origins)
else:
CORS(app) # 未配置时仍允许全部,但记录警告
logger.warning("生产环境未配置 CORS_ORIGINSCORS 允许所有来源")
else:
CORS(app)
# ==================== Repository 依赖注入 ====================
@@ -227,14 +237,14 @@ def _validate_production_config() -> None:
- DASHSCOPE_API_KEY: 大模型调用必需
Raises:
AssertionError: 缺少必需的配置项
RuntimeError: 缺少必需的配置项
"""
import os
from config import DASHSCOPE_API_KEY
# 检查环境变量或配置文件中的 API Key
has_key = os.getenv("DASHSCOPE_API_KEY") or os.environ.get("DASHSCOPE_API_KEY") or DASHSCOPE_API_KEY
assert has_key and has_key != "", \
"Missing DASHSCOPE_API_KEY in production environment"
if not has_key or has_key == "":
raise RuntimeError("Missing DASHSCOPE_API_KEY in production environment")
logger.info("Configuration validated")

View File

@@ -97,7 +97,8 @@ def get_audit_logs():
return jsonify({"logs": logs, "total": total})
except Exception as e:
return jsonify({"error": str(e), "logs": [], "total": 0}), 500
logger.error(f"审计查询异常: {e}")
return jsonify({"error": "查询失败", "logs": [], "total": 0}), 500
def log_audit_event(user_id: str, username: str, action: str,

View File

@@ -41,6 +41,7 @@ from auth.gateway import require_gateway_auth
from auth.security import validate_query, filter_response
from config import RAG_CHAT_MODEL
from core.llm_utils import call_llm, call_llm_stream
from core.prompt_guard import detect_injection, sanitize_user_input
def _get_vlm_cache(image_path: str) -> Optional[str]:
@@ -1953,11 +1954,11 @@ def rag():
yield f"data: {json.dumps(finish_event, ensure_ascii=False)}\n\n"
except Exception as e:
import traceback
import logging as _logging
_logging.getLogger(__name__).error(f"[SSE] RAG 流异常: {e}", exc_info=True)
error_event = {
"type": "error",
"message": str(e),
"traceback": traceback.format_exc()
"message": "服务内部错误,请稍后重试"
}
yield f"data: {json.dumps(error_event, ensure_ascii=False)}\n\n"
@@ -1986,12 +1987,27 @@ def search():
"""
data = request.json or {}
query = data.get('query', '')
query = sanitize_user_input(query)
injection_matches = detect_injection(query)
if injection_matches:
logger.warning(f"[Chat] 检测到可疑注入: {injection_matches}")
top_k = data.get('top_k', 5)
collections = data.get('collections') # 后端传入的知识库列表
if not query:
return jsonify({'error': 'query is required'}), 400
# 输入安全校验(注入检测、违禁词、长度限制)
is_valid, reason = validate_query(query)
if not is_valid:
return jsonify({'error': reason}), 400
# top_k 范围校验
try:
top_k = max(1, min(int(top_k), 50))
except (ValueError, TypeError):
top_k = 5
# 如果没有指定 collections使用默认的公开库
if not collections:
collections = ['public_kb']

View File

@@ -96,6 +96,24 @@ def safe_filename(filename: str) -> str:
return safe_name
def _validate_doc_path(doc_path: str, documents_path: str) -> str:
"""
校验文档路径是否在允许目录内,防止路径遍历攻击。
Returns:
解析后的安全绝对路径
Raises:
ValueError: 路径不合法
"""
filepath = os.path.join(documents_path, doc_path)
real_path = os.path.realpath(filepath)
real_base = os.path.realpath(documents_path)
if not real_path.startswith(real_base + os.sep) and real_path != real_base:
raise ValueError("非法路径")
return real_path
# 延迟初始化缓存
_kb_manager = None
_kb_checked = False
@@ -160,6 +178,10 @@ def serve_document_file(doc_path: str) -> Tuple[Any, int]:
from flask import send_from_directory
filepath = os.path.join(DOCUMENTS_PATH, doc_path)
try:
filepath = _validate_doc_path(doc_path, DOCUMENTS_PATH)
except ValueError:
return jsonify({"error": "非法路径"}), 403
if not os.path.exists(filepath):
return jsonify({"error": "文件不存在"}), 404
@@ -329,7 +351,8 @@ def upload_document() -> Tuple[Any, int]:
sync_service.process_change(change)
sync_status = "已保存并添加到向量库"
except Exception as e:
sync_status = f"已保存,向量化失败: {str(e)}"
logger.warning(f"向量化失败: {e}")
sync_status = "已保存,向量化失败"
return success_response(
data={
@@ -404,6 +427,18 @@ def batch_upload_documents() -> Tuple[Any, int]:
})
continue
# 文件大小校验
file.seek(0, 2)
file_size = file.tell()
file.seek(0)
if file_size > MAX_FILE_SIZE:
results.append({
"filename": file.filename,
"status": "error",
"message": f"文件过大(最大 {MAX_FILE_SIZE // 1024 // 1024}MB"
})
continue
try:
original_filename = file.filename
ext = os.path.splitext(original_filename)[1].lower()
@@ -464,10 +499,11 @@ def batch_upload_documents() -> Tuple[Any, int]:
"replaced": replaced
})
except Exception as e:
logger.error(f"上传处理异常: {e}")
results.append({
"filename": file.filename,
"status": "error",
"message": str(e)
"message": "上传处理失败"
})
return success_response(
@@ -644,12 +680,27 @@ def update_document(doc_path: str) -> Tuple[Any, int]:
if file.filename == '':
return jsonify({"error": "没有选择文件"}), 400
# 文件类型校验
ext = os.path.splitext(file.filename)[1].lower()
if ext not in ALLOWED_EXTENSIONS:
return jsonify({"error": f"不支持的文件类型: {ext}"}), 400
# 文件大小校验
file.seek(0, 2) # 跳到文件末尾获取大小
file_size = file.tell()
file.seek(0) # 回到文件开头
if file_size > MAX_FILE_SIZE:
return jsonify({"error": f"文件过大(最大 {MAX_FILE_SIZE // 1024 // 1024}MB"}), 400
# 解析路径
parts = doc_path.split('/')
if len(parts) < 2:
return jsonify({"error": "无效的文档路径"}), 400
filepath = os.path.join(DOCUMENTS_PATH, doc_path)
try:
filepath = _validate_doc_path(doc_path, DOCUMENTS_PATH)
except ValueError:
return jsonify({"error": "非法路径"}), 403
if not os.path.exists(filepath):
return jsonify({"error": "文件不存在"}), 404
@@ -711,7 +762,10 @@ def delete_document(doc_path: str) -> Tuple[Any, int]:
# 目录名即向量库名
collection = subdir
filepath = os.path.join(DOCUMENTS_PATH, doc_path)
try:
filepath = _validate_doc_path(doc_path, DOCUMENTS_PATH)
except ValueError:
return jsonify({"error": "非法路径"}), 403
if not os.path.exists(filepath):
return jsonify({"error": "文件不存在"}), 404
@@ -730,7 +784,8 @@ def delete_document(doc_path: str) -> Tuple[Any, int]:
})
except Exception as e:
return jsonify({"error": f"删除失败: {str(e)}"}), 500
logger.error(f"删除文档异常: {e}")
return jsonify({"error": "删除失败,请稍后重试"}), 500
@document_bp.route('/documents/<path:doc_path>/chunks', methods=['GET'])
@@ -817,7 +872,10 @@ def preview_document(doc_path: str) -> Tuple[Any, int]:
# 查询参数
chunk_index_str = request.args.get('chunk_index')
context_count = int(request.args.get('context', 2))
try:
context_count = max(0, min(int(request.args.get('context', 2)), 10))
except (ValueError, TypeError):
context_count = 2
# 获取所有切片
all_chunks = kb_manager.get_document_chunks(collection, filename)
@@ -1082,5 +1140,5 @@ def delete_chunks_by_source() -> Tuple[Any, int]:
"message": f"已删除 {deleted_count} 个切片"
})
except Exception as e:
logger.error(f"批量删除切片失败: {e}")
return jsonify({"error": str(e)}), 500
logger.error(f"操作异常: {e}")
return jsonify({"error": "操作失败,请稍后重试"}), 500

View File

@@ -33,6 +33,9 @@
from flask import Blueprint, request, jsonify
from auth.gateway import require_gateway_auth, require_role
import logging
logger = logging.getLogger(__name__)
feedback_bp = Blueprint('feedback', __name__)
@@ -95,7 +98,8 @@ def submit_feedback():
"suggestion_id": result.get('suggestion_id')
})
except Exception as e:
return jsonify({"error": str(e)}), 500
logger.error(f"反馈操作异常: {e}")
return jsonify({"error": "操作失败,请稍后重试"}), 500
@feedback_bp.route('/feedback/stats', methods=['GET'])
@@ -113,7 +117,8 @@ def get_feedback_stats():
"stats": stats
})
except Exception as e:
return jsonify({"error": str(e)}), 500
logger.error(f"反馈操作异常: {e}")
return jsonify({"error": "操作失败,请稍后重试"}), 500
@feedback_bp.route('/feedback/list', methods=['GET'])
@@ -141,7 +146,8 @@ def get_feedback_list():
"total": len(feedbacks)
})
except Exception as e:
return jsonify({"error": str(e)}), 500
logger.error(f"反馈操作异常: {e}")
return jsonify({"error": "操作失败,请稍后重试"}), 500
@feedback_bp.route('/reports/weekly', methods=['GET'])
@@ -156,7 +162,8 @@ def get_weekly_report():
"report": report.to_dict()
})
except Exception as e:
return jsonify({"error": str(e)}), 500
logger.error(f"反馈操作异常: {e}")
return jsonify({"error": "操作失败,请稍后重试"}), 500
@feedback_bp.route('/reports/monthly', methods=['GET'])
@@ -171,7 +178,8 @@ def get_monthly_report():
"report": report.to_dict()
})
except Exception as e:
return jsonify({"error": str(e)}), 500
logger.error(f"反馈操作异常: {e}")
return jsonify({"error": "操作失败,请稍后重试"}), 500
@feedback_bp.route('/faq', methods=['GET'])
@@ -190,7 +198,8 @@ def get_faq_list():
"total": len(faqs)
})
except Exception as e:
return jsonify({"error": str(e)}), 500
logger.error(f"反馈操作异常: {e}")
return jsonify({"error": "操作失败,请稍后重试"}), 500
@feedback_bp.route('/faq', methods=['POST'])
@@ -233,7 +242,8 @@ def create_faq():
"message": "FAQ已创建请通过 /faq/<id>/approve 接口确认后生效"
})
except Exception as e:
return jsonify({"error": str(e)}), 500
logger.error(f"反馈操作异常: {e}")
return jsonify({"error": "操作失败,请稍后重试"}), 500
@feedback_bp.route('/faq/<int:faq_id>/approve', methods=['POST'])
@@ -276,7 +286,8 @@ def approve_faq(faq_id):
"message": "FAQ已批准并同步到知识库"
})
except Exception as e:
return jsonify({"error": str(e)}), 500
logger.error(f"反馈操作异常: {e}")
return jsonify({"error": "操作失败,请稍后重试"}), 500
@feedback_bp.route('/faq/<int:faq_id>', methods=['PUT'])
@@ -326,7 +337,8 @@ def update_faq(faq_id):
"sync_status": sync_status
})
except Exception as e:
return jsonify({"error": str(e)}), 500
logger.error(f"反馈操作异常: {e}")
return jsonify({"error": "操作失败,请稍后重试"}), 500
@feedback_bp.route('/faq/<int:faq_id>', methods=['DELETE'])
@@ -358,7 +370,8 @@ def delete_faq(faq_id):
"message": "FAQ删除成功" if deleted else "FAQ不存在"
})
except Exception as e:
return jsonify({"error": str(e)}), 500
logger.error(f"反馈操作异常: {e}")
return jsonify({"error": "操作失败,请稍后重试"}), 500
@feedback_bp.route('/faq/suggestions', methods=['GET'])
@@ -378,7 +391,8 @@ def get_faq_suggestions():
"total": len(suggestions)
})
except Exception as e:
return jsonify({"error": str(e)}), 500
logger.error(f"反馈操作异常: {e}")
return jsonify({"error": "操作失败,请稍后重试"}), 500
@feedback_bp.route('/faq/suggestions/<int:suggestion_id>/approve', methods=['POST'])
@@ -413,7 +427,8 @@ def approve_faq_suggestion(suggestion_id):
else:
return jsonify({"error": result.get('error', '批准失败')}), 400
except Exception as e:
return jsonify({"error": str(e)}), 500
logger.error(f"反馈操作异常: {e}")
return jsonify({"error": "操作失败,请稍后重试"}), 500
@feedback_bp.route('/faq/suggestions/<int:suggestion_id>/reject', methods=['POST'])
@@ -429,7 +444,8 @@ def reject_faq_suggestion(suggestion_id):
"message": "FAQ建议已拒绝" if rejected else "建议不存在"
})
except Exception as e:
return jsonify({"error": str(e)}), 500
logger.error(f"反馈操作异常: {e}")
return jsonify({"error": "操作失败,请稍后重试"}), 500
# ==================== Bad Case 分析接口 ====================
@@ -472,7 +488,8 @@ def get_bad_cases():
]
})
except Exception as e:
return jsonify({"error": str(e)}), 500
logger.error(f"反馈操作异常: {e}")
return jsonify({"error": "操作失败,请稍后重试"}), 500
@feedback_bp.route('/feedback/blacklist', methods=['GET'])
@@ -497,4 +514,5 @@ def get_chunk_blacklist():
"usage": "在检索时过滤这些来源以提升回答质量"
})
except Exception as e:
return jsonify({"error": str(e)}), 500
logger.error(f"反馈操作异常: {e}")
return jsonify({"error": "操作失败,请稍后重试"}), 500

View File

@@ -9,8 +9,11 @@
"""
import os
import logging
from flask import Blueprint, send_file, jsonify, current_app
logger = logging.getLogger(__name__)
image_bp = Blueprint('images', __name__)
@@ -70,7 +73,8 @@ def get_image(image_id: str):
return send_file(image_path, mimetype=mimetype)
except Exception as e:
return jsonify({"error": f"读取图片失败: {str(e)}"}), 500
logger.error(f"读取图片异常: {e}")
return jsonify({"error": "读取图片失败"}), 500
return jsonify({"error": "图片不存在", "image_id": image_id}), 404
@@ -122,7 +126,8 @@ def get_image_info(image_id: str):
"url": f"/images/{image_id}"
})
except Exception as e:
return jsonify({"error": f"读取图片信息失败: {str(e)}"}), 500
logger.error(f"读取图片信息异常: {e}")
return jsonify({"error": "读取图片信息失败"}), 500
return jsonify({"error": "图片不存在", "image_id": image_id}), 404
@@ -179,7 +184,8 @@ def list_images():
})
except Exception as e:
return jsonify({"error": f"列出图片失败: {str(e)}"}), 500
logger.error(f"列出图片异常: {e}")
return jsonify({"error": "列出图片失败"}), 500
@image_bp.route('/images/stats', methods=['GET'])
@@ -222,4 +228,5 @@ def image_stats():
})
except Exception as e:
return jsonify({"error": f"获取统计信息失败: {str(e)}"}), 500
logger.error(f"获取统计信息异常: {e}")
return jsonify({"error": "获取统计信息失败"}), 500

View File

@@ -391,10 +391,11 @@ def sync_documents() -> Tuple[Any, int]:
}
})
except Exception as e:
logger.error(f"知识库操作异常: {e}")
results.append({
"collection": "all",
"status": "error",
"message": str(e)
"message": "操作失败"
})
else:
# 没有 sync_service返回提示
@@ -462,7 +463,8 @@ def debug_scan() -> Tuple[Any, int]:
result["scanned_count"] = len(scanned)
result["scanned_ids"] = list(scanned.keys())[:10]
except Exception as e:
result["scan_error"] = str(e)
logger.error(f"扫描异常: {e}")
result["scan_error"] = "扫描失败"
return jsonify(result)
@@ -499,9 +501,11 @@ def reindex_collection(kb_name: str) -> Tuple[Any, int]:
from data.db import get_connection
with get_connection("knowledge") as conn:
cursor = conn.cursor()
# 转义 LIKE 通配符,防止 kb_name 中的 % 或 _ 导致非预期匹配
escaped_kb = kb_name.replace('%', '\\%').replace('_', '\\_')
# 删除以 "{kb_name}/" 或 "{kb_name}\" 开头的文档哈希(兼容 Windows 和 Linux
cursor.execute("DELETE FROM document_hashes WHERE document_id LIKE ? OR document_id LIKE ?",
(f"{kb_name}/%", f"{kb_name}\\%"))
cursor.execute("DELETE FROM document_hashes WHERE document_id LIKE ? ESCAPE '\\' OR document_id LIKE ? ESCAPE '\\'",
(f"{escaped_kb}/%", f"{escaped_kb}\\%"))
deleted = cursor.rowcount
logger.info(f"已清除 {deleted} 条哈希记录: {kb_name}")
except Exception as e:
@@ -520,7 +524,8 @@ def reindex_collection(kb_name: str) -> Tuple[Any, int]:
"errors": result.errors
})
except Exception as e:
return jsonify({"error": str(e)}), 500
logger.error(f"reindex 异常: {e}")
return jsonify({"error": "操作失败,请稍后重试"}), 500
else:
return jsonify({"error": "同步服务不可用"}), 503
@@ -633,7 +638,8 @@ def deprecate_document(kb_name: str, filename: str) -> Tuple[Any, int]:
)
return jsonify(result)
except Exception as e:
return jsonify({"success": False, "error": str(e)}), 500
logger.error(f"操作异常: {e}")
return jsonify({"success": False, "error": "操作失败,请稍后重试"}), 500
@kb_bp.route('/collections/<kb_name>/documents/<path:filename>/restore', methods=['POST'])
@@ -664,7 +670,8 @@ def restore_document(kb_name: str, filename: str) -> Tuple[Any, int]:
result = kb_manager.restore_document(kb_name, filename)
return jsonify(result)
except Exception as e:
return jsonify({"success": False, "error": str(e)}), 500
logger.error(f"操作异常: {e}")
return jsonify({"success": False, "error": "操作失败,请稍后重试"}), 500
@kb_bp.route('/collections/<kb_name>/documents/<path:filename>/versions', methods=['GET'])
@@ -720,7 +727,8 @@ def get_document_versions(kb_name: str, filename: str) -> Tuple[Any, int]:
"total": len(versions_data)
})
except Exception as e:
return jsonify({"success": False, "error": str(e)}), 500
logger.error(f"操作异常: {e}")
return jsonify({"success": False, "error": "操作失败,请稍后重试"}), 500
@kb_bp.route('/collections/<kb_name>/update-image-descriptions', methods=['POST'])
@@ -754,7 +762,8 @@ def update_image_descriptions(kb_name: str) -> Tuple[Any, int]:
result = kb_manager.update_image_descriptions(kb_name)
return jsonify(result)
except Exception as e:
return jsonify({"success": False, "error": str(e)}), 500
logger.error(f"操作异常: {e}")
return jsonify({"success": False, "error": "操作失败,请稍后重试"}), 500
@kb_bp.route('/collections/sync-vlm-cache', methods=['POST'])

View File

@@ -33,6 +33,8 @@ def get_sessions():
}
"""
session_manager = current_app.config['SESSION_MANAGER']
if session_manager is None:
return jsonify({"error": "会话服务不可用"}), 503
user_id = request.current_user["user_id"]
sessions = session_manager.get_user_sessions(user_id, limit=20)
@@ -62,6 +64,8 @@ def get_history(session_id):
}
"""
session_manager = current_app.config['SESSION_MANAGER']
if session_manager is None:
return jsonify({"error": "会话服务不可用"}), 503
user_id = request.current_user["user_id"]
# 验证会话归属
@@ -81,6 +85,8 @@ def get_history(session_id):
def delete_session(session_id):
"""删除会话"""
session_manager = current_app.config['SESSION_MANAGER']
if session_manager is None:
return jsonify({"error": "会话服务不可用"}), 503
user_id = request.current_user["user_id"]
# 验证会话归属
@@ -100,6 +106,8 @@ def delete_session(session_id):
def clear_history(session_id):
"""清空会话历史(保留会话)"""
session_manager = current_app.config['SESSION_MANAGER']
if session_manager is None:
return jsonify({"error": "会话服务不可用"}), 503
user_id = request.current_user["user_id"]
# 验证会话归属

View File

@@ -112,10 +112,11 @@ def trigger_sync() -> Tuple[Any, int]:
message="同步完成"
)
except Exception as e:
logger.error(f"同步操作异常: {e}")
return error_response(
error="SYNC_ERROR",
error_code=SYNC_ERROR,
message=str(e),
message="操作失败",
http_status=500
)
@@ -164,11 +165,12 @@ def get_sync_status() -> Tuple[Any, int]:
return jsonify(status)
except Exception as e:
logger.error(f"状态查询异常: {e}")
return jsonify({
"status": "failed",
"status_code": INTERNAL_ERROR,
"enabled": True,
"error": str(e)
"error": "操作失败"
})
@@ -196,10 +198,11 @@ def get_sync_history() -> Tuple[Any, int]:
history = service.get_sync_history(limit=limit) if hasattr(service, 'get_sync_history') else []
return jsonify({"history": history})
except Exception as e:
logger.error(f"同步操作异常: {e}")
return error_response(
error="SYNC_ERROR",
error_code=SYNC_ERROR,
message=str(e),
message="操作失败",
http_status=500
)
@@ -230,10 +233,11 @@ def get_change_logs() -> Tuple[Any, int]:
changes = service.get_change_logs(limit=limit, collection=collection) if hasattr(service, 'get_change_logs') else []
return jsonify({"changes": changes})
except Exception as e:
logger.error(f"同步操作异常: {e}")
return error_response(
error="SYNC_ERROR",
error_code=SYNC_ERROR,
message=str(e),
message="操作失败",
http_status=500
)
@@ -275,10 +279,11 @@ def start_sync_monitor() -> Tuple[Any, int]:
else:
return jsonify({"status": "success", "message": "文件监控功能不可用"})
except Exception as e:
logger.error(f"同步操作异常: {e}")
return error_response(
error="SYNC_ERROR",
error_code=SYNC_ERROR,
message=str(e),
message="操作失败",
http_status=500
)
@@ -303,9 +308,10 @@ def stop_sync_monitor() -> Tuple[Any, int]:
service.stop()
return jsonify({"status": "success", "status_code": SYNC_SUCCESS, "message": "文件监控已停止"})
except Exception as e:
logger.error(f"同步操作异常: {e}")
return error_response(
error="SYNC_ERROR",
error_code=SYNC_ERROR,
message=str(e),
message="操作失败",
http_status=500
)

View File

@@ -14,6 +14,7 @@ BM25 关键词检索索引
import os
import pickle
import threading
import numpy as np
from rank_bm25 import BM25Okapi
import jieba
@@ -105,24 +106,27 @@ class BM25Index:
# ==================== 全局 BM25 索引管理器 ====================
_bm25_indexer: BM25Index = None
_bm25_lock = threading.Lock() # 单例创建锁,防止多线程竞态条件
def get_bm25_indexer() -> BM25Index:
"""
获取全局 BM25 索引器实例
获取全局 BM25 索引器实例(双重检查锁定,线程安全)
Returns:
BM25Index 实例
"""
global _bm25_indexer
if _bm25_indexer is None:
_bm25_indexer = BM25Index()
with _bm25_lock:
if _bm25_indexer is None:
_bm25_indexer = BM25Index()
return _bm25_indexer
def init_bm25_indexer(ids=None, documents=None, metadatas=None) -> BM25Index:
"""
初始化 BM25 索引器并添加文档
初始化 BM25 索引器并添加文档(加锁,线程安全)
Args:
ids: 文档 ID 列表
@@ -133,7 +137,8 @@ def init_bm25_indexer(ids=None, documents=None, metadatas=None) -> BM25Index:
初始化后的 BM25Index 实例
"""
global _bm25_indexer
_bm25_indexer = BM25Index()
if ids and documents:
_bm25_indexer.add_documents(ids, documents, metadatas or [])
with _bm25_lock:
_bm25_indexer = BM25Index()
if ids and documents:
_bm25_indexer.add_documents(ids, documents, metadatas or [])
return _bm25_indexer

View File

@@ -31,6 +31,7 @@ import os
import gc
import time
import logging
import threading
import numpy as np
from typing import List, Tuple, Optional, Dict, Any
@@ -50,6 +51,7 @@ except ImportError:
# 延迟导入,防止循环依赖
_engine_instance = None
_engine_lock = threading.Lock() # 单例创建锁,防止多线程竞态条件
try:
from config import (
@@ -327,10 +329,12 @@ class RAGEngine:
@classmethod
def get_instance(cls) -> 'RAGEngine':
"""获取引擎单例实例"""
"""获取引擎单例实例(双重检查锁定,线程安全)"""
global _engine_instance
if _engine_instance is None:
_engine_instance = cls()
with _engine_lock:
if _engine_instance is None:
_engine_instance = cls()
return _engine_instance
def __init__(self) -> None:

View File

@@ -18,6 +18,7 @@
import json
import logging
import hashlib
import threading
from dataclasses import dataclass, asdict
from typing import List, Dict, Optional, Tuple
@@ -503,12 +504,15 @@ class IntentAnalyzer:
# ==================== 便捷函数 ====================
_analyzer = None
_analyzer_lock = threading.Lock() # 单例创建锁,防止多线程竞态条件
def get_intent_analyzer() -> IntentAnalyzer:
"""获取意图分析器单例"""
"""获取意图分析器单例(双重检查锁定,线程安全)"""
global _analyzer
if _analyzer is None:
_analyzer = IntentAnalyzer()
with _analyzer_lock:
if _analyzer is None:
_analyzer = IntentAnalyzer()
return _analyzer

View File

@@ -237,6 +237,76 @@ def parse_json_list_from_response(content: str) -> Optional[List[dict]]:
return None
def extract_json_object(content: str) -> Optional[dict]:
"""
多策略从 LLM 响应中提取 JSON 对象(增强版)
在 parse_json_from_response 基础上增加 fallback 策略:
1. 先调用 parse_json_from_responsemarkdown 代码块 → 直接解析)
2. 失败后 fallback 到正则匹配最外层 {...} 块
Args:
content: LLM 返回的原始内容
Returns:
解析后的字典,全部策略失败返回 None
"""
if not content:
return None
# 策略1+2markdown 代码块提取 + 直接 json.loads
result = parse_json_from_response(content)
if result is not None and isinstance(result, dict):
return result
# 策略3fallback正则匹配最外层 JSON 对象 {...}
brace_match = re.search(r'\{[\s\S]*\}', content)
if brace_match:
try:
parsed = json.loads(brace_match.group(0))
if isinstance(parsed, dict):
return parsed
except (json.JSONDecodeError, TypeError, ValueError):
pass
return None
def extract_json_list(content: str) -> Optional[list]:
"""
多策略从 LLM 响应中提取 JSON 数组(增强版)
在 parse_json_list_from_response 基础上增加 fallback 策略:
1. 先调用 parse_json_list_from_responsemarkdown 代码块 → 直接解析 → 嵌套提取)
2. 失败后 fallback 到正则匹配最外层 [...] 块
Args:
content: LLM 返回的原始内容
Returns:
解析后的列表,全部策略失败返回 None
"""
if not content:
return None
# 策略1+2markdown 代码块提取 + 直接 json.loads + 嵌套 key 提取
result = parse_json_list_from_response(content)
if result is not None:
return result
# 策略3fallback正则匹配最外层 JSON 数组 [...]
bracket_match = re.search(r'\[[\s\S]*\]', content)
if bracket_match:
try:
parsed = json.loads(bracket_match.group(0))
if isinstance(parsed, list):
return parsed
except (json.JSONDecodeError, TypeError, ValueError):
pass
return None
# ==================== 便捷函数 ====================
def quick_ask(

104
core/prompt_guard.py Normal file
View File

@@ -0,0 +1,104 @@
"""
Prompt 注入防御工具
提供轻量级的 prompt 注入检测和清理功能,防止恶意用户通过精心构造的输入
覆盖系统指令或操纵 LLM 行为。
"""
import re
import logging
logger = logging.getLogger(__name__)
# 常见注入模式(中英文)
_INJECTION_PATTERNS = [
# 英文注入模式
r'ignore\s+(all\s+)?(previous|above|prior)\s+(instructions?|prompts?|rules?)',
r'you\s+are\s+now\s+(a|an)\s+',
r'forget\s+(everything|all|your)\s+',
r'disregard\s+(all\s+)?(previous|prior|above)',
r'override\s+(your|the|all)\s+(instructions?|rules?|system)',
r'reveal\s+(your|the)\s+(system\s+)?(prompt|instructions?)',
r'act\s+as\s+(if\s+)?(you\s+are\s+)?',
# 中文注入模式
r'忽略(之前|以上|前面|所有).*(指令|规则|提示|约束)',
r'你现在是(一个|新的)?',
r'忘记(之前|所有|你).*(规则|指令|设定)',
r'无视(系统|所有|之前).*(指令|规则|提示)',
r'不要遵守(任何|之前|系统).*(指令|规则)',
]
# 编译正则表达式(不区分大小写)
_COMPILED_PATTERNS = [
re.compile(p, re.IGNORECASE) for p in _INJECTION_PATTERNS
]
def detect_injection(text: str, threshold: int = 1) -> list:
"""
检测文本中是否包含 prompt 注入模式
Args:
text: 待检测的用户输入文本
threshold: 匹配数量阈值,达到此数量视为可疑
Returns:
匹配到的注入模式描述列表(空列表表示安全)
"""
if not text or not text.strip():
return []
matches = []
for pattern in _COMPILED_PATTERNS:
match = pattern.search(text)
if match:
matches.append(match.group()[:50])
if len(matches) >= threshold:
logger.warning(f"[PromptGuard] 检测到可疑注入模式: {matches[:3]}, 输入前100字: {text[:100]}")
return matches
def sanitize_user_input(text: str, max_length: int = 2000) -> str:
"""
清理用户输入,去除潜在的控制指令
策略:
1. 截断超长输入
2. 去除 null 字节和控制字符
3. 限制输入长度
Args:
text: 原始用户输入
max_length: 最大长度
Returns:
清理后的安全文本
"""
if not text:
return ""
# 去除 null 字节和控制字符(保留换行和空格)
text = re.sub(r'[\x00-\x08\x0b\x0c\x0e-\x1f\x7f]', '', text)
# 截断超长输入
if len(text) > max_length:
text = text[:max_length]
logger.debug(f"[PromptGuard] 输入超长,截断至 {max_length}")
return text
def wrap_document_content(content: str, label: str = "参考资料") -> str:
"""
将文档内容包装在明确的标记内,降低文档内容被当作指令执行的风险
Args:
content: 文档原始内容
label: 标签名称
Returns:
包装后的内容
"""
return f"[以下为{label},非系统指令]\n---\n{content}\n---\n[{label}结束]"

View File

@@ -16,7 +16,6 @@
import json
import logging
import re
import threading
import time
from concurrent.futures import ThreadPoolExecutor, as_completed
@@ -24,7 +23,7 @@ from functools import wraps
from typing import List, Dict, Any, Optional
# 导入 LLM 工具函数
from core.llm_utils import call_llm
from core.llm_utils import call_llm, extract_json_object
logger = logging.getLogger(__name__)
@@ -378,47 +377,12 @@ class AnswerGrader:
return result
def _extract_json(self, response: str) -> dict:
"""
多策略从 LLM 响应中提取 JSON 对象
策略优先级:
1. markdown 代码块提取 ```json ... ```
2. 直接 json.loads
3. 正则匹配第一个 {...} 块
"""
"""多策略从 LLM 响应中提取 JSON 对象(使用共享工具)"""
result = extract_json_object(response)
if result is not None:
return result
if not response:
raise ValueError("LLM 返回为空")
# 策略1提取 markdown 代码块
json_match = re.search(r'```(?:json)?\s*([\s\S]*?)\s*```', response)
if json_match:
json_str = json_match.group(1).strip()
try:
result = json.loads(json_str)
if isinstance(result, dict):
return result
except json.JSONDecodeError:
pass # 继续下一策略
# 策略2直接解析整个响应
try:
result = json.loads(response.strip())
if isinstance(result, dict):
return result
except json.JSONDecodeError:
pass # 继续下一策略
# 策略3正则匹配最外层 JSON 对象
brace_match = re.search(r'\{[\s\S]*\}', response)
if brace_match:
try:
result = json.loads(brace_match.group(0))
if isinstance(result, dict):
return result
except json.JSONDecodeError:
pass
# 全部策略失败
raise json.JSONDecodeError(
f"无法从 LLM 响应中提取有效 JSON响应前300字: {response[:300]}",
response, 0

View File

@@ -28,7 +28,7 @@ def main():
parser = argparse.ArgumentParser(description='RAG API 服务')
parser.add_argument('--host', default='0.0.0.0', help='监听地址(默认 0.0.0.0')
parser.add_argument('--port', type=int, default=5001, help='监听端口(默认 5001')
parser.add_argument('--debug', action='store_true', default=True, help='调试模式')
parser.add_argument('--debug', action='store_true', default=False, help='调试模式')
parser.add_argument('--no-debug', action='store_true', help='关闭调试模式')
args = parser.parse_args()

View File

@@ -24,6 +24,9 @@ logger = logging.getLogger(__name__)
# 大表切片阈值
MAX_ROWS_PER_CHUNK = 200
# 文件大小限制
MAX_FILE_SIZE = 50 * 1024 * 1024 # 50MB
@dataclass
class UnifiedChunk:
@@ -67,6 +70,13 @@ def parse_excel(
}
"""
filepath = Path(filepath)
# 检查文件大小
if filepath.exists():
file_size = filepath.stat().st_size
if file_size > MAX_FILE_SIZE:
raise ValueError(f"文件过大: {file_size / 1024 / 1024:.1f}MB最大允许 {MAX_FILE_SIZE / 1024 / 1024:.0f}MB")
if not filepath.exists():
raise FileNotFoundError(f"文件不存在: {filepath}")

View File

@@ -46,6 +46,9 @@ import logging
logger = logging.getLogger(__name__)
# 文件大小限制
MAX_PDF_SIZE = 100 * 1024 * 1024 # 100MB
# 支持的文件格式
SUPPORTED_FORMATS = {
'.pdf': 'PDF 文档',
@@ -176,6 +179,11 @@ def parse_with_mineru_online(
if not file_path.exists():
raise FileNotFoundError(f"文件不存在: {file_path}")
# 检查文件大小
file_size = file_path.stat().st_size
if file_size > MAX_PDF_SIZE:
raise ValueError(f"文件过大: {file_size / 1024 / 1024:.1f}MB最大允许 {MAX_PDF_SIZE / 1024 / 1024:.0f}MB")
logger.info(f"使用 MinerU 在线 API 解析: {file_path.name}")
# 读取文件内容
@@ -600,6 +608,11 @@ def parse_with_mineru(
if not file_path.exists():
raise FileNotFoundError(f"文件不存在: {file_path}")
# 检查文件大小
file_size = file_path.stat().st_size
if file_size > MAX_PDF_SIZE:
raise ValueError(f"文件过大: {file_size / 1024 / 1024:.1f}MB最大允许 {MAX_PDF_SIZE / 1024 / 1024:.0f}MB")
# 检查文件格式
suffix = file_path.suffix.lower()
if suffix not in SUPPORTED_FORMATS:
@@ -629,8 +642,17 @@ def parse_with_mineru(
if not mineru_exe.exists():
mineru_exe = "mineru" # 回退到系统 PATH
# 参数白名单校验,防止注入非法参数
ALLOWED_BACKENDS = {'auto', 'pipeline', 'vlm', 'vlm-sglang', 'vlm-auto-engine', 'hybrid-auto-engine', 'ocr'}
ALLOWED_LANGS = {'ch', 'en', 'ch_lite', 'en_lite', 'formula', 'table'}
if backend not in ALLOWED_BACKENDS:
backend = 'pipeline'
if lang not in ALLOWED_LANGS:
lang = 'ch'
cmd = [
str(mineru_exe),
"--",
"-p", str(file_path),
"-o", str(output_dir),
"-m", "auto",
@@ -1458,6 +1480,11 @@ def parse_with_mineru_persistent(
if not file_path.exists():
raise FileNotFoundError(f"文件不存在: {file_path}")
# 检查文件大小
file_size = file_path.stat().st_size
if file_size > MAX_PDF_SIZE:
raise ValueError(f"文件过大: {file_size / 1024 / 1024:.1f}MB最大允许 {MAX_PDF_SIZE / 1024 / 1024:.0f}MB")
# 计算文件 hash用于隔离输出目录
file_hash = compute_file_hash(str(file_path))
output_dir = Path(output_base) / file_hash

View File

@@ -5,12 +5,25 @@ TXT 文本解析器
"""
import logging
import os
logger = logging.getLogger(__name__)
# 文件大小限制
MAX_TXT_SIZE = 20 * 1024 * 1024 # 20MB
def extract_text_from_txt(filepath):
"""从TXT提取文本"""
# 检查文件大小
try:
file_size = os.path.getsize(filepath)
if file_size > MAX_TXT_SIZE:
logger.error(f"TXT文件过大: {file_size / 1024 / 1024:.1f}MB")
return ""
except OSError:
return ""
try:
with open(filepath, 'r', encoding='utf-8') as f:
return f.read()