diff --git a/headroom/ccr/batch_store.py b/headroom/ccr/batch_store.py index 895fdd4ec..4bee79897 100644 --- a/headroom/ccr/batch_store.py +++ b/headroom/ccr/batch_store.py @@ -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) diff --git a/headroom/ccr/tool_injection.py b/headroom/ccr/tool_injection.py index 5a873e6bb..3e340fa26 100644 --- a/headroom/ccr/tool_injection.py +++ b/headroom/ccr/tool_injection.py @@ -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 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 diff --git a/headroom/memory/adapters/sqlite.py b/headroom/memory/adapters/sqlite.py index 073538498..a4cc3185f 100644 --- a/headroom/memory/adapters/sqlite.py +++ b/headroom/memory/adapters/sqlite.py @@ -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) diff --git a/tests/test_ccr_response_handler.py b/tests/test_ccr_response_handler.py index de1d96f8f..5897ae63a 100644 --- a/tests/test_ccr_response_handler.py +++ b/tests/test_ccr_response_handler.py @@ -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", diff --git a/tests/test_ccr_tool_injection.py b/tests/test_ccr_tool_injection.py index 6d3f6a023..7158247f3 100644 --- a/tests/test_ccr_tool_injection.py +++ b/tests/test_ccr_tool_injection.py @@ -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.""" diff --git a/tests/test_security_validations.py b/tests/test_security_validations.py new file mode 100644 index 000000000..01719a758 --- /dev/null +++ b/tests/test_security_validations.py @@ -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)