mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
Fix security vulnerabilities in memory and CCR systems
- Fix race condition in BatchContextStore.stats() by acquiring lock - Add atomic dict snapshot in get_memory_stats() to prevent RuntimeError - Add metadata key validation to prevent JSON path injection in SQLite - Parameterize LIMIT/OFFSET in SQLite queries to prevent SQL injection - Strengthen CCR hash validation to require exactly 24 hex characters - Add comprehensive security validation tests
This commit is contained in:
chopratejas
parent
c90cf39826
commit
8f0754a622
6 changed files with 308 additions and 55 deletions
|
|
@ -221,17 +221,22 @@ class BatchContextStore:
|
|||
|
||||
logger.debug(f"Cleaned up {to_remove} oldest batch contexts")
|
||||
|
||||
def stats(self) -> dict[str, Any]:
|
||||
"""Get store statistics."""
|
||||
return {
|
||||
"total_contexts": len(self._contexts),
|
||||
"max_contexts": self._max_contexts,
|
||||
"ttl_seconds": self._ttl,
|
||||
"providers": self._count_by_provider(),
|
||||
}
|
||||
async def stats(self) -> dict[str, Any]:
|
||||
"""Get store statistics.
|
||||
|
||||
def _count_by_provider(self) -> dict[str, int]:
|
||||
"""Count contexts by provider."""
|
||||
Thread-safe: acquires lock before accessing contexts dict to prevent
|
||||
RuntimeError from concurrent modification during iteration.
|
||||
"""
|
||||
async with self._lock:
|
||||
return {
|
||||
"total_contexts": len(self._contexts),
|
||||
"max_contexts": self._max_contexts,
|
||||
"ttl_seconds": self._ttl,
|
||||
"providers": self._count_by_provider_locked(),
|
||||
}
|
||||
|
||||
def _count_by_provider_locked(self) -> dict[str, int]:
|
||||
"""Count contexts by provider. Must be called with lock held."""
|
||||
counts: dict[str, int] = {}
|
||||
for ctx in self._contexts.values():
|
||||
counts[ctx.provider] = counts.get(ctx.provider, 0) + 1
|
||||
|
|
@ -240,6 +245,9 @@ class BatchContextStore:
|
|||
def get_memory_stats(self) -> ComponentStats:
|
||||
"""Get memory statistics for the MemoryTracker.
|
||||
|
||||
Thread-safe: takes a snapshot of contexts dict to prevent RuntimeError
|
||||
from concurrent modification during iteration. Dict copy is atomic in CPython.
|
||||
|
||||
Returns:
|
||||
ComponentStats with current memory usage.
|
||||
"""
|
||||
|
|
@ -247,10 +255,14 @@ class BatchContextStore:
|
|||
|
||||
from ..memory.tracker import ComponentStats
|
||||
|
||||
# Take atomic snapshot to prevent RuntimeError during iteration
|
||||
# dict.copy() is atomic in CPython due to GIL
|
||||
contexts_snapshot = self._contexts.copy()
|
||||
|
||||
# Calculate size
|
||||
size_bytes = sys.getsizeof(self._contexts)
|
||||
|
||||
for batch_id, ctx in self._contexts.items():
|
||||
for batch_id, ctx in contexts_snapshot.items():
|
||||
size_bytes += len(batch_id)
|
||||
size_bytes += sys.getsizeof(ctx)
|
||||
|
||||
|
|
|
|||
|
|
@ -203,13 +203,17 @@ class CCRToolInjector:
|
|||
# - Generic: any [... compressed ... hash=xxx] pattern
|
||||
_marker_patterns: list[re.Pattern] = field(
|
||||
default_factory=lambda: [
|
||||
# All patterns require exactly 24 hex characters for hash validation
|
||||
# CCR uses SHA256 truncated to 24 hex chars (96 bits) for collision resistance
|
||||
# Requiring exact length prevents hash spoofing attacks with shorter hashes
|
||||
#
|
||||
# Standard format: [N <type> compressed to M. Retrieve more: hash=xxx]
|
||||
# Matches items, lines, matches, or any other type
|
||||
re.compile(r"\[(\d+) \w+ compressed to (\d+)\. Retrieve more: hash=([a-f0-9]+)\]"),
|
||||
re.compile(r"\[(\d+) \w+ compressed to (\d+)\. Retrieve more: hash=([a-f0-9]{24})\]"),
|
||||
# Legacy format without "to M" or "Retrieve more:" (old TextCompressor)
|
||||
re.compile(r"\[(\d+) \w+ compressed\. hash=([a-f0-9]+)\]"),
|
||||
# Generic fallback: any compression marker with hash (8+ chars)
|
||||
re.compile(r"\[.*?compressed.*?hash=([a-f0-9]{8,})\]", re.IGNORECASE),
|
||||
re.compile(r"\[(\d+) \w+ compressed\. hash=([a-f0-9]{24})\]"),
|
||||
# Generic fallback: any compression marker with hash (exactly 24 chars)
|
||||
re.compile(r"\[.*?compressed.*?hash=([a-f0-9]{24})\]", re.IGNORECASE),
|
||||
]
|
||||
)
|
||||
|
||||
|
|
@ -456,4 +460,13 @@ def parse_tool_call(
|
|||
hash_key = input_data.get("hash")
|
||||
query = input_data.get("query")
|
||||
|
||||
# Validate hash format: must be exactly 24 hex characters
|
||||
# This prevents hash spoofing attacks with malformed hashes
|
||||
if hash_key is not None:
|
||||
if not isinstance(hash_key, str) or len(hash_key) != 24:
|
||||
return None, None
|
||||
# Validate hex characters only
|
||||
if not all(c in "0123456789abcdef" for c in hash_key.lower()):
|
||||
return None, None
|
||||
|
||||
return hash_key, query
|
||||
|
|
|
|||
|
|
@ -10,6 +10,7 @@ Provides persistent storage for Memory objects with full support for:
|
|||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import re
|
||||
import sqlite3
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
|
|
@ -20,6 +21,27 @@ import numpy as np
|
|||
from ..models import Memory, ScopeLevel
|
||||
from ..ports import MemoryFilter
|
||||
|
||||
# Regex pattern for safe metadata keys: alphanumeric, underscores, hyphens only
|
||||
# This prevents JSON path injection attacks via malicious key names
|
||||
_SAFE_METADATA_KEY_PATTERN = re.compile(r"^[a-zA-Z_][a-zA-Z0-9_\-]*$")
|
||||
|
||||
|
||||
def _validate_metadata_key(key: str) -> bool:
|
||||
"""Validate that a metadata key is safe for use in JSON path expressions.
|
||||
|
||||
Prevents JSON path injection by ensuring keys contain only safe characters.
|
||||
Valid keys: start with letter or underscore, contain only alphanumeric, underscore, hyphen.
|
||||
|
||||
Args:
|
||||
key: The metadata key to validate.
|
||||
|
||||
Returns:
|
||||
True if the key is safe, False otherwise.
|
||||
"""
|
||||
if not key or len(key) > 255:
|
||||
return False
|
||||
return _SAFE_METADATA_KEY_PATTERN.match(key) is not None
|
||||
|
||||
|
||||
class SQLiteMemoryStore:
|
||||
"""SQLite-based memory store implementing the MemoryStore protocol.
|
||||
|
|
@ -481,9 +503,14 @@ class SQLiteMemoryStore:
|
|||
else:
|
||||
conditions.append("promoted_from IS NULL")
|
||||
|
||||
# Metadata filtering
|
||||
# Metadata filtering with key validation to prevent JSON path injection
|
||||
if filter.metadata_filters:
|
||||
for key, value in filter.metadata_filters.items():
|
||||
# Validate key to prevent JSON path injection attacks
|
||||
# Invalid keys are silently skipped to avoid breaking legitimate queries
|
||||
# while blocking malicious attempts like "'] OR 1=1--"
|
||||
if not _validate_metadata_key(key):
|
||||
continue
|
||||
# Use JSON extraction for metadata filtering
|
||||
conditions.append(f"json_extract(metadata, '$.{key}') = ?")
|
||||
params.append(json.dumps(value) if not isinstance(value, str) else value)
|
||||
|
|
@ -523,12 +550,14 @@ class SQLiteMemoryStore:
|
|||
ORDER BY {order_column} {order_direction}
|
||||
"""
|
||||
|
||||
# Add pagination
|
||||
# Add pagination using parameterized queries to prevent injection
|
||||
if filter.limit is not None:
|
||||
query += f" LIMIT {filter.limit}"
|
||||
query += " LIMIT ?"
|
||||
params.append(filter.limit)
|
||||
|
||||
if filter.offset > 0:
|
||||
query += f" OFFSET {filter.offset}"
|
||||
query += " OFFSET ?"
|
||||
params.append(filter.offset)
|
||||
|
||||
with self._get_conn() as conn:
|
||||
cursor = conn.execute(query, params)
|
||||
|
|
|
|||
|
|
@ -155,7 +155,7 @@ class TestCCRToolCallParsing:
|
|||
"type": "tool_use",
|
||||
"id": "tool_123",
|
||||
"name": CCR_TOOL_NAME,
|
||||
"input": {"hash": "abc123"},
|
||||
"input": {"hash": "abc123def456abc123def456"},
|
||||
}
|
||||
]
|
||||
}
|
||||
|
|
@ -164,7 +164,7 @@ class TestCCRToolCallParsing:
|
|||
|
||||
assert len(ccr_calls) == 1
|
||||
assert ccr_calls[0].tool_call_id == "tool_123"
|
||||
assert ccr_calls[0].hash_key == "abc123"
|
||||
assert ccr_calls[0].hash_key == "abc123def456abc123def456"
|
||||
assert ccr_calls[0].query is None
|
||||
assert len(other_calls) == 0
|
||||
|
||||
|
|
@ -178,7 +178,7 @@ class TestCCRToolCallParsing:
|
|||
"type": "tool_use",
|
||||
"id": "tool_456",
|
||||
"name": CCR_TOOL_NAME,
|
||||
"input": {"hash": "def456", "query": "authentication error"},
|
||||
"input": {"hash": "def456abc123def456abc123", "query": "authentication error"},
|
||||
}
|
||||
]
|
||||
}
|
||||
|
|
@ -186,7 +186,7 @@ class TestCCRToolCallParsing:
|
|||
ccr_calls, other_calls = handler._parse_ccr_tool_calls(response, "anthropic")
|
||||
|
||||
assert len(ccr_calls) == 1
|
||||
assert ccr_calls[0].hash_key == "def456"
|
||||
assert ccr_calls[0].hash_key == "def456abc123def456abc123"
|
||||
assert ccr_calls[0].query == "authentication error"
|
||||
|
||||
def test_parse_mixed_tool_calls(self):
|
||||
|
|
@ -199,7 +199,7 @@ class TestCCRToolCallParsing:
|
|||
"type": "tool_use",
|
||||
"id": "tool_1",
|
||||
"name": CCR_TOOL_NAME,
|
||||
"input": {"hash": "abc123"},
|
||||
"input": {"hash": "abc123def456abc123def456"},
|
||||
},
|
||||
{
|
||||
"type": "tool_use",
|
||||
|
|
|
|||
|
|
@ -54,7 +54,7 @@ class TestCCRToolInjector:
|
|||
{"role": "user", "content": "Find errors"},
|
||||
{
|
||||
"role": "tool",
|
||||
"content": '[{"id": 1}]\n[100 items compressed to 10. Retrieve more: hash=abc123def456]',
|
||||
"content": '[{"id": 1}]\n[100 items compressed to 10. Retrieve more: hash=abc123def456abc123def456]',
|
||||
},
|
||||
]
|
||||
|
||||
|
|
@ -62,7 +62,7 @@ class TestCCRToolInjector:
|
|||
hashes = injector.scan_for_markers(messages)
|
||||
|
||||
assert len(hashes) == 1
|
||||
assert "abc123def456" in hashes
|
||||
assert "abc123def456abc123def456" in hashes
|
||||
assert injector.has_compressed_content
|
||||
|
||||
def test_scan_for_markers_multiple_hashes(self):
|
||||
|
|
@ -70,11 +70,11 @@ class TestCCRToolInjector:
|
|||
messages = [
|
||||
{
|
||||
"role": "tool",
|
||||
"content": "[50 items compressed to 5. Retrieve more: hash=aaa111111111]",
|
||||
"content": "[50 items compressed to 5. Retrieve more: hash=aaa111111111aaa111111111]",
|
||||
},
|
||||
{
|
||||
"role": "tool",
|
||||
"content": "[200 items compressed to 20. Retrieve more: hash=bbb222222222]",
|
||||
"content": "[200 items compressed to 20. Retrieve more: hash=bbb222222222bbb222222222]",
|
||||
},
|
||||
]
|
||||
|
||||
|
|
@ -82,19 +82,19 @@ class TestCCRToolInjector:
|
|||
hashes = injector.scan_for_markers(messages)
|
||||
|
||||
assert len(hashes) == 2
|
||||
assert "aaa111111111" in hashes
|
||||
assert "bbb222222222" in hashes
|
||||
assert "aaa111111111aaa111111111" in hashes
|
||||
assert "bbb222222222bbb222222222" in hashes
|
||||
|
||||
def test_scan_no_duplicates(self):
|
||||
"""Scanner deduplicates repeated hashes."""
|
||||
messages = [
|
||||
{
|
||||
"role": "tool",
|
||||
"content": "[100 items compressed to 10. Retrieve more: hash=aabbcc123456]",
|
||||
"content": "[100 items compressed to 10. Retrieve more: hash=aabbcc123456aabbcc123456]",
|
||||
},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": "I see [100 items compressed to 10. Retrieve more: hash=aabbcc123456]",
|
||||
"content": "I see [100 items compressed to 10. Retrieve more: hash=aabbcc123456aabbcc123456]",
|
||||
},
|
||||
]
|
||||
|
||||
|
|
@ -117,7 +117,7 @@ class TestCCRToolInjector:
|
|||
"content": [
|
||||
{
|
||||
"type": "tool_result",
|
||||
"content": "[100 items compressed to 10. Retrieve more: hash=b10cf0a2b3c4]",
|
||||
"content": "[100 items compressed to 10. Retrieve more: hash=b10cf0a2b3c4b10cf0a2b3c4]",
|
||||
},
|
||||
],
|
||||
},
|
||||
|
|
@ -126,14 +126,14 @@ class TestCCRToolInjector:
|
|||
injector = CCRToolInjector()
|
||||
hashes = injector.scan_for_markers(messages)
|
||||
|
||||
assert "b10cf0a2b3c4" in hashes
|
||||
assert "b10cf0a2b3c4b10cf0a2b3c4" in hashes
|
||||
|
||||
def test_inject_tool_when_compression_detected(self):
|
||||
"""Tool is injected when compression markers are found."""
|
||||
messages = [
|
||||
{
|
||||
"role": "tool",
|
||||
"content": "[100 items compressed to 10. Retrieve more: hash=abc123def456]",
|
||||
"content": "[100 items compressed to 10. Retrieve more: hash=abc123def456abc123def456]",
|
||||
},
|
||||
]
|
||||
|
||||
|
|
@ -150,7 +150,7 @@ class TestCCRToolInjector:
|
|||
messages = [
|
||||
{
|
||||
"role": "tool",
|
||||
"content": "[100 items compressed to 10. Retrieve more: hash=e1e2e3f4f5f6]",
|
||||
"content": "[100 items compressed to 10. Retrieve more: hash=e1e2e3f4f5f6e1e2e3f4f5f6]",
|
||||
},
|
||||
]
|
||||
existing_tools = [{"name": "other_tool", "input_schema": {}}]
|
||||
|
|
@ -169,7 +169,7 @@ class TestCCRToolInjector:
|
|||
messages = [
|
||||
{
|
||||
"role": "tool",
|
||||
"content": "[100 items compressed to 10. Retrieve more: hash=aac123456789]",
|
||||
"content": "[100 items compressed to 10. Retrieve more: hash=aac123456789aac123456789]",
|
||||
},
|
||||
]
|
||||
# Tool already present (e.g., from MCP)
|
||||
|
|
@ -187,7 +187,7 @@ class TestCCRToolInjector:
|
|||
messages = [
|
||||
{
|
||||
"role": "tool",
|
||||
"content": "[100 items compressed to 10. Retrieve more: hash=bbc456789012]",
|
||||
"content": "[100 items compressed to 10. Retrieve more: hash=bbc456789012bbc456789012]",
|
||||
},
|
||||
]
|
||||
# OpenAI format tool already present
|
||||
|
|
@ -222,7 +222,7 @@ class TestCCRToolInjector:
|
|||
{"role": "system", "content": "You are helpful."},
|
||||
{
|
||||
"role": "tool",
|
||||
"content": "[100 items compressed to 10. Retrieve more: hash=abc123def456]",
|
||||
"content": "[100 items compressed to 10. Retrieve more: hash=abc123def456abc123def456]",
|
||||
},
|
||||
]
|
||||
|
||||
|
|
@ -231,7 +231,7 @@ class TestCCRToolInjector:
|
|||
updated = injector.inject_into_system_message(messages)
|
||||
|
||||
assert "Compressed Context Available" in updated[0]["content"]
|
||||
assert "abc123def456" in updated[0]["content"]
|
||||
assert "abc123def456abc123def456" in updated[0]["content"]
|
||||
|
||||
def test_process_request_full_flow(self):
|
||||
"""process_request handles complete injection flow."""
|
||||
|
|
@ -240,7 +240,7 @@ class TestCCRToolInjector:
|
|||
{"role": "user", "content": "Search for errors"},
|
||||
{
|
||||
"role": "tool",
|
||||
"content": "[500 items compressed to 25. Retrieve more: hash=f011f10abcde]",
|
||||
"content": "[500 items compressed to 25. Retrieve more: hash=f011f10abcdef011f10abcde]",
|
||||
},
|
||||
]
|
||||
|
||||
|
|
@ -266,12 +266,12 @@ class TestParseToolCall:
|
|||
tool_call = {
|
||||
"id": "toolu_123",
|
||||
"name": CCR_TOOL_NAME,
|
||||
"input": {"hash": "abc123", "query": "errors"},
|
||||
"input": {"hash": "abc123def456abc123def456", "query": "errors"},
|
||||
}
|
||||
|
||||
hash_key, query = parse_tool_call(tool_call, "anthropic")
|
||||
|
||||
assert hash_key == "abc123"
|
||||
assert hash_key == "abc123def456abc123def456"
|
||||
assert query == "errors"
|
||||
|
||||
def test_parse_openai_format(self):
|
||||
|
|
@ -280,13 +280,13 @@ class TestParseToolCall:
|
|||
"id": "call_123",
|
||||
"function": {
|
||||
"name": CCR_TOOL_NAME,
|
||||
"arguments": json.dumps({"hash": "def456", "query": None}),
|
||||
"arguments": json.dumps({"hash": "def456abc123def456abc123", "query": None}),
|
||||
},
|
||||
}
|
||||
|
||||
hash_key, query = parse_tool_call(tool_call, "openai")
|
||||
|
||||
assert hash_key == "def456"
|
||||
assert hash_key == "def456abc123def456abc123"
|
||||
assert query is None
|
||||
|
||||
def test_parse_non_ccr_tool(self):
|
||||
|
|
@ -316,6 +316,65 @@ class TestParseToolCall:
|
|||
assert hash_key is None
|
||||
|
||||
|
||||
class TestHashSecurityValidation:
|
||||
"""Test hash validation security measures.
|
||||
|
||||
CCR hashes must be exactly 24 hex characters (96 bits of SHA256).
|
||||
This prevents hash spoofing attacks with shorter or malformed hashes.
|
||||
"""
|
||||
|
||||
def test_rejects_short_hash(self):
|
||||
"""Rejects hash that's too short (potential spoofing attack)."""
|
||||
tool_call = {
|
||||
"name": CCR_TOOL_NAME,
|
||||
"input": {"hash": "abc123"}, # Only 6 chars
|
||||
}
|
||||
|
||||
hash_key, query = parse_tool_call(tool_call, "anthropic")
|
||||
assert hash_key is None # Rejected
|
||||
|
||||
def test_rejects_long_hash(self):
|
||||
"""Rejects hash that's too long."""
|
||||
tool_call = {
|
||||
"name": CCR_TOOL_NAME,
|
||||
"input": {"hash": "abc123def456abc123def456abc123"}, # 30 chars
|
||||
}
|
||||
|
||||
hash_key, query = parse_tool_call(tool_call, "anthropic")
|
||||
assert hash_key is None # Rejected
|
||||
|
||||
def test_rejects_non_hex_characters(self):
|
||||
"""Rejects hash with non-hex characters."""
|
||||
tool_call = {
|
||||
"name": CCR_TOOL_NAME,
|
||||
"input": {"hash": "abc123xyz456abc123xyz456"}, # Contains xyz
|
||||
}
|
||||
|
||||
hash_key, query = parse_tool_call(tool_call, "anthropic")
|
||||
assert hash_key is None # Rejected
|
||||
|
||||
def test_accepts_valid_24_char_hash(self):
|
||||
"""Accepts properly formatted 24-char hex hash."""
|
||||
tool_call = {
|
||||
"name": CCR_TOOL_NAME,
|
||||
"input": {"hash": "abc123def456abc123def456"},
|
||||
}
|
||||
|
||||
hash_key, query = parse_tool_call(tool_call, "anthropic")
|
||||
assert hash_key == "abc123def456abc123def456"
|
||||
|
||||
def test_accepts_uppercase_hex(self):
|
||||
"""Accepts uppercase hex characters (normalized to lowercase internally)."""
|
||||
tool_call = {
|
||||
"name": CCR_TOOL_NAME,
|
||||
"input": {"hash": "ABC123DEF456ABC123DEF456"},
|
||||
}
|
||||
|
||||
hash_key, query = parse_tool_call(tool_call, "anthropic")
|
||||
# Note: validation accepts uppercase since we use .lower() for hex check
|
||||
assert hash_key == "ABC123DEF456ABC123DEF456"
|
||||
|
||||
|
||||
class TestSystemInstructions:
|
||||
"""Test system instruction generation."""
|
||||
|
||||
|
|
@ -364,7 +423,7 @@ class TestAlternativeMarkerFormats:
|
|||
messages = [
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": "Build output:\n[500 lines compressed to 50. Retrieve more: hash=aabbccddeeff00112233]",
|
||||
"content": "Build output:\n[500 lines compressed to 50. Retrieve more: hash=aabbccddeeff001122334455]",
|
||||
},
|
||||
]
|
||||
|
||||
|
|
@ -372,14 +431,14 @@ class TestAlternativeMarkerFormats:
|
|||
hashes = injector.scan_for_markers(messages)
|
||||
|
||||
assert len(hashes) == 1
|
||||
assert "aabbccddeeff00112233" in hashes
|
||||
assert "aabbccddeeff001122334455" in hashes
|
||||
|
||||
def test_searchcompressor_format(self):
|
||||
"""Detects SearchCompressor marker format (matches)."""
|
||||
messages = [
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": "Search results:\n[100 matches compressed to 10. Retrieve more: hash=1122334455667788]",
|
||||
"content": "Search results:\n[100 matches compressed to 10. Retrieve more: hash=112233445566778899001122]",
|
||||
},
|
||||
]
|
||||
|
||||
|
|
@ -387,22 +446,22 @@ class TestAlternativeMarkerFormats:
|
|||
hashes = injector.scan_for_markers(messages)
|
||||
|
||||
assert len(hashes) == 1
|
||||
assert "1122334455667788" in hashes
|
||||
assert "112233445566778899001122" in hashes
|
||||
|
||||
def test_mixed_compressor_formats(self):
|
||||
"""Detects multiple marker formats in same conversation."""
|
||||
messages = [
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": "Search results:\n[50 matches compressed to 5. Retrieve more: hash=aaaa11111111]",
|
||||
"content": "Search results:\n[50 matches compressed to 5. Retrieve more: hash=aaaa11111111aaaa11111111]",
|
||||
},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": "Build logs:\n[200 lines compressed to 20. Retrieve more: hash=bbbb22222222]",
|
||||
"content": "Build logs:\n[200 lines compressed to 20. Retrieve more: hash=bbbb22222222bbbb22222222]",
|
||||
},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": "Database:\n[1000 items compressed to 100. Retrieve more: hash=cccc33333333]",
|
||||
"content": "Database:\n[1000 items compressed to 100. Retrieve more: hash=cccc33333333cccc33333333]",
|
||||
},
|
||||
]
|
||||
|
||||
|
|
@ -410,9 +469,9 @@ class TestAlternativeMarkerFormats:
|
|||
hashes = injector.scan_for_markers(messages)
|
||||
|
||||
assert len(hashes) == 3
|
||||
assert "aaaa11111111" in hashes
|
||||
assert "bbbb22222222" in hashes
|
||||
assert "cccc33333333" in hashes
|
||||
assert "aaaa11111111aaaa11111111" in hashes
|
||||
assert "bbbb22222222bbbb22222222" in hashes
|
||||
assert "cccc33333333cccc33333333" in hashes
|
||||
|
||||
def test_generic_compressed_marker(self):
|
||||
"""Detects generic compression markers via fallback pattern."""
|
||||
|
|
|
|||
140
tests/test_security_validations.py
Normal file
140
tests/test_security_validations.py
Normal file
|
|
@ -0,0 +1,140 @@
|
|||
"""Security validation tests for Headroom.
|
||||
|
||||
These tests verify security measures against common attack vectors:
|
||||
- SQL injection via metadata keys
|
||||
- CCR hash spoofing attacks
|
||||
- JSON path injection
|
||||
|
||||
These tests exist as regression tests to ensure security fixes remain effective.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
|
||||
from headroom.memory.adapters.sqlite import _validate_metadata_key
|
||||
|
||||
|
||||
class TestSQLiteMetadataKeyValidation:
|
||||
"""Test metadata key validation to prevent JSON path injection.
|
||||
|
||||
Metadata keys are interpolated into json_extract() SQL expressions.
|
||||
Without validation, malicious keys could escape the JSON path and
|
||||
inject arbitrary SQL.
|
||||
"""
|
||||
|
||||
def test_valid_simple_key(self):
|
||||
"""Accept simple alphanumeric keys."""
|
||||
assert _validate_metadata_key("status") is True
|
||||
assert _validate_metadata_key("user_id") is True
|
||||
assert _validate_metadata_key("item_count") is True
|
||||
|
||||
def test_valid_key_with_hyphens(self):
|
||||
"""Accept keys with hyphens (common in APIs)."""
|
||||
assert _validate_metadata_key("content-type") is True
|
||||
assert _validate_metadata_key("x-custom-header") is True
|
||||
|
||||
def test_valid_key_starting_with_underscore(self):
|
||||
"""Accept keys starting with underscore."""
|
||||
assert _validate_metadata_key("_internal") is True
|
||||
assert _validate_metadata_key("_id") is True
|
||||
|
||||
def test_rejects_empty_key(self):
|
||||
"""Reject empty keys."""
|
||||
assert _validate_metadata_key("") is False
|
||||
|
||||
def test_rejects_json_path_injection(self):
|
||||
"""Reject keys that could escape JSON path expression."""
|
||||
# Attempt to close the JSON path and add SQL
|
||||
assert _validate_metadata_key("'] OR 1=1--") is False
|
||||
assert _validate_metadata_key("key') OR 1=1--") is False
|
||||
assert _validate_metadata_key('key") OR 1=1--') is False
|
||||
|
||||
def test_rejects_sql_injection_patterns(self):
|
||||
"""Reject keys with SQL injection patterns."""
|
||||
assert _validate_metadata_key("key; DROP TABLE memories;--") is False
|
||||
assert _validate_metadata_key("key UNION SELECT * FROM users") is False
|
||||
assert _validate_metadata_key("1=1") is False
|
||||
|
||||
def test_rejects_special_characters(self):
|
||||
"""Reject keys with special characters."""
|
||||
assert _validate_metadata_key("key.nested") is False # Dots
|
||||
assert _validate_metadata_key("key[0]") is False # Brackets
|
||||
assert _validate_metadata_key("key$") is False # Dollar sign
|
||||
assert _validate_metadata_key("key@domain") is False # At sign
|
||||
assert _validate_metadata_key("key/path") is False # Slashes
|
||||
assert _validate_metadata_key("key\\path") is False # Backslashes
|
||||
assert _validate_metadata_key("key'quote") is False # Quotes
|
||||
assert _validate_metadata_key('key"quote') is False # Double quotes
|
||||
|
||||
def test_rejects_keys_starting_with_number(self):
|
||||
"""Reject keys starting with a number."""
|
||||
assert _validate_metadata_key("123key") is False
|
||||
assert _validate_metadata_key("0_prefix") is False
|
||||
|
||||
def test_rejects_very_long_keys(self):
|
||||
"""Reject excessively long keys (potential DoS)."""
|
||||
long_key = "a" * 256
|
||||
assert _validate_metadata_key(long_key) is False
|
||||
|
||||
# 255 chars should be acceptable
|
||||
valid_long_key = "a" * 255
|
||||
assert _validate_metadata_key(valid_long_key) is True
|
||||
|
||||
def test_rejects_unicode_bypass_attempts(self):
|
||||
"""Reject Unicode characters that might bypass filtering."""
|
||||
# Various Unicode quote-like characters
|
||||
assert _validate_metadata_key("key\u2019") is False # Right single quote
|
||||
assert _validate_metadata_key("key\u201c") is False # Left double quote
|
||||
assert _validate_metadata_key("key\u0000") is False # Null byte
|
||||
|
||||
|
||||
class TestSQLiteMetadataFilteringIntegration:
|
||||
"""Integration tests for metadata filtering with validation."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_malicious_metadata_filter_is_skipped(self):
|
||||
"""Malicious metadata keys should be silently skipped, not cause errors."""
|
||||
import tempfile
|
||||
from pathlib import Path
|
||||
|
||||
from headroom.memory.adapters.sqlite import SQLiteMemoryStore
|
||||
from headroom.memory.models import Memory
|
||||
from headroom.memory.ports import MemoryFilter
|
||||
|
||||
with tempfile.NamedTemporaryFile(suffix=".db", delete=False) as f:
|
||||
db_path = Path(f.name)
|
||||
|
||||
try:
|
||||
store = SQLiteMemoryStore(db_path)
|
||||
|
||||
# Create a memory with safe metadata
|
||||
memory = Memory(
|
||||
id="test-1",
|
||||
content="Test content",
|
||||
user_id="alice",
|
||||
metadata={"safe_key": "value"},
|
||||
)
|
||||
await store.save(memory)
|
||||
|
||||
# Attempt to query with malicious metadata key - should not raise
|
||||
# The malicious key should be silently skipped
|
||||
malicious_filter = MemoryFilter(
|
||||
user_id="alice",
|
||||
metadata_filters={"'] OR 1=1--": "malicious"},
|
||||
)
|
||||
results = await store.query(malicious_filter)
|
||||
|
||||
# Query should succeed (malicious key skipped)
|
||||
# Results may or may not include our memory depending on other conditions
|
||||
assert isinstance(results, list)
|
||||
|
||||
# Query with valid metadata filter should work normally
|
||||
valid_filter = MemoryFilter(
|
||||
user_id="alice",
|
||||
metadata_filters={"safe_key": "value"},
|
||||
)
|
||||
results = await store.query(valid_filter)
|
||||
assert len(results) == 1
|
||||
assert results[0].id == "test-1"
|
||||
|
||||
finally:
|
||||
db_path.unlink(missing_ok=True)
|
||||
Loading…
Add table
Add a link
Reference in a new issue