Harden search filters against Cypher injection (#1312)

* harden search filter inputs

* validate entity node labels on save

* tighten security regression coverage

Author: danielchalef

graphiti_core/driver/falkordb_driver.py

@@ -66,6 +66,7 @@
 from graphiti_core.driver.operations.saga_node_ops import SagaNodeOperations
 from graphiti_core.driver.operations.search_ops import SearchOperations
 from graphiti_core.graph_queries import get_fulltext_indices, get_range_indices
+from graphiti_core.helpers import validate_group_ids
 from graphiti_core.utils.datetime_utils import convert_datetimes_to_strings
 
 logger = logging.getLogger(__name__)
@@ -397,6 +398,8 @@ def build_fulltext_query(
         - AND is implicit with space: (@group_id:value) (text)
         - OR uses pipe within parentheses: (@group_id:value1|value2)
         """
+        validate_group_ids(group_ids)
+
         if group_ids is None or len(group_ids) == 0:
             group_filter = ''
         else:

graphiti_core/driver/neo4j/operations/search_ops.py

@@ -32,7 +32,7 @@
     get_relationships_query,
     get_vector_cosine_func_query,
 )
-from graphiti_core.helpers import lucene_sanitize
+from graphiti_core.helpers import lucene_sanitize, validate_group_ids
 from graphiti_core.models.edges.edge_db_queries import get_entity_edge_return_query
 from graphiti_core.models.nodes.node_db_queries import (
     COMMUNITY_NODE_RETURN,
@@ -56,6 +56,8 @@ def _build_neo4j_fulltext_query(
     group_ids: list[str] | None = None,
     max_query_length: int = MAX_QUERY_LENGTH,
 ) -> str:
+    validate_group_ids(group_ids)
+
     group_ids_filter_list = [f'group_id:"{g}"' for g in group_ids] if group_ids is not None else []
     group_ids_filter = ''
     for f in group_ids_filter_list:

graphiti_core/errors.py

@@ -81,3 +81,15 @@ class GroupIdValidationError(GraphitiError):
     def __init__(self, group_id: str):
         self.message = f'group_id "{group_id}" must contain only alphanumeric characters, dashes, or underscores'
         super().__init__(self.message)
+
+
+class NodeLabelValidationError(GraphitiError, ValueError):
+    """Raised when a node label contains invalid characters."""
+
+    def __init__(self, node_labels: list[str]):
+        label_list = ', '.join(f'"{label}"' for label in node_labels)
+        self.message = (
+            'node_labels must start with a letter or underscore and contain only '
+            f'alphanumeric characters or underscores: {label_list}'
+        )
+        super().__init__(self.message)

graphiti_core/helpers.py

@@ -28,10 +28,12 @@
 from pydantic import BaseModel
 
 from graphiti_core.driver.driver import GraphProvider
-from graphiti_core.errors import GroupIdValidationError
+from graphiti_core.errors import GroupIdValidationError, NodeLabelValidationError
 
 load_dotenv()
 
+SAFE_CYPHER_IDENTIFIER_PATTERN = re.compile(r'^[A-Za-z_][A-Za-z0-9_]*$')
+
 USE_PARALLEL_RUNTIME = bool(os.getenv('USE_PARALLEL_RUNTIME', False))
 SEMAPHORE_LIMIT = int(os.getenv('SEMAPHORE_LIMIT', 20))
 DEFAULT_PAGE_LIMIT = 20
@@ -157,6 +159,33 @@ def validate_group_id(group_id: str | None) -> bool:
     return True
 
 
+def validate_group_ids(group_ids: list[str] | None) -> bool:
+    """Validate a list of group ids used by search paths."""
+
+    if group_ids is None:
+        return True
+
+    for group_id in group_ids:
+        validate_group_id(group_id)
+
+    return True
+
+
+def validate_node_labels(node_labels: list[str] | None) -> bool:
+    """Validate that node labels are safe to interpolate into Cypher label expressions."""
+
+    if not node_labels:
+        return True
+
+    invalid_labels = [
+        label for label in node_labels if not SAFE_CYPHER_IDENTIFIER_PATTERN.match(label)
+    ]
+    if invalid_labels:
+        raise NodeLabelValidationError(invalid_labels)
+
+    return True
+
+
 def validate_excluded_entity_types(
     excluded_entity_types: list[str] | None, entity_types: dict[str, type[BaseModel]] | None = None
 ) -> bool:

graphiti_core/models/nodes/node_db_queries.py

@@ -17,6 +17,14 @@
 from typing import Any
 
 from graphiti_core.driver.driver import GraphProvider
+from graphiti_core.helpers import validate_node_labels
+
+
+def _validate_entity_labels(labels: str | list[str]) -> list[str]:
+    resolved_labels = labels.split(':') if isinstance(labels, str) else labels
+    filtered_labels = [label for label in resolved_labels if label]
+    validate_node_labels(filtered_labels)
+    return filtered_labels
 
 
 def get_episode_node_save_query(provider: GraphProvider) -> str:
@@ -127,6 +135,9 @@ def get_episode_node_save_bulk_query(provider: GraphProvider) -> str:
 
 
 def get_entity_node_save_query(provider: GraphProvider, labels: str, has_aoss: bool = False) -> str:
+    validated_labels = _validate_entity_labels(labels)
+    labels = ':'.join(validated_labels)
+
     match provider:
         case GraphProvider.FALKORDB:
             return f"""
@@ -152,7 +163,7 @@ def get_entity_node_save_query(provider: GraphProvider, labels: str, has_aoss: b
             """
         case GraphProvider.NEPTUNE:
             label_subquery = ''
-            for label in labels.split(':'):
+            for label in validated_labels:
                 label_subquery += f' SET n:{label}\n'
             return f"""
                 MERGE (n:Entity {{uuid: $entity_data.uuid}})
@@ -183,6 +194,9 @@ def get_entity_node_save_query(provider: GraphProvider, labels: str, has_aoss: b
 def get_entity_node_save_bulk_query(
     provider: GraphProvider, nodes: list[dict], has_aoss: bool = False
 ) -> str | Any:
+    for node in nodes:
+        _validate_entity_labels(node.get('labels', []))
+
     match provider:
         case GraphProvider.FALKORDB:
             queries = []

graphiti_core/nodes.py

@@ -23,7 +23,7 @@
 from typing import Any
 from uuid import uuid4
 
-from pydantic import BaseModel, Field
+from pydantic import BaseModel, ConfigDict, Field, field_validator
 from typing_extensions import LiteralString
 
 from graphiti_core.driver.driver import (
@@ -32,7 +32,7 @@
 )
 from graphiti_core.embedder import EmbedderClient
 from graphiti_core.errors import NodeNotFoundError
-from graphiti_core.helpers import parse_db_date
+from graphiti_core.helpers import parse_db_date, validate_node_labels
 from graphiti_core.models.nodes.node_db_queries import (
     COMMUNITY_NODE_RETURN,
     COMMUNITY_NODE_RETURN_NEPTUNE,
@@ -94,6 +94,14 @@ class Node(BaseModel, ABC):
     labels: list[str] = Field(default_factory=list)
     created_at: datetime = Field(default_factory=lambda: utc_now())
 
+    model_config = ConfigDict(validate_assignment=True)
+
+    @field_validator('labels')
+    @classmethod
+    def validate_labels(cls, value: list[str]) -> list[str]:
+        validate_node_labels(value)
+        return value
+
     @abstractmethod
     async def save(self, driver: GraphDriver): ...
 

graphiti_core/search/search.py

@@ -24,7 +24,7 @@
 from graphiti_core.embedder.client import EMBEDDING_DIM
 from graphiti_core.errors import SearchRerankerError
 from graphiti_core.graphiti_types import GraphitiClients
-from graphiti_core.helpers import semaphore_gather
+from graphiti_core.helpers import semaphore_gather, validate_group_ids
 from graphiti_core.nodes import CommunityNode, EntityNode, EpisodicNode
 from graphiti_core.search.search_config import (
     DEFAULT_SEARCH_LIMIT,
@@ -77,6 +77,7 @@ async def search(
     driver: GraphDriver | None = None,
 ) -> SearchResults:
     start = time()
+    validate_group_ids(group_ids)
 
     driver = driver or clients.driver
     embedder = clients.embedder

graphiti_core/search/search_filters.py

@@ -18,9 +18,10 @@
 from enum import Enum
 from typing import Any
 
-from pydantic import BaseModel, Field
+from pydantic import BaseModel, Field, field_validator
 
 from graphiti_core.driver.driver import GraphProvider
+from graphiti_core.helpers import validate_node_labels
 
 
 class ComparisonOperator(Enum):
@@ -65,6 +66,12 @@ class SearchFilters(BaseModel):
     edge_uuids: list[str] | None = Field(default=None)
     property_filters: list[PropertyFilter] | None = Field(default=None)
 
+    @field_validator('node_labels')
+    @classmethod
+    def validate_node_label_filters(cls, value: list[str] | None) -> list[str] | None:
+        validate_node_labels(value)
+        return value
+
 
 def cypher_to_opensearch_operator(op: ComparisonOperator) -> str:
     mapping = {
@@ -84,6 +91,8 @@ def node_search_filter_query_constructor(
     filter_params: dict[str, Any] = {}
 
     if filters.node_labels is not None:
+        # Defense-in-depth for model_construct()/other validation bypasses.
+        validate_node_labels(filters.node_labels)
         if provider == GraphProvider.KUZU:
             node_label_filter = 'list_has_all(n.labels, $labels)'
             filter_params['labels'] = filters.node_labels
@@ -125,6 +134,8 @@ def edge_search_filter_query_constructor(
         filter_params['edge_uuids'] = filters.edge_uuids
 
     if filters.node_labels is not None:
+        # Defense-in-depth for model_construct()/other validation bypasses.
+        validate_node_labels(filters.node_labels)
         if provider == GraphProvider.KUZU:
             node_label_filter = (
                 'list_has_all(n.labels, $labels) AND list_has_all(m.labels, $labels)'

graphiti_core/search/search_utils.py

@@ -37,6 +37,7 @@
     lucene_sanitize,
     normalize_l2,
     semaphore_gather,
+    validate_group_ids,
 )
 from graphiti_core.models.edges.edge_db_queries import get_entity_edge_return_query
 from graphiti_core.models.nodes.node_db_queries import (
@@ -82,6 +83,8 @@ def calculate_cosine_similarity(vector1: list[float], vector2: list[float]) -> f
 
 
 def fulltext_query(query: str, group_ids: list[str] | None, driver: GraphDriver):
+    validate_group_ids(group_ids)
+
     if driver.provider == GraphProvider.KUZU:
         # Kuzu only supports simple queries.
         if len(query.split(' ')) > MAX_QUERY_LENGTH:

tests/test_node_label_security.py

@@ -0,0 +1,56 @@
+import pytest
+from pydantic import ValidationError
+
+from graphiti_core.driver.driver import GraphProvider
+from graphiti_core.errors import NodeLabelValidationError
+from graphiti_core.models.nodes.node_db_queries import (
+    get_entity_node_save_bulk_query,
+    get_entity_node_save_query,
+)
+from graphiti_core.nodes import EntityNode
+
+
+def test_entity_node_rejects_unsafe_labels():
+    with pytest.raises(ValidationError, match='node_labels must start with a letter or underscore'):
+        EntityNode(
+            name='Alice',
+            group_id='group',
+            labels=['Entity`) WITH n MATCH (x) DETACH DELETE x //'],
+        )
+
+
+def test_entity_node_assignment_rejects_unsafe_labels():
+    node = EntityNode(name='Alice', group_id='group', labels=['Person'])
+
+    with pytest.raises(ValidationError, match='node_labels must start with a letter or underscore'):
+        node.labels = ['Entity`) WITH n MATCH (x) DETACH DELETE x //']
+
+
+def test_entity_node_save_query_rejects_unsafe_labels_when_validation_is_bypassed():
+    with pytest.raises(
+        NodeLabelValidationError, match='node_labels must start with a letter or underscore'
+    ):
+        get_entity_node_save_query(
+            GraphProvider.NEO4J,
+            'Entity:Entity`) WITH n MATCH (x) DETACH DELETE x //',
+        )
+
+
+def test_entity_node_save_bulk_query_rejects_unsafe_labels_when_validation_is_bypassed():
+    with pytest.raises(
+        NodeLabelValidationError, match='node_labels must start with a letter or underscore'
+    ):
+        get_entity_node_save_bulk_query(
+            GraphProvider.FALKORDB,
+            [
+                {
+                    'uuid': 'node-1',
+                    'name': 'Alice',
+                    'group_id': 'group',
+                    'summary': 'summary',
+                    'created_at': '2024-01-01T00:00:00Z',
+                    'name_embedding': [0.1, 0.2],
+                    'labels': ['Entity', 'Entity`) WITH n MATCH (x) DETACH DELETE x //'],
+                }
+            ],
+        )