mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-10 14:27:00 -04:00
## Description Reported by [@JoaoMarcos44](https://github.com/JoaoMarcos44) via an independent security audit — thanks for the careful, well-documented report. Fixes two confirmed findings from a June 2026 security audit of `headroom/compression/` (the `UniversalCompressor` utility). Both are real defects in shipped, public, tested code; note that this module is **not** on the proxy hot path (the proxy uses `headroom/transforms/`), so real-world blast radius is module-local rather than proxy-wide. - **SEC-01 (entropy bypass):** `use_entropy_preservation` was a silent no-op. `compress()` tokenized content at character level (`list(content)`) and fed single-char tokens to `compute_entropy_mask`, whose `min_token_length` guard skipped every one — so high-entropy secrets (API keys, OAuth tokens, UUIDs, hashes) were never preserved despite the feature being enabled. - **SEC-02 (JSON corruption):** the `_simple_compress` truncation fallback (used when Kompress is unavailable or raises) inserted a separator containing raw newlines. When that fallback ran on a span inside a JSON string value it produced invalid JSON (RFC 8259 §7), crashing downstream `json.loads()`. The other three audited items need no code change and were verified, not assumed: SEC-03 (surrogate DoS) is already caught by the `try/except` in `code_handler._extract_mask` and falls back to regex — non-reproducible even with `tree_sitter_language_pack` installed; SEC-04 (prompt injection) is out of a compressor's scope; SEC-05 (SQLite race) is a misread (`CompressionStore` defaults to `InMemoryBackend`; the SQLite backend uses WAL + busy_timeout + a lock). Closes # ## Type of Change - [x] Bug fix (non-breaking change that fixes an issue) ## Changes Made - Add `compute_entropy_mask_for_content()` (`masks.py`): scores whitespace-delimited words and maps high-entropy ones back to character positions, returning a char-aligned mask. The existing token-level `compute_entropy_mask` is left intact. - Introduce `SECRET_ENTROPY_MIN_LENGTH = 20` as the default word-length floor. Normalized Shannon entropy rates short-but-diverse words (e.g. "detailed") nearly as high as a real secret, so a length floor is the discriminator; 20 matches the entropy-detection floor used by secret scanners (trufflehog, detect-secrets) and prevents over-preserving prose (which would otherwise block legitimate compression). - Wire the content-level entropy pass into `UniversalCompressor.compress()` (scores `content`, not the char-level `tokens`). - Replace the `_simple_compress` separator `"\n...[compressed]...\n"` with the control-char-free `" ...[compressed]... "`. - Add regression tests at the mask level and end-to-end. ## Testing - [x] Unit tests pass (`pytest`) - [x] Linting passes (`ruff check .`) - [x] Type checking passes (`mypy headroom`) - [x] New tests added for new functionality - [x] Manual testing performed ### Test Output ```text $ ruff check headroom/compression/ All checks passed! $ mypy headroom/compression/masks.py headroom/compression/universal.py Success: no issues found in 2 source files $ pytest tests/test_compression/test_masks.py tests/test_compression/test_universal.py \ tests/test_compression/test_json_handler.py tests/test_compression/test_code_handler.py -q ======================= 111 passed, 2 warnings in 10.76s ======================= ``` ## Real Behavior Proof - Environment: macOS, Python 3.12 in repo `.venv`; `tree_sitter_language_pack` and Kompress present. - Exact command / steps: reproduced each finding by calling `UniversalCompressor.compress()` directly before/after the fix — SEC-01: `compute_entropy_mask(list("k="+secret))` preserved 0 of N tokens (inert); after fix `compute_entropy_mask_for_content` preserves the secret's char range and the end-to-end test shows a 43-char secret dropped with preservation off / kept with it on. SEC-02: `compress(json.dumps({...long value...}), content_type=JSON)` with `use_kompress=False` raised `JSONDecodeError` before the fix and round-trips through `json.loads()` after. - Observed result: SEC-01 entropy preservation now functions; SEC-02 output is valid JSON on both the Kompress and fallback paths; the previously-failing `test_compression_reduces_tokens` passes again (no over-preservation). - Not tested: `tests/test_compression/test_evals.py` and `test_llm_eval.py` (require external API/model access); the proxy/transforms live path is unaffected since it does not import `UniversalCompressor`. ## Review Readiness - [x] I have performed a self-review - [x] This PR is ready for human review ## Checklist - [x] My code follows the project's style guidelines - [x] I have performed a self-review of my code - [x] I have commented my code, particularly in hard-to-understand areas - [x] I have made corresponding changes to the documentation - [x] My changes generate no new warnings - [x] I have added tests that prove my fix is effective or that my feature works - [x] New and existing unit tests pass locally with my changes - [ ] I have updated the CHANGELOG.md if applicable ## Additional Notes CHANGELOG not updated (handled by the release tooling). The audit also flagged SEC-03/04/05 — left unchanged by design, with verification rationale in the Description.
399 lines
13 KiB
Python
399 lines
13 KiB
Python
"""Tests for UniversalCompressor."""
|
|
|
|
import json
|
|
|
|
import pytest
|
|
|
|
from headroom.compression.detector import ContentType
|
|
from headroom.compression.handlers.base import NoOpHandler
|
|
from headroom.compression.universal import (
|
|
CompressionResult,
|
|
UniversalCompressor,
|
|
UniversalCompressorConfig,
|
|
compress,
|
|
)
|
|
|
|
|
|
class TestUniversalCompressorConfig:
|
|
"""Tests for UniversalCompressorConfig."""
|
|
|
|
def test_default_config(self):
|
|
"""Test default configuration values."""
|
|
config = UniversalCompressorConfig()
|
|
|
|
assert config.use_magika is True
|
|
assert config.use_kompress is True
|
|
assert config.use_entropy_preservation is True
|
|
assert config.entropy_threshold == 0.85
|
|
assert config.min_content_length == 100
|
|
assert config.compression_ratio_target == 0.3
|
|
|
|
def test_custom_config(self):
|
|
"""Test custom configuration."""
|
|
config = UniversalCompressorConfig(
|
|
use_magika=False,
|
|
compression_ratio_target=0.5,
|
|
)
|
|
|
|
assert config.use_magika is False
|
|
assert config.compression_ratio_target == 0.5
|
|
|
|
|
|
class TestCompressionResult:
|
|
"""Tests for CompressionResult."""
|
|
|
|
def test_tokens_saved(self):
|
|
"""Test tokens_saved calculation."""
|
|
result = CompressionResult(
|
|
compressed="short",
|
|
original="much longer original content",
|
|
compression_ratio=0.5,
|
|
tokens_before=100,
|
|
tokens_after=50,
|
|
content_type=ContentType.TEXT,
|
|
detection_confidence=0.9,
|
|
handler_used="test",
|
|
preservation_ratio=0.5,
|
|
)
|
|
|
|
assert result.tokens_saved == 50
|
|
|
|
def test_savings_percentage(self):
|
|
"""Test savings_percentage calculation."""
|
|
result = CompressionResult(
|
|
compressed="short",
|
|
original="longer",
|
|
compression_ratio=0.5,
|
|
tokens_before=100,
|
|
tokens_after=25,
|
|
content_type=ContentType.TEXT,
|
|
detection_confidence=0.9,
|
|
handler_used="test",
|
|
preservation_ratio=0.5,
|
|
)
|
|
|
|
assert result.savings_percentage == 75.0
|
|
|
|
def test_zero_tokens_before(self):
|
|
"""Test handling of zero tokens_before."""
|
|
result = CompressionResult(
|
|
compressed="",
|
|
original="",
|
|
compression_ratio=1.0,
|
|
tokens_before=0,
|
|
tokens_after=0,
|
|
content_type=ContentType.UNKNOWN,
|
|
detection_confidence=0.0,
|
|
handler_used="none",
|
|
preservation_ratio=1.0,
|
|
)
|
|
|
|
assert result.savings_percentage == 0.0
|
|
|
|
|
|
class TestUniversalCompressor:
|
|
"""Tests for UniversalCompressor."""
|
|
|
|
@pytest.fixture
|
|
def compressor(self):
|
|
"""Create compressor with fallback detector (no Magika required)."""
|
|
config = UniversalCompressorConfig(
|
|
use_magika=False, # Use fallback detector
|
|
use_kompress=False, # Use simple compression
|
|
ccr_enabled=False, # Skip CCR
|
|
)
|
|
return UniversalCompressor(config=config)
|
|
|
|
def test_compress_short_content_unchanged(self, compressor):
|
|
"""Test that short content is not compressed."""
|
|
content = "short"
|
|
result = compressor.compress(content)
|
|
|
|
assert result.compressed == content
|
|
assert result.compression_ratio == 1.0
|
|
assert "skipped" in result.metadata
|
|
|
|
def test_compress_empty_content(self, compressor):
|
|
"""Test handling of empty content."""
|
|
result = compressor.compress("")
|
|
|
|
assert result.compressed == ""
|
|
assert result.content_type == ContentType.UNKNOWN
|
|
|
|
def test_compress_json_content(self, compressor):
|
|
"""Test compression of JSON content."""
|
|
content = json.dumps(
|
|
{"users": [{"id": i, "name": f"User {i}", "bio": "x" * 100} for i in range(10)]}
|
|
)
|
|
|
|
result = compressor.compress(content)
|
|
|
|
assert result.content_type == ContentType.JSON
|
|
assert result.handler_used == "json"
|
|
# Compression should reduce size
|
|
assert len(result.compressed) < len(content)
|
|
|
|
def test_compress_code_content(self, compressor):
|
|
"""Test compression of code content."""
|
|
content = (
|
|
'''
|
|
def hello_world():
|
|
"""Say hello to the world."""
|
|
message = "Hello, World!"
|
|
print(message)
|
|
return message
|
|
|
|
def another_function():
|
|
"""Another function with a long body."""
|
|
x = 1
|
|
y = 2
|
|
z = x + y
|
|
'''
|
|
+ "result = z * " * 50
|
|
+ """
|
|
return result
|
|
"""
|
|
)
|
|
result = compressor.compress(content)
|
|
|
|
assert result.content_type == ContentType.CODE
|
|
assert result.handler_used == "code"
|
|
|
|
def test_compress_plain_text(self, compressor):
|
|
"""Test compression of plain text."""
|
|
content = "This is plain text without any special structure. " * 20
|
|
|
|
result = compressor.compress(content)
|
|
|
|
assert result.content_type == ContentType.TEXT
|
|
|
|
def test_compress_with_override_type(self, compressor):
|
|
"""Test compression with overridden content type."""
|
|
content = '{"key": "value"}' + " " * 100 # Pad to meet min length
|
|
|
|
result = compressor.compress(content, content_type=ContentType.TEXT)
|
|
|
|
# Should use TEXT even though it looks like JSON
|
|
assert result.content_type == ContentType.TEXT
|
|
|
|
def test_compression_result_has_metadata(self, compressor):
|
|
"""Test that result includes metadata."""
|
|
content = json.dumps({"items": [{"id": i} for i in range(20)]})
|
|
|
|
result = compressor.compress(content)
|
|
|
|
assert "detection" in result.metadata
|
|
assert "handler" in result.metadata
|
|
|
|
def test_register_custom_handler(self, compressor):
|
|
"""Test registering a custom handler."""
|
|
custom_handler = NoOpHandler()
|
|
compressor.register_handler(ContentType.JSON, custom_handler)
|
|
|
|
content = '{"key": "value"}' + " " * 100
|
|
|
|
result = compressor.compress(content)
|
|
|
|
# Should use our custom handler
|
|
assert result.handler_used == "noop"
|
|
|
|
def test_get_handler(self, compressor):
|
|
"""Test getting handler for content type."""
|
|
json_handler = compressor.get_handler(ContentType.JSON)
|
|
assert json_handler is not None
|
|
assert json_handler.name == "json"
|
|
|
|
unknown_handler = compressor.get_handler(ContentType.UNKNOWN)
|
|
assert unknown_handler.name == "noop"
|
|
|
|
|
|
class TestUniversalCompressorBatch:
|
|
"""Tests for batch compression."""
|
|
|
|
@pytest.fixture
|
|
def compressor(self):
|
|
"""Create compressor with fallback detector."""
|
|
config = UniversalCompressorConfig(
|
|
use_magika=False,
|
|
use_kompress=False,
|
|
ccr_enabled=False,
|
|
)
|
|
return UniversalCompressor(config=config)
|
|
|
|
def test_compress_batch_empty(self, compressor):
|
|
"""Test batch compression with empty list."""
|
|
results = compressor.compress_batch([])
|
|
assert results == []
|
|
|
|
def test_compress_batch_mixed_content(self, compressor):
|
|
"""Test batch compression with mixed content types."""
|
|
contents = [
|
|
json.dumps({"id": 1, "data": "x" * 100}),
|
|
"def foo(): pass\n" * 10,
|
|
"Plain text content " * 10,
|
|
]
|
|
|
|
results = compressor.compress_batch(contents)
|
|
|
|
assert len(results) == 3
|
|
assert results[0].content_type == ContentType.JSON
|
|
assert results[1].content_type == ContentType.CODE
|
|
assert results[2].content_type == ContentType.TEXT
|
|
|
|
|
|
class TestCompressFunction:
|
|
"""Tests for the convenience compress function."""
|
|
|
|
def test_compress_function(self):
|
|
"""Test one-off compression function."""
|
|
content = json.dumps({"items": [{"id": i} for i in range(20)]})
|
|
|
|
result = compress(content)
|
|
|
|
assert isinstance(result, CompressionResult)
|
|
assert result.content_type == ContentType.JSON
|
|
|
|
|
|
class TestStructurePreservation:
|
|
"""Integration tests for structure preservation."""
|
|
|
|
@pytest.fixture
|
|
def compressor(self):
|
|
"""Create compressor."""
|
|
config = UniversalCompressorConfig(
|
|
use_magika=False,
|
|
use_kompress=False,
|
|
ccr_enabled=False,
|
|
)
|
|
return UniversalCompressor(config=config)
|
|
|
|
def test_json_keys_preserved(self, compressor):
|
|
"""Test that JSON keys are visible after compression."""
|
|
data = {
|
|
"user_id": "12345",
|
|
"user_name": "Alice",
|
|
"user_email": "alice@example.com",
|
|
"user_bio": "A very long biography that goes on and on " * 10,
|
|
}
|
|
content = json.dumps(data)
|
|
|
|
result = compressor.compress(content)
|
|
|
|
# All keys should be visible in compressed output
|
|
for key in data.keys():
|
|
assert key in result.compressed, f"Key {key} should be in compressed output"
|
|
|
|
def test_code_signatures_preserved(self, compressor):
|
|
"""Test that code signatures are visible after compression."""
|
|
content = (
|
|
'''
|
|
def calculate_total(items, tax_rate=0.1):
|
|
"""Calculate total with tax."""
|
|
subtotal = sum(item.price for item in items)
|
|
tax = subtotal * tax_rate
|
|
total = subtotal + tax
|
|
'''
|
|
+ "# padding " * 50
|
|
+ '''
|
|
return total
|
|
|
|
class ShoppingCart:
|
|
"""Shopping cart implementation."""
|
|
|
|
def __init__(self):
|
|
self.items = []
|
|
'''
|
|
+ "# more padding " * 30
|
|
+ '''
|
|
|
|
def add_item(self, item):
|
|
"""Add item to cart."""
|
|
self.items.append(item)
|
|
'''
|
|
)
|
|
result = compressor.compress(content)
|
|
|
|
# Function and class names should be visible
|
|
assert "calculate_total" in result.compressed
|
|
assert "ShoppingCart" in result.compressed
|
|
assert "add_item" in result.compressed
|
|
|
|
def test_compression_reduces_tokens(self, compressor):
|
|
"""Test that compression actually reduces token count."""
|
|
# Large content that should be compressible
|
|
data = {
|
|
"results": [
|
|
{
|
|
"id": i,
|
|
"title": f"Result {i}",
|
|
"description": f"This is a detailed description for result {i}. " * 5,
|
|
}
|
|
for i in range(50)
|
|
]
|
|
}
|
|
content = json.dumps(data)
|
|
|
|
result = compressor.compress(content)
|
|
|
|
# Should achieve some compression
|
|
assert result.tokens_after < result.tokens_before
|
|
assert result.compression_ratio < 1.0
|
|
|
|
|
|
class TestSecurityAuditRegressions:
|
|
"""End-to-end regressions for the June 2026 security audit findings."""
|
|
|
|
def test_sec01_entropy_preservation_keeps_high_entropy_secret(self):
|
|
"""SEC-01: an enabled entropy pass must preserve high-entropy strings.
|
|
|
|
With the character-level tokenization bug, use_entropy_preservation was
|
|
a silent no-op, so a secret buried in a compressible region was dropped.
|
|
Uses the deterministic truncation fallback (use_kompress=False) so the
|
|
middle of the content is genuinely removed when preservation is off.
|
|
"""
|
|
secret = "Zx9Kq3Wm7Pv2Lr8Nt4Bc6Df1Gh5Jy" # gitleaks:allow synthetic test fixture
|
|
filler = "the quick brown fox jumps over the lazy dog and runs on. " * 6
|
|
content = filler + " " + secret + " " + filler
|
|
|
|
off = UniversalCompressor(
|
|
config=UniversalCompressorConfig(
|
|
use_magika=False,
|
|
use_kompress=False,
|
|
use_entropy_preservation=False,
|
|
ccr_enabled=False,
|
|
)
|
|
).compress(content, content_type=ContentType.TEXT)
|
|
on = UniversalCompressor(
|
|
config=UniversalCompressorConfig(
|
|
use_magika=False,
|
|
use_kompress=False,
|
|
use_entropy_preservation=True,
|
|
ccr_enabled=False,
|
|
)
|
|
).compress(content, content_type=ContentType.TEXT)
|
|
|
|
assert secret not in off.compressed # dropped without preservation
|
|
assert secret in on.compressed # preserved with the fix
|
|
|
|
def test_sec02_json_fallback_stays_valid_json(self):
|
|
"""SEC-02: the truncation fallback must not emit control chars into JSON.
|
|
|
|
The old separator embedded raw newlines, producing invalid JSON inside a
|
|
string value when Kompress was unavailable. Validate both paths.
|
|
"""
|
|
payload = json.dumps(
|
|
{
|
|
"service": "api",
|
|
"payload": "This is a fictional long string value " * 4 + "END",
|
|
}
|
|
)
|
|
for use_kompress in (False, True):
|
|
res = UniversalCompressor(
|
|
config=UniversalCompressorConfig(
|
|
use_magika=False,
|
|
use_kompress=use_kompress,
|
|
ccr_enabled=False,
|
|
)
|
|
).compress(payload, content_type=ContentType.JSON)
|
|
assert "\n" not in res.compressed, f"raw newline leaked (use_kompress={use_kompress})"
|
|
json.loads(res.compressed) # must not raise
|