Compare commits

...
Author SHA1 Message Date
Eugene Yurtsev 18a2a2a832 x 2025-12-09 16:24:45 -05:00
Eugene Yurtsev 5ae0aff522 x 2025-12-09 16:24:38 -05:00
Eugene Yurtsev b7d068677c x 2025-12-09 11:28:15 -05:00
Eugene Yurtsev 706f3e981e x 2025-12-08 22:35:50 -05:00
Eugene Yurtsev 344ab65351 checkpoint sqlite 2025-12-08 21:24:00 -05:00
7 changed files with 374 additions and 97 deletions
@@ -329,8 +329,9 @@ class SqliteSaver(BaseCheckpointSaver[str]):
FROM checkpoints
{where}
ORDER BY checkpoint_id DESC"""
if limit:
query += f" LIMIT {limit}"
if limit is not None:
query += " LIMIT ?"
param_values = (*param_values, limit)
with self.cursor(transaction=False) as cur, closing(self.conn.cursor()) as wcur:
cur.execute(query, param_values)
for (
@@ -425,8 +425,9 @@ class AsyncSqliteSaver(BaseCheckpointSaver[str]):
FROM checkpoints
{where}
ORDER BY checkpoint_id DESC"""
if limit:
query += f" LIMIT {limit}"
if limit is not None:
query += " LIMIT ?"
params = (*params, limit)
async with (
self.lock,
self.conn.execute(query, params) as cur,
@@ -1,6 +1,7 @@
from __future__ import annotations
import json
import re
from collections.abc import Sequence
from typing import Any
@@ -8,6 +9,23 @@ from langchain_core.runnables import RunnableConfig
from langgraph.checkpoint.base import get_checkpoint_id
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 _metadata_predicate(
metadata_filter: dict[str, Any],
) -> tuple[Sequence[str], Sequence[Any]]:
@@ -43,6 +61,7 @@ def _metadata_predicate(
# process metadata query
for query_key, query_value in metadata_filter.items():
_validate_filter_key(query_key)
operator, param_value = _where_value(query_value)
predicates.append(
f"json_extract(CAST(metadata AS TEXT), '$.{query_key}') {operator}"
@@ -404,12 +404,9 @@ class BaseSqliteStore:
# SQLite json_extract returns unquoted string values
if isinstance(value, str):
filter_conditions.append(
"json_extract(value, '$."
+ key
+ "') = '"
+ value.replace("'", "''")
+ "'"
"json_extract(value, '$." + key + "') = ?"
)
filter_params.append(value)
elif value is None:
filter_conditions.append(
"json_extract(value, '$." + key + "') IS NULL"
@@ -423,9 +420,11 @@ class BaseSqliteStore:
+ ("1" if value else "0")
)
elif isinstance(value, (int, float)):
# Use parameterized query to handle special floats and large integers
filter_conditions.append(
"json_extract(value, '$." + key + "') = " + str(value)
"json_extract(value, '$." + key + "') = ?"
)
filter_params.append(float(value))
else:
# Complex objects (list, dict, …) compare JSON text
filter_conditions.append(
@@ -636,85 +635,66 @@ class BaseSqliteStore:
# We need to properly format values for SQLite JSON extraction comparison
if op == "$eq":
if isinstance(value, str):
# Direct string comparison with proper quoting for unquoted json_extract result
return (
f"json_extract(value, '$.{key}') = '"
+ value.replace("'", "''")
+ "'",
[],
)
return f"json_extract(value, '$.{key}') = ?", [value]
elif value is None:
return f"json_extract(value, '$.{key}') IS NULL", []
elif isinstance(value, bool):
# SQLite JSON stores booleans as integers
return f"json_extract(value, '$.{key}') = {1 if value else 0}", []
elif isinstance(value, (int, float)):
return f"json_extract(value, '$.{key}') = {value}", []
# Convert to float to handle inf, -inf, nan, and very large integers
# SQLite REAL can handle these cases better than INTEGER
return f"json_extract(value, '$.{key}') = ?", [float(value)]
else:
return f"json_extract(value, '$.{key}') = ?", [orjson.dumps(value)]
elif op == "$gt":
# For numeric values, SQLite needs to compare as numbers, not strings
if isinstance(value, (int, float)):
return f"CAST(json_extract(value, '$.{key}') AS REAL) > {value}", []
# Convert to float to handle special values and very large integers
return f"CAST(json_extract(value, '$.{key}') AS REAL) > ?", [
float(value)
]
elif isinstance(value, str):
return (
f"json_extract(value, '$.{key}') > '"
+ value.replace("'", "''")
+ "'",
[],
)
return f"json_extract(value, '$.{key}') > ?", [value]
else:
return f"json_extract(value, '$.{key}') > ?", [orjson.dumps(value)]
elif op == "$gte":
if isinstance(value, (int, float)):
return f"CAST(json_extract(value, '$.{key}') AS REAL) >= {value}", []
return f"CAST(json_extract(value, '$.{key}') AS REAL) >= ?", [
float(value)
]
elif isinstance(value, str):
return (
f"json_extract(value, '$.{key}') >= '"
+ value.replace("'", "''")
+ "'",
[],
)
return f"json_extract(value, '$.{key}') >= ?", [value]
else:
return f"json_extract(value, '$.{key}') >= ?", [orjson.dumps(value)]
elif op == "$lt":
if isinstance(value, (int, float)):
return f"CAST(json_extract(value, '$.{key}') AS REAL) < {value}", []
return f"CAST(json_extract(value, '$.{key}') AS REAL) < ?", [
float(value)
]
elif isinstance(value, str):
return (
f"json_extract(value, '$.{key}') < '"
+ value.replace("'", "''")
+ "'",
[],
)
return f"json_extract(value, '$.{key}') < ?", [value]
else:
return f"json_extract(value, '$.{key}') < ?", [orjson.dumps(value)]
elif op == "$lte":
if isinstance(value, (int, float)):
return f"CAST(json_extract(value, '$.{key}') AS REAL) <= {value}", []
return f"CAST(json_extract(value, '$.{key}') AS REAL) <= ?", [
float(value)
]
elif isinstance(value, str):
return (
f"json_extract(value, '$.{key}') <= '"
+ value.replace("'", "''")
+ "'",
[],
)
return f"json_extract(value, '$.{key}') <= ?", [value]
else:
return f"json_extract(value, '$.{key}') <= ?", [orjson.dumps(value)]
elif op == "$ne":
if isinstance(value, str):
return (
f"json_extract(value, '$.{key}') != '"
+ value.replace("'", "''")
+ "'",
[],
)
return f"json_extract(value, '$.{key}') != ?", [value]
elif value is None:
return f"json_extract(value, '$.{key}') IS NOT NULL", []
elif isinstance(value, bool):
return f"json_extract(value, '$.{key}') != {1 if value else 0}", []
elif isinstance(value, (int, float)):
return f"json_extract(value, '$.{key}') != {value}", []
# Convert to float for consistency
return f"json_extract(value, '$.{key}') != ?", [float(value)]
else:
return f"json_extract(value, '$.{key}') != ?", [orjson.dumps(value)]
else:
@@ -874,85 +854,66 @@ class SqliteStore(BaseSqliteStore, BaseStore):
# We need to properly format values for SQLite JSON extraction comparison
if op == "$eq":
if isinstance(value, str):
# Direct string comparison with proper quoting for unquoted json_extract result
return (
f"json_extract(value, '$.{key}') = '"
+ value.replace("'", "''")
+ "'",
[],
)
return f"json_extract(value, '$.{key}') = ?", [value]
elif value is None:
return f"json_extract(value, '$.{key}') IS NULL", []
elif isinstance(value, bool):
# SQLite JSON stores booleans as integers
return f"json_extract(value, '$.{key}') = {1 if value else 0}", []
elif isinstance(value, (int, float)):
return f"json_extract(value, '$.{key}') = {value}", []
# Convert to float to handle inf, -inf, nan, and very large integers
# SQLite REAL can handle these cases better than INTEGER
return f"json_extract(value, '$.{key}') = ?", [float(value)]
else:
return f"json_extract(value, '$.{key}') = ?", [orjson.dumps(value)]
elif op == "$gt":
# For numeric values, SQLite needs to compare as numbers, not strings
if isinstance(value, (int, float)):
return f"CAST(json_extract(value, '$.{key}') AS REAL) > {value}", []
# Convert to float to handle special values and very large integers
return f"CAST(json_extract(value, '$.{key}') AS REAL) > ?", [
float(value)
]
elif isinstance(value, str):
return (
f"json_extract(value, '$.{key}') > '"
+ value.replace("'", "''")
+ "'",
[],
)
return f"json_extract(value, '$.{key}') > ?", [value]
else:
return f"json_extract(value, '$.{key}') > ?", [orjson.dumps(value)]
elif op == "$gte":
if isinstance(value, (int, float)):
return f"CAST(json_extract(value, '$.{key}') AS REAL) >= {value}", []
return f"CAST(json_extract(value, '$.{key}') AS REAL) >= ?", [
float(value)
]
elif isinstance(value, str):
return (
f"json_extract(value, '$.{key}') >= '"
+ value.replace("'", "''")
+ "'",
[],
)
return f"json_extract(value, '$.{key}') >= ?", [value]
else:
return f"json_extract(value, '$.{key}') >= ?", [orjson.dumps(value)]
elif op == "$lt":
if isinstance(value, (int, float)):
return f"CAST(json_extract(value, '$.{key}') AS REAL) < {value}", []
return f"CAST(json_extract(value, '$.{key}') AS REAL) < ?", [
float(value)
]
elif isinstance(value, str):
return (
f"json_extract(value, '$.{key}') < '"
+ value.replace("'", "''")
+ "'",
[],
)
return f"json_extract(value, '$.{key}') < ?", [value]
else:
return f"json_extract(value, '$.{key}') < ?", [orjson.dumps(value)]
elif op == "$lte":
if isinstance(value, (int, float)):
return f"CAST(json_extract(value, '$.{key}') AS REAL) <= {value}", []
return f"CAST(json_extract(value, '$.{key}') AS REAL) <= ?", [
float(value)
]
elif isinstance(value, str):
return (
f"json_extract(value, '$.{key}') <= '"
+ value.replace("'", "''")
+ "'",
[],
)
return f"json_extract(value, '$.{key}') <= ?", [value]
else:
return f"json_extract(value, '$.{key}') <= ?", [orjson.dumps(value)]
elif op == "$ne":
if isinstance(value, str):
return (
f"json_extract(value, '$.{key}') != '"
+ value.replace("'", "''")
+ "'",
[],
)
return f"json_extract(value, '$.{key}') != ?", [value]
elif value is None:
return f"json_extract(value, '$.{key}') IS NOT NULL", []
elif isinstance(value, bool):
return f"json_extract(value, '$.{key}') != {1 if value else 0}", []
elif isinstance(value, (int, float)):
return f"json_extract(value, '$.{key}') != {value}", []
# Convert to float for consistency
return f"json_extract(value, '$.{key}') != ?", [float(value)]
else:
return f"json_extract(value, '$.{key}') != ?", [orjson.dumps(value)]
else:
+75 -1
View File
@@ -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
+117
View File
@@ -182,3 +182,120 @@ 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
# 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 = list(saver.list(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 = list(saver.list(None, limit=None))
assert len(results) == 5
+104
View File
@@ -1069,6 +1069,110 @@ 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
@pytest.mark.parametrize("distance_type", VECTOR_TYPES)
def test_non_ascii(
fake_embeddings: CharacterEmbeddings,