fix: prevent arbitrary code execution via pickle in FAISS vector store (#4833)
This commit is contained in:
Kartik
+155
-17
@@ -1,10 +1,11 @@
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import pickle
|
||||
import uuid
|
||||
import warnings
|
||||
from pathlib import Path
|
||||
from typing import Dict, List, Optional
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
import numpy as np
|
||||
from pydantic import BaseModel
|
||||
@@ -13,7 +14,7 @@ try:
|
||||
# Suppress SWIG deprecation warnings from FAISS
|
||||
warnings.filterwarnings("ignore", category=DeprecationWarning, message=".*SwigPy.*")
|
||||
warnings.filterwarnings("ignore", category=DeprecationWarning, message=".*swigvarlink.*")
|
||||
|
||||
|
||||
logging.getLogger("faiss").setLevel(logging.WARNING)
|
||||
logging.getLogger("faiss.loader").setLevel(logging.WARNING)
|
||||
|
||||
@@ -30,6 +31,93 @@ from mem0.vector_stores.base import VectorStoreBase
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class SafeUnpickler(pickle.Unpickler):
|
||||
"""
|
||||
Restricted unpickler that only allows safe built-in types.
|
||||
|
||||
This prevents arbitrary code execution via pickle deserialization by only
|
||||
allowing a whitelist of safe types (dict, list, str, int, float, bool, tuple, None).
|
||||
"""
|
||||
|
||||
# Only allow builtins module
|
||||
SAFE_MODULES = frozenset({"builtins", "__builtin__"})
|
||||
# Only allow safe basic types
|
||||
SAFE_NAMES = frozenset({"dict", "list", "str", "int", "float", "bool", "tuple", "set", "frozenset", "NoneType"})
|
||||
|
||||
def find_class(self, module: str, name: str) -> Any:
|
||||
"""Override find_class to only allow safe types."""
|
||||
if module in self.SAFE_MODULES and name in self.SAFE_NAMES:
|
||||
import builtins
|
||||
|
||||
if hasattr(builtins, name):
|
||||
return getattr(builtins, name)
|
||||
# NoneType special case
|
||||
if name == "NoneType":
|
||||
return type(None)
|
||||
raise pickle.UnpicklingError(
|
||||
f"Unsafe pickle: attempted to load '{module}.{name}'. "
|
||||
f"Only basic Python types are allowed for security reasons."
|
||||
)
|
||||
|
||||
|
||||
def _safe_pickle_load(file_path: str) -> Any:
|
||||
"""
|
||||
Safely load a pickle file using restricted unpickler.
|
||||
|
||||
Args:
|
||||
file_path: Path to the pickle file.
|
||||
|
||||
Returns:
|
||||
The deserialized object (only basic Python types allowed).
|
||||
|
||||
Raises:
|
||||
pickle.UnpicklingError: If the pickle contains unsafe types.
|
||||
"""
|
||||
with open(file_path, "rb") as f:
|
||||
return SafeUnpickler(f).load()
|
||||
|
||||
|
||||
def _validate_docstore_structure(data: Any) -> tuple:
|
||||
"""
|
||||
Validate that loaded data has the expected structure.
|
||||
|
||||
Args:
|
||||
data: The loaded data to validate.
|
||||
|
||||
Returns:
|
||||
Tuple of (docstore, index_to_id) if valid.
|
||||
|
||||
Raises:
|
||||
ValueError: If the data structure is invalid.
|
||||
"""
|
||||
if not isinstance(data, tuple) or len(data) != 2:
|
||||
raise ValueError("Invalid docstore format: expected tuple of (docstore, index_to_id)")
|
||||
|
||||
docstore, index_to_id = data
|
||||
|
||||
if not isinstance(docstore, dict):
|
||||
raise ValueError("Invalid docstore format: docstore must be a dict")
|
||||
|
||||
if not isinstance(index_to_id, dict):
|
||||
raise ValueError("Invalid docstore format: index_to_id must be a dict")
|
||||
|
||||
# Validate docstore entries
|
||||
for key, value in docstore.items():
|
||||
if not isinstance(key, str):
|
||||
raise ValueError(f"Invalid docstore key type: {type(key)}, expected str")
|
||||
if not isinstance(value, dict):
|
||||
raise ValueError(f"Invalid docstore value type: {type(value)}, expected dict")
|
||||
|
||||
# Validate index_to_id entries
|
||||
for key, value in index_to_id.items():
|
||||
if not isinstance(key, int):
|
||||
raise ValueError(f"Invalid index_to_id key type: {type(key)}, expected int")
|
||||
if not isinstance(value, str):
|
||||
raise ValueError(f"Invalid index_to_id value type: {type(value)}, expected str")
|
||||
|
||||
return docstore, index_to_id
|
||||
|
||||
|
||||
class OutputData(BaseModel):
|
||||
id: Optional[str] # memory id
|
||||
score: Optional[float] # distance
|
||||
@@ -73,9 +161,13 @@ class FAISS(VectorStoreBase):
|
||||
|
||||
# Try to load existing index if available
|
||||
index_path = f"{self.path}/{collection_name}.faiss"
|
||||
docstore_path = f"{self.path}/{collection_name}.pkl"
|
||||
if os.path.exists(index_path) and os.path.exists(docstore_path):
|
||||
self._load(index_path, docstore_path)
|
||||
json_docstore_path = f"{self.path}/{collection_name}.json"
|
||||
pkl_docstore_path = f"{self.path}/{collection_name}.pkl"
|
||||
|
||||
# Check for index file and either JSON (preferred) or legacy pickle docstore
|
||||
if os.path.exists(index_path) and (os.path.exists(json_docstore_path) or os.path.exists(pkl_docstore_path)):
|
||||
# _load will prefer JSON over pickle and auto-migrate
|
||||
self._load(index_path, pkl_docstore_path)
|
||||
else:
|
||||
self.create_col(collection_name)
|
||||
|
||||
@@ -83,34 +175,76 @@ class FAISS(VectorStoreBase):
|
||||
"""
|
||||
Load FAISS index and docstore from disk.
|
||||
|
||||
Supports both JSON (preferred) and legacy pickle formats. Pickle files are loaded
|
||||
using a restricted unpickler that only allows basic Python types to prevent
|
||||
arbitrary code execution (CVE mitigation).
|
||||
|
||||
Args:
|
||||
index_path (str): Path to FAISS index file.
|
||||
docstore_path (str): Path to docstore pickle file.
|
||||
docstore_path (str): Path to docstore file (.json or legacy .pkl).
|
||||
"""
|
||||
try:
|
||||
self.index = faiss.read_index(index_path)
|
||||
with open(docstore_path, "rb") as f:
|
||||
self.docstore, self.index_to_id = pickle.load(f)
|
||||
logger.info(f"Loaded FAISS index from {index_path} with {self.index.ntotal} vectors")
|
||||
|
||||
# Determine docstore format - prefer JSON over pickle
|
||||
json_docstore_path = docstore_path.replace(".pkl", ".json")
|
||||
|
||||
if os.path.exists(json_docstore_path):
|
||||
# Load from JSON (safe, preferred format)
|
||||
with open(json_docstore_path, "r", encoding="utf-8") as f:
|
||||
data = json.load(f)
|
||||
self.docstore = data.get("docstore", {})
|
||||
# JSON keys are always strings, convert back to int
|
||||
self.index_to_id = {int(k): v for k, v in data.get("index_to_id", {}).items()}
|
||||
logger.info(f"Loaded FAISS index from {index_path} with {self.index.ntotal} vectors (JSON format)")
|
||||
|
||||
elif os.path.exists(docstore_path):
|
||||
# Load from legacy pickle using safe unpickler
|
||||
# This prevents arbitrary code execution from malicious pickle files
|
||||
logger.warning(
|
||||
f"Loading legacy pickle docstore from {docstore_path}. "
|
||||
f"Consider migrating to JSON format for better security."
|
||||
)
|
||||
data = _safe_pickle_load(docstore_path)
|
||||
self.docstore, self.index_to_id = _validate_docstore_structure(data)
|
||||
logger.info(f"Loaded FAISS index from {index_path} with {self.index.ntotal} vectors (pickle format)")
|
||||
|
||||
# Auto-migrate to JSON format
|
||||
self._save()
|
||||
logger.info(f"Migrated docstore to JSON format: {json_docstore_path}")
|
||||
|
||||
else:
|
||||
raise FileNotFoundError(f"No docstore found at {docstore_path} or {json_docstore_path}")
|
||||
|
||||
except pickle.UnpicklingError as e:
|
||||
logger.error(f"Security error loading FAISS docstore: {e}")
|
||||
raise ValueError(f"Failed to load FAISS docstore: potentially malicious pickle file. {e}") from e
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to load FAISS index: {e}")
|
||||
|
||||
self.docstore = {}
|
||||
self.index_to_id = {}
|
||||
|
||||
def _save(self):
|
||||
"""Save FAISS index and docstore to disk."""
|
||||
"""Save FAISS index and docstore to disk using JSON format (secure)."""
|
||||
if not self.path or not self.index:
|
||||
return
|
||||
|
||||
try:
|
||||
os.makedirs(self.path, exist_ok=True)
|
||||
index_path = f"{self.path}/{self.collection_name}.faiss"
|
||||
docstore_path = f"{self.path}/{self.collection_name}.pkl"
|
||||
json_docstore_path = f"{self.path}/{self.collection_name}.json"
|
||||
|
||||
faiss.write_index(self.index, index_path)
|
||||
with open(docstore_path, "wb") as f:
|
||||
pickle.dump((self.docstore, self.index_to_id), f)
|
||||
|
||||
# Save docstore as JSON (safe format, no code execution risk)
|
||||
# JSON keys must be strings, so convert int keys to str
|
||||
data = {
|
||||
"docstore": self.docstore,
|
||||
"index_to_id": {str(k): v for k, v in self.index_to_id.items()},
|
||||
}
|
||||
with open(json_docstore_path, "w", encoding="utf-8") as f:
|
||||
json.dump(data, f, indent=2)
|
||||
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to save FAISS index: {e}")
|
||||
|
||||
@@ -417,12 +551,16 @@ class FAISS(VectorStoreBase):
|
||||
if self.path:
|
||||
try:
|
||||
index_path = f"{self.path}/{self.collection_name}.faiss"
|
||||
docstore_path = f"{self.path}/{self.collection_name}.pkl"
|
||||
json_docstore_path = f"{self.path}/{self.collection_name}.json"
|
||||
pkl_docstore_path = f"{self.path}/{self.collection_name}.pkl"
|
||||
|
||||
if os.path.exists(index_path):
|
||||
os.remove(index_path)
|
||||
if os.path.exists(docstore_path):
|
||||
os.remove(docstore_path)
|
||||
if os.path.exists(json_docstore_path):
|
||||
os.remove(json_docstore_path)
|
||||
# Also clean up legacy pickle files if they exist
|
||||
if os.path.exists(pkl_docstore_path):
|
||||
os.remove(pkl_docstore_path)
|
||||
|
||||
logger.info(f"Deleted collection {self.collection_name}")
|
||||
except Exception as e:
|
||||
|
||||
Reference in New Issue
Block a user