From bc9d45b476101e441cb1cc602dea03eb29232de4 Mon Sep 17 00:00:00 2001
From: Eugene Yurtsev <eyurtsev@gmail.com>
Date: Fri, 25 Jul 2025 13:01:13 -0400
Subject: [PATCH] fix(checkpoint-sqlite): add validation to filter keys in sql
 store (#5666)

This PR adds validation to keys used in filtering logic in the SQLite store implementation.
---
 .../langgraph/store/sqlite/base.py            | 24 +++++++++++++++++++
 libs/checkpoint-sqlite/tests/test_store.py    | 20 ++++++++++++++++
 2 files changed, 44 insertions(+)

diff --git a/libs/checkpoint-sqlite/langgraph/store/sqlite/base.py b/libs/checkpoint-sqlite/langgraph/store/sqlite/base.py
index db161d667..36e18d334 100644
--- a/libs/checkpoint-sqlite/langgraph/store/sqlite/base.py
+++ b/libs/checkpoint-sqlite/langgraph/store/sqlite/base.py
@@ -3,6 +3,7 @@ from __future__ import annotations
 import concurrent.futures
 import datetime
 import logging
+import re
 import sqlite3
 import threading
 from collections import defaultdict
@@ -107,6 +108,23 @@ def _decode_ns_text(namespace: str) -> tuple[str, ...]:
     return tuple(namespace.split("."))
 
 
+def _validate_filter_key(key: str) -> None:
+    """Validate that a filter key is safe for use in SQL queries.
+
+    Args:
+        key: The filter key to validate
+
+    Raises:
+        ValueError: If the key contains invalid characters that could enable SQL injection
+    """
+    # Allow alphanumeric characters, underscores, dots, and hyphens
+    # This covers typical JSON property names while preventing SQL injection
+    if not re.match(r"^[a-zA-Z0-9_.-]+$", key):
+        raise ValueError(
+            f"Invalid filter key: '{key}'. Filter keys must contain only alphanumeric characters, underscores, dots, and hyphens."
+        )
+
+
 def _json_loads(content: bytes | str | orjson.Fragment) -> Any:
     if isinstance(content, orjson.Fragment):
         if hasattr(content, "buf"):
@@ -372,6 +390,8 @@ class BaseSqliteStore:
             filter_conditions = []
             if op.filter:
                 for key, value in op.filter.items():
+                    _validate_filter_key(key)
+
                     if isinstance(value, dict):
                         for op_name, val in value.items():
                             condition, filter_params_ = self._get_filter_condition(
@@ -622,6 +642,8 @@ class BaseSqliteStore:
 
     def _get_filter_condition(self, key: str, op: str, value: Any) -> tuple[str, list]:
         """Helper to generate filter conditions."""
+        _validate_filter_key(key)
+
         # We need to properly format values for SQLite JSON extraction comparison
         if op == "$eq":
             if isinstance(value, str):
@@ -858,6 +880,8 @@ class SqliteStore(BaseSqliteStore, BaseStore):
 
     def _get_filter_condition(self, key: str, op: str, value: Any) -> tuple[str, list]:
         """Helper to generate filter conditions."""
+        _validate_filter_key(key)
+
         # We need to properly format values for SQLite JSON extraction comparison
         if op == "$eq":
             if isinstance(value, str):
diff --git a/libs/checkpoint-sqlite/tests/test_store.py b/libs/checkpoint-sqlite/tests/test_store.py
index 135624dc9..590394375 100644
--- a/libs/checkpoint-sqlite/tests/test_store.py
+++ b/libs/checkpoint-sqlite/tests/test_store.py
@@ -1047,3 +1047,23 @@ def test_search_items(
         for ns in test_namespaces:
             key = f"item_{ns[-1]}"
             store.delete(ns, key)
+
+
+def test_sql_injection_vulnerability(store: SqliteStore) -> None:
+    """Test that SQL injection via malicious filter keys is prevented."""
+    # Add public and private documents
+    store.put(("docs",), "public", {"access": "public", "data": "public info"})
+    store.put(
+        ("docs",), "private", {"access": "private", "data": "secret", "password": "123"}
+    )
+
+    # Normal query - returns 1 public document
+    normal = store.search(("docs",), filter={"access": "public"})
+    assert len(normal) == 1
+    assert normal[0].value["access"] == "public"
+
+    # SQL injection attempt via malicious key should raise ValueError
+    malicious_key = "access') = 'public' OR '1'='1' OR json_extract(value, '$."
+
+    with pytest.raises(ValueError, match="Invalid filter key"):
+        store.search(("docs",), filter={malicious_key: "dummy"})