headroom/tests/test_compression/test_universal.py
Tejas Chopra 840871cb96
fix(compression): repair entropy preservation + JSON-safe truncation fallback (#1536)
## 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.
2026-06-28 10:39:02 -07:00

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