github/langchain
1
0
Fork 0
mirror of https://github.com/hwchase17/langchain.git synced 2025-08-06 03:27:55 +00:00
Code Issues Packages Projects Releases Wiki Activity

community[patch]: Patch tdidf retriever (CVE-2024-2057) (#18695)

Browse Source 
This is a patch for `CVE-2024-2057`:
https://www.cve.org/CVERecord?id=CVE-2024-2057

This affects users that: 

* Use the  `TFIDFRetriever`
* Attempt to de-serialize it from an untrusted source that contains a
malicious payload
This commit is contained in:
Eugene Yurtsev 2024-03-06 15:49:04 -05:00 committed by GitHub
parent 81cbf0f2fd
commit 0e52961562
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 35 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

30
libs/community/langchain_community/retrievers/tfidf.py
View File

@ -106,8 +106,26 @@ class TFIDFRetriever(BaseRetriever):
def load_local( def load_local(
cls, cls,
folder_path: str, folder_path: str,
*,
allow_dangerous_deserialization: bool = False,
file_name: str = "tfidf_vectorizer", file_name: str = "tfidf_vectorizer",
) -> TFIDFRetriever: ) -> TFIDFRetriever:
"""Load the retriever from local storage.
Args:
folder_path: Folder path to load from.
allow_dangerous_deserialization: Whether to allow dangerous deserialization.
Defaults to False.
The deserialization relies on .joblib and .pkl files, which can be
modified to deliver a malicious payload that results in execution of
arbitrary code on your machine. You will need to set this to `True` to
use deserialization. If you do this, make sure you trust the source of
the file.
file_name: File name to load from. Defaults to "tfidf_vectorizer".
Returns:
TFIDFRetriever: Loaded retriever.
"""
try: try:
import joblib import joblib
except ImportError: except ImportError:
@ -115,6 +133,18 @@ class TFIDFRetriever(BaseRetriever):
"Could not import joblib, please install with `pip install joblib`." "Could not import joblib, please install with `pip install joblib`."
) )
if not allow_dangerous_deserialization:
raise ValueError(
"The de-serialization of this retriever is based on .joblib and "
".pkl files."
"Such files can be modified to deliver a malicious payload that "
"results in execution of arbitrary code on your machine."
"You will need to set `allow_dangerous_deserialization` to `True` to "
"load this retriever. If you do this, make sure you trust the source "
"of the file, and you are responsible for validating the the file "
"came from a trusted source."
)
path = Path(folder_path) path = Path(folder_path)
# Load vectorizer with joblib load. # Load vectorizer with joblib load.

5
libs/community/tests/unit_tests/retrievers/test_tfidf.py
View File

@ -56,6 +56,11 @@ def test_save_local_load_local() -> None:
loaded_tfidf_retriever = TFIDFRetriever.load_local( loaded_tfidf_retriever = TFIDFRetriever.load_local(
folder_path=temp_folder, folder_path=temp_folder,
file_name=file_name, file_name=file_name,
# Not a realistic security risk in this case.
# OK to allow for testing purposes.
# If the file has been compromised during this test, there's
# a much bigger problem.
allow_dangerous_deserialization=True,
) )
assert len(loaded_tfidf_retriever.docs) == 3 assert len(loaded_tfidf_retriever.docs) == 3
assert loaded_tfidf_retriever.tfidf_array.toarray().shape == (3, 5) assert loaded_tfidf_retriever.tfidf_array.toarray().shape == (3, 5)