fix(checkpoint-sqlite): harden (#6565)

harden

libs/checkpoint-sqlite/tests/test_aiosqlite.py

@@ -113,4 +113,78 @@ class TestAsyncSqliteSaver:
                 search_results_5[1].config["configurable"]["checkpoint_ns"],
             } == {"", "inner"}
 
-            # TODO: test before and limit params
+            # Test limit param
+            search_results_6 = [
+                c
+                async for c in saver.alist(
+                    {"configurable": {"thread_id": "thread-2"}}, limit=1
+                )
+            ]
+            assert len(search_results_6) == 1
+            assert search_results_6[0].config["configurable"]["thread_id"] == "thread-2"
+
+            # Test before param
+            search_results_7 = [
+                c async for c in saver.alist(None, before=search_results_5[1].config)
+            ]
+            assert len(search_results_7) == 1
+            assert search_results_7[0].config["configurable"]["thread_id"] == "thread-1"
+
+    async def test_limit_parameter_sql_injection_prevention(self) -> None:
+        """Test that the limit parameter properly uses parameterized queries to prevent SQL injection."""
+        async with AsyncSqliteSaver.from_conn_string(":memory:") as saver:
+            # Setup: Create multiple checkpoints
+            for i in range(5):
+                config: RunnableConfig = {
+                    "configurable": {
+                        "thread_id": f"thread-{i}",
+                        "checkpoint_ns": "",
+                    }
+                }
+                checkpoint = empty_checkpoint()
+                metadata: CheckpointMetadata = {"index": i}
+                await saver.aput(config, checkpoint, metadata, {})
+
+            # Test that limit works correctly with valid integer
+            results = [c async for c in saver.alist(None, limit=2)]
+            assert len(results) == 2
+
+            # Test that limit=0 returns no results
+            results = [c async for c in saver.alist(None, limit=0)]
+            assert len(results) == 0
+
+            # Test that limit=None returns all results
+            results = [c async for c in saver.alist(None, limit=None)]
+            assert len(results) == 5
+
+            # Test explicit SQL injection attempt via limit parameter
+            # Even if type checking is bypassed and a malicious string is passed,
+            # the parameterized query will treat it as a value, not SQL code
+            # This would cause an error (can't convert string to int for LIMIT),
+            # which is the correct secure behavior
+            malicious_limits = [
+                "1; DROP TABLE checkpoints; --",
+                "1 OR 1=1",
+                "999999 UNION SELECT * FROM checkpoints",
+            ]
+
+            for malicious_limit in malicious_limits:
+                # The parameterized query should safely reject non-integer limits
+                # or convert them in a way that prevents SQL injection
+                try:
+                    # Bypass type checking by casting
+                    results = [
+                        c
+                        async for c in saver.alist(None, limit=malicious_limit)  # type: ignore
+                    ]
+                    # If it doesn't raise an error, it should at least not execute the injection
+                    # SQLite's parameter binding will try to convert the string to an integer
+                    # which will either fail or treat it as 0
+                except Exception:
+                    # Expected: SQLite should reject invalid limit values
+                    pass
+
+            # Verify the checkpoints table still exists and has all data
+            # (would have been dropped if injection succeeded)
+            results = [c async for c in saver.alist(None, limit=None)]
+            assert len(results) == 5

libs/checkpoint-sqlite/tests/test_sqlite.py

@@ -182,3 +182,128 @@ class TestSqliteSaver:
             with pytest.raises(NotImplementedError, match="AsyncSqliteSaver"):
                 async for _ in saver.alist(self.config_1):
                     pass
+
+    def test_metadata_predicate_sql_injection_prevention(self) -> None:
+        """Test that _metadata_predicate rejects malicious filter keys."""
+        # Test various SQL injection payloads
+        malicious_keys = [
+            "x') OR '1'='1",  # Boolean-based injection
+            "x') OR 1=1 --",  # Comment-based injection
+            "x') UNION SELECT 1,2,3,4,5,6,7 --",  # UNION-based injection
+            "access') = 'public' OR '1'='1' OR json_extract(value, '$.",  # Complex injection
+            "'; DROP TABLE checkpoints; --",  # Destructive injection
+        ]
+
+        for malicious_key in malicious_keys:
+            with pytest.raises(ValueError, match="Invalid filter key"):
+                _metadata_predicate({malicious_key: "dummy"})
+
+    def test_checkpoint_search_sql_injection_prevention(self) -> None:
+        """Test that SQL injection via malicious filter keys is prevented in checkpoint search."""
+        with SqliteSaver.from_conn_string(":memory:") as saver:
+            # Setup: Create checkpoints with different metadata
+            config_public: RunnableConfig = {
+                "configurable": {
+                    "thread_id": "thread-public",
+                    "checkpoint_ns": "",
+                }
+            }
+            config_private: RunnableConfig = {
+                "configurable": {
+                    "thread_id": "thread-private",
+                    "checkpoint_ns": "",
+                }
+            }
+
+            checkpoint_public = empty_checkpoint()
+            checkpoint_private = empty_checkpoint()
+
+            metadata_public: CheckpointMetadata = {
+                "access": "public",
+                "data": "public information",
+            }
+            metadata_private: CheckpointMetadata = {
+                "access": "private",
+                "data": "secret information",
+                "password": "secret123",
+            }
+
+            saver.put(config_public, checkpoint_public, metadata_public, {})
+            saver.put(config_private, checkpoint_private, metadata_private, {})
+
+            # Normal query - should return only public checkpoint
+            normal_results = list(saver.list(None, filter={"access": "public"}))
+            assert len(normal_results) == 1
+            assert normal_results[0].metadata["access"] == "public"
+
+            # SQL injection attempt should raise ValueError
+            malicious_key = (
+                "access') = 'public' OR '1'='1' OR json_extract(metadata, '$."
+            )
+
+            with pytest.raises(ValueError, match="Invalid filter key"):
+                list(saver.list(None, filter={malicious_key: "dummy"}))
+
+    def test_limit_parameter_sql_injection_prevention(self) -> None:
+        """Test that the limit parameter properly uses parameterized queries to prevent SQL injection."""
+        with SqliteSaver.from_conn_string(":memory:") as saver:
+            # Setup: Create multiple checkpoints
+            for i in range(5):
+                config: RunnableConfig = {
+                    "configurable": {
+                        "thread_id": f"thread-{i}",
+                        "checkpoint_ns": "",
+                    }
+                }
+                checkpoint = empty_checkpoint()
+                metadata: CheckpointMetadata = {"index": i}
+                saver.put(config, checkpoint, metadata, {})
+
+            # Test that limit works correctly with valid integer
+            results = list(saver.list(None, limit=2))
+            assert len(results) == 2
+
+            # Test that limit=0 returns no results
+            results = list(saver.list(None, limit=0))
+            assert len(results) == 0
+
+            # Test that limit=None returns all results
+            results = list(saver.list(None, limit=None))
+            assert len(results) == 5
+
+    def test_metadata_filter_keys_with_hyphens_and_digits(self) -> None:
+        """Metadata keys with hyphens and digit-start should be filterable.
+
+        This exposes incorrect JSON path handling (unquoted segments) by asserting
+        that such filters successfully match saved checkpoints.
+        """
+        with SqliteSaver.from_conn_string(":memory:") as saver:
+            config: RunnableConfig = {
+                "configurable": {
+                    "thread_id": "thread-hyphen-digit",
+                    "checkpoint_ns": "",
+                }
+            }
+            checkpoint = empty_checkpoint()
+            metadata: CheckpointMetadata = {
+                "access-level": "public",
+                "user": {"access-level": "nested", "123abc": "ok2"},
+                "123abc": "ok",
+            }
+            saver.put(config, checkpoint, metadata, {})
+
+            # Top-level hyphenated key
+            results = list(saver.list(None, filter={"access-level": "public"}))
+            assert len(results) == 1
+
+            # Nested hyphenated key via dotted path
+            results = list(saver.list(None, filter={"user.access-level": "nested"}))
+            assert len(results) == 1
+
+            # Top-level digit-starting key
+            results = list(saver.list(None, filter={"123abc": "ok"}))
+            assert len(results) == 1
+
+            # Nested digit-starting key via dotted path
+            results = list(saver.list(None, filter={"user.123abc": "ok2"}))
+            assert len(results) == 1

libs/checkpoint-sqlite/tests/test_store.py

@@ -1069,6 +1069,141 @@ def test_sql_injection_vulnerability(store: SqliteStore) -> None:
         store.search(("docs",), filter={malicious_key: "dummy"})
 
 
+def test_sql_injection_filter_values(store: SqliteStore) -> None:
+    """Test that SQL injection via malicious filter values is properly escaped."""
+    # Setup: Create documents with different access levels
+    store.put(("docs",), "doc1", {"access": "public", "title": "Public Document"})
+    store.put(("docs",), "doc2", {"access": "private", "title": "Private Document"})
+    store.put(("docs",), "doc3", {"access": "secret", "title": "Secret Document"})
+
+    # Test 1: Basic SQL injection attempt with single quote
+    malicious_value = "public' OR '1'='1"
+    results = store.search(("docs",), filter={"access": malicious_value})
+    # Should return 0 results because the malicious value is escaped and won't match anything
+    assert len(results) == 0, "SQL injection via string value should be blocked"
+
+    # Test 2: SQL injection with comment
+    malicious_value = "public'; --"
+    results = store.search(("docs",), filter={"access": malicious_value})
+    assert len(results) == 0, "SQL comment injection should be blocked"
+
+    # Test 3: UNION injection attempt
+    malicious_value = "public' UNION SELECT * FROM store --"
+    results = store.search(("docs",), filter={"access": malicious_value})
+    assert len(results) == 0, "UNION injection should be blocked"
+
+    # Test 4: Parameterized queries handle strings with null bytes and SQL injection attempts safely
+    malicious_value = "public\x00' OR '1'='1"
+    results = store.search(("docs",), filter={"access": malicious_value})
+    assert len(results) == 0, (
+        "Parameterized queries treat injection attempts as literal strings"
+    )
+
+    # Test 5: Multiple single quotes
+    malicious_value = "''''"
+    results = store.search(("docs",), filter={"access": malicious_value})
+    assert len(results) == 0, "Multiple quotes should be handled safely"
+
+    # Test 6: Legitimate value with single quote should work
+    store.put(("docs",), "doc4", {"title": "O'Brien's Document", "access": "public"})
+    results = store.search(("docs",), filter={"title": "O'Brien's Document"})
+    assert len(results) == 1, "Legitimate single quotes should work"
+    assert results[0].value["title"] == "O'Brien's Document"
+
+    # Test 7: Unicode characters with injection attempt
+    malicious_value = "public' OR 'א'='א"
+    results = store.search(("docs",), filter={"access": malicious_value})
+    assert len(results) == 0, "Unicode-based injection should be blocked"
+
+
+def test_numeric_filter_safety(store: SqliteStore) -> None:
+    """Test that numeric filter values are handled safely."""
+    # Setup: Create documents with numeric fields
+    store.put(("items",), "item1", {"price": 10, "quantity": 5})
+    store.put(("items",), "item2", {"price": 20, "quantity": 3})
+    store.put(("items",), "item3", {"price": 30, "quantity": 1})
+
+    # Test 1: Normal numeric comparison
+    results = store.search(("items",), filter={"price": {"$gt": 15}})
+    assert len(results) == 2
+    assert all(r.value["price"] > 15 for r in results)
+
+    # Test 2: Special float values (infinity)
+    results = store.search(("items",), filter={"price": {"$lt": float("inf")}})
+    assert len(results) == 3, "All finite values should be less than infinity"
+
+    # Test 3: Special float values (negative infinity)
+    results = store.search(("items",), filter={"price": {"$gt": float("-inf")}})
+    assert len(results) == 3, (
+        "All finite values should be greater than negative infinity"
+    )
+
+    # Test 4: NaN handling - NaN comparisons should not cause errors
+    try:
+        results = store.search(("items",), filter={"price": {"$eq": float("nan")}})
+        # NaN never equals anything, including itself, so should return 0 results
+        assert len(results) == 0
+    except Exception as e:
+        pytest.fail(f"NaN handling should not raise exception: {e}")
+
+    # Test 5: Very large numbers
+    results = store.search(("items",), filter={"price": {"$lt": 10**100}})
+    assert len(results) == 3, "Very large numbers should be handled safely"
+
+    # Test 6: Negative numbers
+    store.put(("items",), "item4", {"price": -10, "quantity": 0})
+    results = store.search(("items",), filter={"price": {"$lt": 0}})
+    assert len(results) == 1
+    assert results[0].key == "item4"
+
+
+def test_boolean_filter_safety(store: SqliteStore) -> None:
+    """Test that boolean filter values are handled safely."""
+    store.put(("flags",), "flag1", {"active": True, "name": "Feature A"})
+    store.put(("flags",), "flag2", {"active": False, "name": "Feature B"})
+    store.put(("flags",), "flag3", {"active": True, "name": "Feature C"})
+
+    # Test boolean filters
+    results = store.search(("flags",), filter={"active": True})
+    assert len(results) == 2
+    assert all(r.value["active"] is True for r in results)
+
+    results = store.search(("flags",), filter={"active": False})
+    assert len(results) == 1
+    assert results[0].value["active"] is False
+
+
+def test_filter_keys_with_hyphens_and_digits(store: SqliteStore) -> None:
+    """Keys with hyphens or leading digits should be queryable via filters.
+
+    Current unquoted JSON path construction (e.g., '$.access-level' or '$.123abc')
+    is not valid JSON1 syntax, so this test will catch regressions in path handling.
+    """
+    # Documents with top-level and nested keys requiring bracket-quoted JSON paths
+    store.put(
+        ("docs",),
+        "hyphen",
+        {"access-level": "public", "user": {"access-level": "nested"}},
+    )
+    store.put(("docs",), "digit", {"123abc": "ok", "user": {"123abc": "ok2"}})
+
+    # Top-level hyphenated key
+    results = store.search(("docs",), filter={"access-level": "public"})
+    assert [r.key for r in results] == ["hyphen"]
+
+    # Nested hyphenated key via dotted path
+    results = store.search(("docs",), filter={"user.access-level": "nested"})
+    assert [r.key for r in results] == ["hyphen"]
+
+    # Top-level digit-starting key
+    results = store.search(("docs",), filter={"123abc": "ok"})
+    assert [r.key for r in results] == ["digit"]
+
+    # Nested digit-starting key via dotted path
+    results = store.search(("docs",), filter={"user.123abc": "ok2"})
+    assert [r.key for r in results] == ["digit"]
+
+
 @pytest.mark.parametrize("distance_type", VECTOR_TYPES)
 def test_non_ascii(
     fake_embeddings: CharacterEmbeddings,