fix(elasticsearch): validate filter keys and values to prevent term injection (#5980)
Signed-off-by: Hrushikesh Yadav <yadavhrushikesh65@gmail.com>
This commit is contained in:
@@ -1,4 +1,5 @@
|
|||||||
import logging
|
import logging
|
||||||
|
import re
|
||||||
from typing import Any, Dict, List, Optional
|
from typing import Any, Dict, List, Optional
|
||||||
|
|
||||||
try:
|
try:
|
||||||
@@ -21,6 +22,19 @@ class OutputData(BaseModel):
|
|||||||
payload: Dict
|
payload: Dict
|
||||||
|
|
||||||
|
|
||||||
|
_SAFE_FILTER_KEY = re.compile(r"^[a-zA-Z_][a-zA-Z0-9_]*$")
|
||||||
|
|
||||||
|
|
||||||
|
def _validate_filter(key: str, value: Any) -> None:
|
||||||
|
if not isinstance(key, str) or not _SAFE_FILTER_KEY.match(key):
|
||||||
|
raise ValueError(f"Invalid filter key: {key!r}")
|
||||||
|
if not isinstance(value, (str, int, float, bool)):
|
||||||
|
raise ValueError(
|
||||||
|
f"Filter value for {key!r} must be str, int, float, or bool, "
|
||||||
|
f"got {type(value).__name__}"
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
class ElasticsearchDB(VectorStoreBase):
|
class ElasticsearchDB(VectorStoreBase):
|
||||||
def __init__(self, **kwargs):
|
def __init__(self, **kwargs):
|
||||||
config = ElasticsearchConfig(**kwargs)
|
config = ElasticsearchConfig(**kwargs)
|
||||||
@@ -154,6 +168,7 @@ class ElasticsearchDB(VectorStoreBase):
|
|||||||
if filters:
|
if filters:
|
||||||
filter_conditions = []
|
filter_conditions = []
|
||||||
for key, value in filters.items():
|
for key, value in filters.items():
|
||||||
|
_validate_filter(key, value)
|
||||||
filter_conditions.append({"term": {f"metadata.{key}": value}})
|
filter_conditions.append({"term": {f"metadata.{key}": value}})
|
||||||
search_query["knn"]["filter"] = {"bool": {"must": filter_conditions}}
|
search_query["knn"]["filter"] = {"bool": {"must": filter_conditions}}
|
||||||
|
|
||||||
@@ -192,6 +207,7 @@ class ElasticsearchDB(VectorStoreBase):
|
|||||||
if filters:
|
if filters:
|
||||||
filter_conditions = []
|
filter_conditions = []
|
||||||
for key, value in filters.items():
|
for key, value in filters.items():
|
||||||
|
_validate_filter(key, value)
|
||||||
filter_conditions.append({"term": {f"metadata.{key}": value}})
|
filter_conditions.append({"term": {f"metadata.{key}": value}})
|
||||||
bool_query["filter"] = filter_conditions
|
bool_query["filter"] = filter_conditions
|
||||||
|
|
||||||
@@ -262,6 +278,7 @@ class ElasticsearchDB(VectorStoreBase):
|
|||||||
if filters:
|
if filters:
|
||||||
filter_conditions = []
|
filter_conditions = []
|
||||||
for key, value in filters.items():
|
for key, value in filters.items():
|
||||||
|
_validate_filter(key, value)
|
||||||
filter_conditions.append({"term": {f"metadata.{key}": value}})
|
filter_conditions.append({"term": {f"metadata.{key}": value}})
|
||||||
query["query"] = {"bool": {"must": filter_conditions}}
|
query["query"] = {"bool": {"must": filter_conditions}}
|
||||||
|
|
||||||
|
|||||||
@@ -10,7 +10,7 @@ except ImportError:
|
|||||||
raise ImportError("Elasticsearch requires extra dependencies. Install with `pip install elasticsearch`") from None
|
raise ImportError("Elasticsearch requires extra dependencies. Install with `pip install elasticsearch`") from None
|
||||||
|
|
||||||
from mem0.configs.vector_stores.elasticsearch import ElasticsearchConfig
|
from mem0.configs.vector_stores.elasticsearch import ElasticsearchConfig
|
||||||
from mem0.vector_stores.elasticsearch import ElasticsearchDB, OutputData
|
from mem0.vector_stores.elasticsearch import ElasticsearchDB, OutputData, _validate_filter
|
||||||
|
|
||||||
|
|
||||||
class TestElasticsearchDB(unittest.TestCase):
|
class TestElasticsearchDB(unittest.TestCase):
|
||||||
@@ -360,3 +360,33 @@ class TestElasticsearchDB(unittest.TestCase):
|
|||||||
with self.assertRaises(ValueError):
|
with self.assertRaises(ValueError):
|
||||||
config = {**base_config, "headers": headers}
|
config = {**base_config, "headers": headers}
|
||||||
ElasticsearchConfig(**config)
|
ElasticsearchConfig(**config)
|
||||||
|
|
||||||
|
def test_filter_rejects_dict_value(self):
|
||||||
|
with self.assertRaises(ValueError):
|
||||||
|
_validate_filter("user_id", {"value": "alice", "boost": 0})
|
||||||
|
|
||||||
|
def test_filter_rejects_list_value(self):
|
||||||
|
with self.assertRaises(ValueError):
|
||||||
|
_validate_filter("user_id", ["alice", "bob"])
|
||||||
|
|
||||||
|
def test_filter_rejects_invalid_key(self):
|
||||||
|
with self.assertRaises(ValueError):
|
||||||
|
_validate_filter("user_id; DROP", "alice")
|
||||||
|
|
||||||
|
def test_filter_accepts_scalar_values(self):
|
||||||
|
_validate_filter("user_id", "alice")
|
||||||
|
_validate_filter("count", 42)
|
||||||
|
_validate_filter("score", 0.95)
|
||||||
|
_validate_filter("active", True)
|
||||||
|
|
||||||
|
def test_search_with_dict_filter_raises(self):
|
||||||
|
with self.assertRaises(ValueError):
|
||||||
|
self.es_db.search(
|
||||||
|
query="test",
|
||||||
|
vectors=[0.1] * 1536,
|
||||||
|
filters={"user_id": {"value": "*", "case_insensitive": True}},
|
||||||
|
)
|
||||||
|
|
||||||
|
def test_list_with_dict_filter_raises(self):
|
||||||
|
with self.assertRaises(ValueError):
|
||||||
|
self.es_db.list(filters={"user_id": {"$ne": ""}})
|
||||||
|
|||||||
Reference in New Issue
Block a user