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