diff --git a/mem0/configs/vector_stores/azure_mysql.py b/mem0/configs/vector_stores/azure_mysql.py index 861452ac8..40919e676 100644 --- a/mem0/configs/vector_stores/azure_mysql.py +++ b/mem0/configs/vector_stores/azure_mysql.py @@ -1,6 +1,11 @@ +import re from typing import Any, Dict, Optional -from pydantic import BaseModel, ConfigDict, Field, model_validator +from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator + +# Rejecting any identifier that doesn't match this pattern is what keeps the +# f-string table/column interpolations in mem0/vector_stores/azure_mysql.py safe. +_VALID_SQL_IDENTIFIER = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*$") class AzureMySQLConfig(BaseModel): @@ -26,6 +31,15 @@ class AzureMySQLConfig(BaseModel): description="Pre-configured connection pool object (overrides other connection parameters)" ) + @field_validator("collection_name") + def validate_collection_name(cls, v): + if not _VALID_SQL_IDENTIFIER.match(v): + raise ValueError( + f"Invalid collection_name: {v!r}. Must start with a letter or underscore and " + "contain only letters, digits, and underscores." + ) + return v + @model_validator(mode="before") @classmethod def check_auth(cls, values: Dict[str, Any]) -> Dict[str, Any]: diff --git a/mem0/configs/vector_stores/neptune.py b/mem0/configs/vector_stores/neptune.py index 06132f9ac..3d441c1a5 100644 --- a/mem0/configs/vector_stores/neptune.py +++ b/mem0/configs/vector_stores/neptune.py @@ -5,7 +5,13 @@ This module provides configuration settings for integrating with Amazon Neptune as a vector store backend for Mem0's memory layer. """ -from pydantic import BaseModel, ConfigDict, Field +import re + +from pydantic import BaseModel, ConfigDict, Field, field_validator + +# Rejecting any identifier that doesn't match this pattern is what keeps the +# f-string label interpolations in mem0/vector_stores/neptune_analytics.py safe. +_VALID_CYPHER_IDENTIFIER = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*$") class NeptuneAnalyticsConfig(BaseModel): @@ -22,4 +28,13 @@ class NeptuneAnalyticsConfig(BaseModel): collection_name: str = Field("mem0", description="Default name for the collection") endpoint: str = Field("endpoint", description="Graph ID for the runtime") + @field_validator("collection_name") + def validate_collection_name(cls, v): + if not _VALID_CYPHER_IDENTIFIER.match(v): + raise ValueError( + f"Invalid collection_name: {v!r}. Must start with a letter or underscore and " + "contain only letters, digits, and underscores." + ) + return v + model_config = ConfigDict(arbitrary_types_allowed=False) diff --git a/mem0/vector_stores/azure_mysql.py b/mem0/vector_stores/azure_mysql.py index 263ffc12c..b4e6dc4cb 100644 --- a/mem0/vector_stores/azure_mysql.py +++ b/mem0/vector_stores/azure_mysql.py @@ -37,6 +37,10 @@ def _validate_identifier(name: str, label: str = "identifier") -> str: ) return name +# Filter keys are interpolated into `JSON_EXTRACT(payload, '$.')` strings, +# so we whitelist identifier-safe keys and skip anything else (logged as a warning). +_VALID_FILTER_KEY = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*$") + class OutputData(BaseModel): id: Optional[str] @@ -290,6 +294,9 @@ class AzureMySQL(VectorStoreBase): if filters: for k, v in filters.items(): + if not _VALID_FILTER_KEY.match(k): + logger.warning("Skipping invalid filter key: %r", k) + continue filter_conditions.append("JSON_EXTRACT(payload, %s) = %s") filter_params.extend([f"$.{k}", json.dumps(v)]) @@ -346,6 +353,9 @@ class AzureMySQL(VectorStoreBase): if filters: for k, v in filters.items(): + if not _VALID_FILTER_KEY.match(k): + logger.warning("Skipping invalid filter key: %r", k) + continue filter_conditions.append("JSON_EXTRACT(payload, %s) = %s") filter_params.extend([f"$.{k}", json.dumps(v)]) @@ -502,6 +512,9 @@ class AzureMySQL(VectorStoreBase): if filters: for k, v in filters.items(): + if not _VALID_FILTER_KEY.match(k): + logger.warning("Skipping invalid filter key: %r", k) + continue filter_conditions.append("JSON_EXTRACT(payload, %s) = %s") filter_params.extend([f"$.{k}", json.dumps(v)]) diff --git a/mem0/vector_stores/neptune_analytics.py b/mem0/vector_stores/neptune_analytics.py index 3b46f8af6..905b5de57 100644 --- a/mem0/vector_stores/neptune_analytics.py +++ b/mem0/vector_stores/neptune_analytics.py @@ -16,6 +16,7 @@ from mem0.vector_stores.base import VectorStoreBase logger = logging.getLogger(__name__) _SAFE_FILTER_KEY = re.compile(r"^[a-zA-Z_~][a-zA-Z0-9_]*$") +_VALID_IDENTIFIER = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*$") def _validate_filter(key: str, value: Any) -> None: @@ -73,6 +74,12 @@ class NeptuneAnalyticsVector(VectorStoreBase): if not endpoint.startswith("neptune-graph://"): raise ValueError("Please provide 'endpoint' with the format as 'neptune-graph://'.") + if not _VALID_IDENTIFIER.match(collection_name): + raise ValueError( + f"Invalid collection_name: {collection_name!r}. Must start with a letter or underscore and " + "contain only letters, digits, and underscores." + ) + graph_id = endpoint.replace("neptune-graph://", "") self.graph = NeptuneAnalyticsGraph(graph_id) self.collection_name = self._COLLECTION_PREFIX + collection_name diff --git a/tests/vector_stores/test_azure_mysql.py b/tests/vector_stores/test_azure_mysql.py index 05b61dc4c..c110d1edd 100644 --- a/tests/vector_stores/test_azure_mysql.py +++ b/tests/vector_stores/test_azure_mysql.py @@ -2,7 +2,9 @@ import json from unittest.mock import Mock, patch import pytest +from pydantic import ValidationError +from mem0.configs.vector_stores.azure_mysql import AzureMySQLConfig from mem0.vector_stores.azure_mysql import AzureMySQL, OutputData @@ -268,3 +270,63 @@ def test_output_data_model(): assert data.id == "test_id" assert data.score == 0.95 assert data.payload == {"text": "test"} + + +INJECTION_PAYLOADS = [ + "memories; DROP TABLE users; --", + "memories` OR 1=1; --", + "memories:Label {prop: 'val'}) DELETE n; --", + "valid_name OR 1=1", + "1_starts_with_digit", + "has space", + "", +] + +class TestAzureMySQLConfigCollectionNameValidation: + def test_accepts_valid_identifier(self): + config = AzureMySQLConfig( + collection_name="valid_name", + host="host", + user="user", + database="db", + password="pw", + ) + assert config.collection_name == "valid_name" + + @pytest.mark.parametrize("payload", INJECTION_PAYLOADS) + def test_rejects_injection_payload(self, payload): + with pytest.raises(ValidationError, match="Invalid collection_name"): + AzureMySQLConfig( + collection_name=payload, + host="host", + user="user", + database="db", + password="pw", + ) + + +class TestAzureMySQLFilterKeySanitization: + """Non-identifier filter keys are interpolated into the JSON_EXTRACT path, + so every method that builds a filter clause must drop them before they run.""" + + _FILTERS = {"user_id": "u1", "evil') OR '1'='1": "x"} + + @staticmethod + def _executed_blob(instance): + cursor = instance.connection_pool.connection().cursor() + return " ".join(repr(c.args) + repr(c.kwargs) for c in cursor.execute.call_args_list) + + def test_search_drops_invalid_key(self, azure_mysql_instance): + azure_mysql_instance.search(query="q", vectors=[0.1, 0.2, 0.3], filters=self._FILTERS) + blob = self._executed_blob(azure_mysql_instance) + assert "$.user_id" in blob and "OR '1'='1" not in blob + + def test_keyword_search_drops_invalid_key(self, azure_mysql_instance): + azure_mysql_instance.keyword_search(query="q", filters=self._FILTERS) + blob = self._executed_blob(azure_mysql_instance) + assert "$.user_id" in blob and "OR '1'='1" not in blob + + def test_list_drops_invalid_key(self, azure_mysql_instance): + azure_mysql_instance.list(filters=self._FILTERS) + blob = self._executed_blob(azure_mysql_instance) + assert "$.user_id" in blob and "OR '1'='1" not in blob diff --git a/tests/vector_stores/test_neptune_analytics.py b/tests/vector_stores/test_neptune_analytics.py index 0b067a1b8..41cd841df 100644 --- a/tests/vector_stores/test_neptune_analytics.py +++ b/tests/vector_stores/test_neptune_analytics.py @@ -4,7 +4,9 @@ import sys import pytest from dotenv import load_dotenv +from pydantic import ValidationError +from mem0.configs.vector_stores.neptune import NeptuneAnalyticsConfig from mem0.utils.factory import VectorStoreFactory from mem0.vector_stores.neptune_analytics import ( NeptuneAnalyticsVector, @@ -237,3 +239,45 @@ class TestNeptuneFilterValidation: NeptuneAnalyticsVector._get_node_filter_clause( {"user_id": {"$ne": ""}} ) + +INJECTION_PAYLOADS = [ + "memories; DROP TABLE users; --", + "memories` OR 1=1; --", + "memories:Label {prop: 'val'}) DELETE n; --", + "valid_name OR 1=1", + "1_starts_with_digit", + "has space", + "", +] + +class TestNeptuneAnalyticsConfigCollectionNameValidation: + def test_accepts_valid_identifier(self): + config = NeptuneAnalyticsConfig(collection_name="valid_name") + assert config.collection_name == "valid_name" + + @pytest.mark.parametrize("payload", INJECTION_PAYLOADS) + def test_rejects_injection_payload(self, payload): + with pytest.raises(ValidationError, match="Invalid collection_name"): + NeptuneAnalyticsConfig(collection_name=payload) + +class TestNeptuneAnalyticsVectorInitValidation: + def test_accepts_valid_identifier(self, monkeypatch): + from mem0.vector_stores.neptune_analytics import NeptuneAnalyticsVector + monkeypatch.setattr("mem0.vector_stores.neptune_analytics.NeptuneAnalyticsGraph", lambda *args, **kwargs: None) + + vec = NeptuneAnalyticsVector( + endpoint="neptune-graph://test", + collection_name="valid_name" + ) + assert vec.collection_name.endswith("valid_name") + + @pytest.mark.parametrize("payload", INJECTION_PAYLOADS) + def test_rejects_injection_payload_in_init(self, payload, monkeypatch): + from mem0.vector_stores.neptune_analytics import NeptuneAnalyticsVector + monkeypatch.setattr("mem0.vector_stores.neptune_analytics.NeptuneAnalyticsGraph", lambda *args, **kwargs: None) + + with pytest.raises(ValueError, match="Invalid collection_name"): + NeptuneAnalyticsVector( + endpoint="neptune-graph://test", + collection_name=payload + )