security: fix SQL and Cypher injection vulnerabilities in PGVector, Azure MySQL, and Neptune (#4878)

Co-authored-by: kartik-mem0 <kartik.labhshetwar@mem0.ai>

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]:

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)

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, '$.<key>')` 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)])
 

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://<graphid>'.")
 
+        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

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

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
+            )