This commit is contained in:
Eugene Yurtsev
2025-12-09 11:28:15 -05:00
parent 706f3e981e
commit b7d068677c
2 changed files with 156 additions and 92 deletions
@@ -3,6 +3,7 @@ from __future__ import annotations
import concurrent.futures
import datetime
import logging
import math
import re
import sqlite3
import threading
@@ -404,12 +405,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 +421,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 +636,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 +855,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:
+102
View File
@@ -1069,6 +1069,108 @@ 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,