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

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

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