diff --git a/api/helpers.py b/api/helpers.py index c53b402b..e049c1ab 100644 --- a/api/helpers.py +++ b/api/helpers.py @@ -9,6 +9,7 @@ import logging import os import re as _re import ssl +import sys from pathlib import Path from api.config import IMAGE_EXTS, MD_EXTS @@ -472,31 +473,26 @@ _redact_fn_lru = functools.lru_cache(maxsize=4096)(_redact_fn_uncached) # thousands of small recurring strings that actually benefit, or balloon RSS. _REDACT_CACHE_MAX_TEXT_LEN = 16384 -# Strings above that threshold used to bypass memoization entirely and re-run -# the full ~15-pass redactor on every request. On a real 22MB session that was -# 59 large strings (29 unique, 0.76MB) costing 1.68s per request — 99.9% of the -# recurring redaction cost once the small-string cache is warm, paid again by -# every tab and every poll for a byte-identical result. +# Strings above that threshold deliberately stay UNCACHED and re-run the full +# redactor on every request, even though they dominate the recurring cost of a +# large session (measured: 59 large strings, 29 unique, 1.68s per request on a +# real 22MB session). Memoizing them by input text alone is not safe: +# ``agent.redact.register_redaction_patterns()`` lets a plugin extend the +# secret matcher at runtime, and the installed registry exposes neither a +# policy generation the cache key could include nor a hook that could clear +# it. A large blob primed before such a registration would keep being served +# with the newly-registered secret intact through session, SSE and +# public-share projections. Until the agent registry exposes a generation, +# large strings fail closed. # -# They get their own small cache instead of sharing the 4096-entry one, so a -# large blob can never evict the thousands of small recurring strings. Both the -# entry count and the per-entry size are capped, which bounds retained bytes to -# roughly _REDACT_LARGE_CACHE_SIZE * _REDACT_LARGE_CACHE_MAX_TEXT_LEN (key plus -# value). Anything larger stays uncached: unbounded growth is the failure mode -# this threshold exists to prevent. -_REDACT_LARGE_CACHE_MAX_TEXT_LEN = 131072 -_REDACT_LARGE_CACHE_SIZE = 64 - -_redact_fn_large_lru = functools.lru_cache(maxsize=_REDACT_LARGE_CACHE_SIZE)( - _redact_fn_uncached -) +# tests/test_redact_large_string_cache.py pins this against the real agent +# registry: prime a large string, register a pattern, the next response must +# be redacted. def _redact_fn_cached(text): if len(text) > _REDACT_CACHE_MAX_TEXT_LEN: - if len(text) > _REDACT_LARGE_CACHE_MAX_TEXT_LEN: - return _redact_fn_uncached(text) - return _redact_fn_large_lru(text) + return _redact_fn_uncached(text) return _redact_fn_lru(text) @@ -593,11 +589,21 @@ _SENSITIVE_PHONE_MARKER_RE = _re.compile(r"(? bool: @@ -626,12 +632,13 @@ _might_contain_sensitive_text_lru = functools.lru_cache( def _might_contain_sensitive_text(text: str) -> bool: """Memoized wrapper around the prefilter. - Falls back to the uncached scan for non-strings and oversized text so the - cache can never be poisoned by an unhashable value or grow unbounded. + Falls back to the uncached scan for non-strings and for any string whose + object size exceeds the per-entry byte gate, so the cache can never be + poisoned by an unhashable value and its retained bytes stay hard-bounded. """ if not isinstance(text, str) or not text: return False - if len(text) > _SENSITIVE_PREFILTER_MAX_TEXT_LEN: + if sys.getsizeof(text) > _SENSITIVE_PREFILTER_MAX_ENTRY_BYTES: return _might_contain_sensitive_text_uncached(text) return _might_contain_sensitive_text_lru(text) diff --git a/tests/test_redact_large_string_cache.py b/tests/test_redact_large_string_cache.py index 400b0d1b..36572acf 100644 --- a/tests/test_redact_large_string_cache.py +++ b/tests/test_redact_large_string_cache.py @@ -1,21 +1,24 @@ -"""Large strings: memoize instead of re-paying ~15 regexes on every request. +"""Large strings must NOT be memoized by input text alone. -``_redact_fn_cached`` excluded every string > 16,384 characters from the LRU, -to avoid evicting the thousands of small recurring strings and inflating RSS. -Measured consequence on a real 22 MB session: 59 large strings (29 unique, -0.76 MB) went through the full regex set on EVERY request, for an always -identical result -- 1.68 s per request, i.e. 99.9% of the recurring redaction -cost once the small-string cache is warm. +``_redact_fn_cached`` memoizes strings up to 16,384 characters. An earlier +revision of this branch added a second LRU for larger strings, keyed only on +the text. That is unsafe: ``agent.redact.register_redaction_patterns()`` lets +a plugin extend the secret matcher at runtime, and the installed registry +exposes no policy generation the cache key could include and no hook that +could clear the cache. A large blob primed BEFORE such a registration would +keep being served with the newly-registered secret intact through session, +SSE and public-share projections. -The fix adds a SECOND cache, dedicated to large strings, bounded both by entry -count AND by a maximum size per entry, so that RSS stays capped. Giant strings -(> cap) deliberately remain uncached. +Until the agent registry exposes a generation, strings above the small-cache +threshold fail closed: they run the full redactor on every call. Contract verified here: - 1. a large string is redacted exactly as before; - 2. the same large string is not recomputed on the second call; - 3. the two caches stay separate (no cross-eviction); - 4. a string above the cap is never retained. + 1. a large string is redacted exactly like the uncached redactor; + 2. no module-level cache retains a large string (only the small LRU exists + and it never sees a string above the threshold); + 3. production-registry regression: prime a large string carrying a token + the redactor does not know yet, register a runtime pattern for it, and + the NEXT ``redact_session_data`` response is redacted. """ import sys from pathlib import Path @@ -27,36 +30,25 @@ sys.path.insert(0, str(Path(__file__).resolve().parents[1])) from api import helpers # noqa: E402 from api.helpers import ( # noqa: E402 _REDACT_CACHE_MAX_TEXT_LEN, - _REDACT_LARGE_CACHE_MAX_TEXT_LEN, _redact_fn_cached, _redact_fn_uncached, + _redact_text, + redact_session_data, ) SECRET = "gh" + "p_" + "0123456789abcdefghijklmnopqrstuvwxyzAB" @pytest.fixture(autouse=True) -def _isolate_large_cache(): - """Do not let this file pollute the process-wide global state. - - The cache is a module singleton shared by the whole suite. Clearing it - without restoring it makes the following tests start cold, which disturbs - those that depend on freshness windows or response times. We start from an - empty cache AND clear it again on exit, so that this file has no observable - effect outside of itself. - """ - helpers._redact_fn_large_lru.cache_clear() +def _isolate_small_cache(): + """Start and finish with an empty small-string LRU.""" + helpers._redact_fn_lru.cache_clear() yield - helpers._redact_fn_large_lru.cache_clear() + helpers._redact_fn_lru.cache_clear() def _big(secret: str, size: int) -> str: - """String > small-cache threshold, carrying a secret. - - The filler must be long enough to reach ``size`` even beyond the large - cache cap, otherwise truncation yields a string that is too short and the - test no longer verifies what it claims. - """ + """String of exactly ``size`` characters carrying ``secret``.""" unit = "lorem ipsum dolor sit amet " filler = unit * (size // len(unit) + 1) text = (filler + secret + filler)[:size] @@ -71,44 +63,84 @@ def test_large_string_redaction_matches_uncached(): assert SECRET not in _redact_fn_cached(text) -def test_large_string_is_memoized(): - helpers._redact_fn_large_lru.cache_clear() +def test_large_string_is_never_memoized(): + """Above the threshold the redactor runs every time: no cache hit, no + cache entry, on any module-level LRU.""" text = _big(SECRET, _REDACT_CACHE_MAX_TEXT_LEN + 7000) first = _redact_fn_cached(text) - misses_after_first = helpers._redact_fn_large_lru.cache_info().misses - second = _redact_fn_cached(text) - info = helpers._redact_fn_large_lru.cache_info() assert second == first - assert info.hits >= 1, "large string recomputed on second call" - assert info.misses == misses_after_first, "unexpected additional miss" + info = helpers._redact_fn_lru.cache_info() + assert info.hits == 0 and info.misses == 0 and info.currsize == 0, ( + "large string reached the small-string LRU" + ) + assert not hasattr(helpers, "_redact_fn_large_lru"), ( + "a text-keyed large-string cache is back; it must carry a redaction " + "policy generation before it can exist" + ) -def test_small_strings_do_not_use_the_large_cache(): - helpers._redact_fn_large_lru.cache_clear() +def test_small_strings_still_use_the_small_cache(): small = f"small text {SECRET}" assert len(small) <= _REDACT_CACHE_MAX_TEXT_LEN - _redact_fn_cached(small) - info = helpers._redact_fn_large_lru.cache_info() - assert info.hits == 0 and info.misses == 0, "small string routed to the large cache" + first = _redact_fn_cached(small) + second = _redact_fn_cached(small) + assert first == second + assert SECRET not in first + assert helpers._redact_fn_lru.cache_info().hits >= 1 -def test_giant_string_is_not_retained(): - """Above the cap: correct, but never cached (bounded RSS).""" - helpers._redact_fn_large_lru.cache_clear() - giant = _big(SECRET, _REDACT_LARGE_CACHE_MAX_TEXT_LEN + 10000) - assert len(giant) > _REDACT_LARGE_CACHE_MAX_TEXT_LEN +# ── Production-registry regression ───────────────────────────────────────── +# Uses the REAL installed agent.redact registry (skipped when the agent is not +# importable): the module-level `_redact_fn_uncached` is `_combined_redact` +# only when `agent.redact` imported at helpers import time. - out = _redact_fn_cached(giant) - - assert out == _redact_fn_uncached(giant) - assert SECRET not in out - info = helpers._redact_fn_large_lru.cache_info() - assert info.currsize == 0, "giant string retained in cache" +_RUNTIME_TOKEN = "nvapi-" + "A1b2C3d4E5f6G7h8I9j0K1l2M3n4O5p6Q7r8S9t0U1v2W3x4Y5z6A7b8C9d0" +_RUNTIME_PATTERN = r"nvapi-[A-Za-z0-9]{60}" -def test_large_cache_is_bounded(): - assert helpers._redact_fn_large_lru.cache_info().maxsize is not None - assert helpers._redact_fn_large_lru.cache_info().maxsize <= 128 +@pytest.fixture +def agent_registry(): + redact = pytest.importorskip("agent.redact", reason="hermes-agent not installed") + if not hasattr(redact, "register_redaction_patterns"): + pytest.skip("installed agent.redact has no runtime pattern registry") + if helpers._redact_fn_uncached.__name__ != "_combined_redact": + pytest.skip("api.helpers was built without the agent redactor") + reset = getattr(redact, "_reset_plugin_redaction_patterns", None) + if reset is None: + pytest.skip("installed agent.redact has no registry reset seam") + reset() + yield redact + reset() + + +def test_runtime_pattern_registration_is_honored_for_large_strings(agent_registry): + """Prime → register → the next response must be redacted (fail-closed).""" + # A marker the WebUI prefilter already trips on, so the large string goes + # through the redactor both before and after registration. + text = ("see https://example.invalid/docs lorem ipsum " * 500) + _RUNTIME_TOKEN + " tail" + assert len(text) > _REDACT_CACHE_MAX_TEXT_LEN + assert helpers._might_contain_sensitive_text(text) + + # Sanity: the installed registry does not know this token shape yet. + # If it ever does, the test still proves the cache path stays honest. + primed = _redact_text(text, _enabled=True) + session_before = redact_session_data( + {"messages": [{"role": "assistant", "content": text}]} + ) + assert session_before["messages"][0]["content"] == primed + + accepted = agent_registry.register_redaction_patterns([_RUNTIME_PATTERN], source="test-7276") + assert accepted == 1, "the registry rejected the probe pattern" + + oracle = _redact_fn_uncached(text) + assert _RUNTIME_TOKEN not in oracle, "the registered pattern is not effective" + + assert _redact_text(text, _enabled=True) == oracle + session_after = redact_session_data( + {"messages": [{"role": "assistant", "content": text}]} + ) + assert _RUNTIME_TOKEN not in session_after["messages"][0]["content"] + assert session_after["messages"][0]["content"] == oracle diff --git a/tests/test_sensitive_prefilter_equivalence.py b/tests/test_sensitive_prefilter_equivalence.py index 75f1dc07..376ff7e5 100644 --- a/tests/test_sensitive_prefilter_equivalence.py +++ b/tests/test_sensitive_prefilter_equivalence.py @@ -11,9 +11,15 @@ lists or the cache wrapper is caught here rather than in production. They must not leave global state behind: the LRU is shared process-wide, so each test that touches it restores it (see the autouse fixture). + +The cache retains its keys, so it is also bounded in BYTES, not just in entry +count: ``sys.getsizeof(text)`` gates every entry (a character cap is not a +byte bound -- 16k four-byte code points is ~64 KiB) and the entry cap gives a +hard aggregate ceiling. The memory tests below pin that. """ import random +import subprocess import sys from pathlib import Path @@ -25,9 +31,17 @@ from api import helpers # noqa: E402 from api.helpers import ( # noqa: E402 _SENSITIVE_CASE_MARKERS, _SENSITIVE_LOWER_MARKERS, + _SENSITIVE_PREFILTER_CACHE_SIZE, + _SENSITIVE_PREFILTER_MAX_ENTRY_BYTES, + _SENSITIVE_PREFILTER_MAX_RETAINED_KEY_BYTES, _might_contain_sensitive_text, ) +REPO_ROOT = Path(__file__).resolve().parents[1] + +# Widest code point CPython stores (UCS-4): every character costs 4 bytes. +_WIDE = "\U0001d54f" + @pytest.fixture(autouse=True) def _isolate_prefilter_cache(): @@ -130,9 +144,10 @@ def test_repeated_calls_are_stable_and_cached(): def test_oversized_text_bypasses_the_cache_but_keeps_the_verdict(): """Huge strings must stay correct while never entering the LRU.""" - limit = helpers._SENSITIVE_PREFILTER_MAX_TEXT_LEN - big_secret = "x" * (limit + 1) + "sk-" + "b" * 40 + limit = _SENSITIVE_PREFILTER_MAX_ENTRY_BYTES + big_secret = "x" * limit + "sk-" + "b" * 40 big_clean = "y" * (limit + 100) + assert sys.getsizeof(big_secret) > limit and sys.getsizeof(big_clean) > limit before = helpers._might_contain_sensitive_text_lru.cache_info().currsize assert _might_contain_sensitive_text(big_secret) is True @@ -159,7 +174,175 @@ def test_unhashable_input_does_not_reach_the_cache(): def test_cache_is_bounded(): """The LRU must cap its own size rather than grow with traffic.""" - limit = helpers._SENSITIVE_PREFILTER_CACHE_SIZE + limit = _SENSITIVE_PREFILTER_CACHE_SIZE for i in range(limit + 500): _might_contain_sensitive_text(f"unique-clean-string-{i}") assert helpers._might_contain_sensitive_text_lru.cache_info().currsize <= limit + + +# ── Byte bound regressions ───────────────────────────────────────────────── + + +def _retained_key_bytes(inserted_in_order: list) -> int: + """Actual ``sys.getsizeof`` sum of the keys the LRU currently holds. + + ``functools.lru_cache`` exposes no key iterator, but it is a strict LRU + and the tests below never re-touch a key before calling this, so the + resident set is exactly the last ``currsize`` keys that were inserted. + """ + info = helpers._might_contain_sensitive_text_lru.cache_info() + resident = inserted_in_order[len(inserted_in_order) - info.currsize:] + assert len(resident) == info.currsize + return sum(sys.getsizeof(key) for key in resident) + + +def test_entry_gate_is_object_bytes_not_characters(): + """A short string of 4-byte code points must be judged on its byte size.""" + max_ascii = "a" * (_SENSITIVE_PREFILTER_MAX_ENTRY_BYTES - sys.getsizeof("")) + assert sys.getsizeof(max_ascii) <= _SENSITIVE_PREFILTER_MAX_ENTRY_BYTES + # Same CHARACTER count in the widest storage is ~4x the bytes: must bypass. + wide_same_len = _WIDE * len(max_ascii) + assert len(wide_same_len) == len(max_ascii) + assert sys.getsizeof(wide_same_len) > _SENSITIVE_PREFILTER_MAX_ENTRY_BYTES + + before = helpers._might_contain_sensitive_text_lru.cache_info().currsize + assert _might_contain_sensitive_text(max_ascii) is False + assert helpers._might_contain_sensitive_text_lru.cache_info().currsize == before + 1 + assert _might_contain_sensitive_text(wide_same_len) is False + assert helpers._might_contain_sensitive_text_lru.cache_info().currsize == before + 1, ( + "a wide-character string above the byte gate was retained" + ) + assert _might_contain_sensitive_text(wide_same_len) == _reference(wide_same_len) + + +def test_maximum_width_unicode_fills_cache_within_byte_budget(): + """Fill the cache with the largest admissible 4-byte-per-char keys and + check the aggregate stays under the hard byte budget.""" + # Largest wide string that still passes the gate. + n = (_SENSITIVE_PREFILTER_MAX_ENTRY_BYTES - sys.getsizeof(_WIDE)) // 4 + 1 + while sys.getsizeof(_WIDE * n) > _SENSITIVE_PREFILTER_MAX_ENTRY_BYTES: + n -= 1 + assert n > 0 + base = _WIDE * n + assert sys.getsizeof(base) <= _SENSITIVE_PREFILTER_MAX_ENTRY_BYTES + + # Unique keys: vary the LAST code point (keeps width 4, same byte size). + total = _SENSITIVE_PREFILTER_CACHE_SIZE + 64 + keys = [base[:-1] + chr(0x1D400 + i) for i in range(total)] + assert all(sys.getsizeof(k) == sys.getsizeof(base) for k in keys[:8]) + for key in keys: + assert _might_contain_sensitive_text(key) is False + + info = helpers._might_contain_sensitive_text_lru.cache_info() + assert info.currsize == _SENSITIVE_PREFILTER_CACHE_SIZE + retained = _retained_key_bytes(keys) + assert retained <= _SENSITIVE_PREFILTER_MAX_RETAINED_KEY_BYTES + # The bound is tight, not slack: max-width keys fill most of it. + assert retained > _SENSITIVE_PREFILTER_MAX_RETAINED_KEY_BYTES * 0.95 + # Concretely: 16 MiB, not the ~512 MiB a character cap would have allowed. + assert _SENSITIVE_PREFILTER_MAX_RETAINED_KEY_BYTES == 16 * 1024 * 1024 + + +def test_sensitive_output_pairs_are_byte_gated_too(): + """A SENSITIVE verdict is cached under the same byte gate: an oversized + secret-bearing string is neither retained nor misjudged, and an + admissible one is cached with the right verdict.""" + small_secret = "token=" + "s" * 40 + assert sys.getsizeof(small_secret) <= _SENSITIVE_PREFILTER_MAX_ENTRY_BYTES + wide_secret = _WIDE * 600 + " sk-" + "c" * 40 + assert sys.getsizeof(wide_secret) > _SENSITIVE_PREFILTER_MAX_ENTRY_BYTES + + before = helpers._might_contain_sensitive_text_lru.cache_info().currsize + assert _might_contain_sensitive_text(small_secret) is True + assert _might_contain_sensitive_text(small_secret) is True + assert helpers._might_contain_sensitive_text_lru.cache_info().currsize == before + 1 + assert helpers._might_contain_sensitive_text_lru.cache_info().hits >= 1 + + assert _might_contain_sensitive_text(wide_secret) is True + assert _might_contain_sensitive_text(wide_secret) is True + assert helpers._might_contain_sensitive_text_lru.cache_info().currsize == before + 1, ( + "oversized sensitive string retained" + ) + assert _might_contain_sensitive_text(wide_secret) == _reference(wide_secret) + + +def test_eviction_keeps_aggregate_under_budget_and_verdicts_correct(): + """Churn well past the entry cap with max-size keys of mixed verdicts: + the cache must evict (never exceed the cap) and every verdict must still + match the oracle, including for keys that were evicted and re-scanned.""" + fill = _SENSITIVE_PREFILTER_MAX_ENTRY_BYTES - sys.getsizeof("") - 64 + total = _SENSITIVE_PREFILTER_CACHE_SIZE * 2 + 17 + peak = 0 + inserted = [] + samples = [] + for i in range(total): + tag = f"{i:08d}" + if i % 3 == 0: + text = "k" * fill + " api_key=" + tag + else: + text = "k" * fill + " plain " + tag + assert sys.getsizeof(text) <= _SENSITIVE_PREFILTER_MAX_ENTRY_BYTES + assert _might_contain_sensitive_text(text) == _reference(text) + inserted.append(text) + peak = max(peak, helpers._might_contain_sensitive_text_lru.cache_info().currsize) + if i % 997 == 0: + samples.append(text) + assert peak <= _SENSITIVE_PREFILTER_CACHE_SIZE + assert helpers._might_contain_sensitive_text_lru.cache_info().currsize == _SENSITIVE_PREFILTER_CACHE_SIZE + assert _retained_key_bytes(inserted) <= _SENSITIVE_PREFILTER_MAX_RETAINED_KEY_BYTES + # Early samples were evicted; re-querying them must still be correct. + misses_before = helpers._might_contain_sensitive_text_lru.cache_info().misses + for text in samples: + assert _might_contain_sensitive_text(text) == _reference(text) + evicted = [t for t in samples if t not in inserted[-_SENSITIVE_PREFILTER_CACHE_SIZE:]] + assert evicted, "test did not churn past the cap" + assert helpers._might_contain_sensitive_text_lru.cache_info().misses - misses_before >= len(evicted) + + +_LOW_MEMORY_PROBE = r""" +import resource, sys +sys.path.insert(0, sys.argv[1]) +from api import helpers +# Address-space headroom well below the ~512 MiB a character-capped cache of +# wide keys would pin, but comfortably above the byte-bounded 16 MiB ceiling. +headroom = int(sys.argv[2]) * 1024 * 1024 +soft, hard = resource.getrlimit(resource.RLIMIT_AS) +current = 0 +for line in open("/proc/self/status"): + if line.startswith("VmSize:"): + current = int(line.split()[1]) * 1024 +resource.setrlimit(resource.RLIMIT_AS, (current + headroom, hard)) +wide = "\U0001d54f" +n = helpers._SENSITIVE_PREFILTER_MAX_ENTRY_BYTES // 4 +total = helpers._SENSITIVE_PREFILTER_CACHE_SIZE + 256 +entered = 0 +for i in range(total): + # Unique wide key of ~16k characters (~64 KiB each): what the old + # character cap admitted. Must be scanned, must NOT be retained. + text = wide * 16000 + chr(0x1D400 + (i % 500)) + f"{i:06d}" + verdict = helpers._might_contain_sensitive_text(text) + assert verdict is False, verdict + # And a unique key that DOES pass the byte gate, to keep the cache busy. + small = wide * (n - 8) + f"{i:06d}" + assert helpers._might_contain_sensitive_text(small) is False + entered += 1 +info = helpers._might_contain_sensitive_text_lru.cache_info() +assert info.currsize <= helpers._SENSITIVE_PREFILTER_CACHE_SIZE, info +print("OK", entered, info.currsize) +""" + + +@pytest.mark.skipif(sys.platform != "linux", reason="RLIMIT_AS + /proc probe is Linux-only") +def test_low_memory_headroom_survives_wide_unicode_churn(): + """Under ~96 MiB of address-space headroom, pushing more than the entry + cap of 16k-character wide strings through the prefilter must complete + without MemoryError. Before the byte gate this workload retained ~512 MiB + and died after a few thousand entries.""" + proc = subprocess.run( + [sys.executable, "-c", _LOW_MEMORY_PROBE, str(REPO_ROOT), "96"], + capture_output=True, + text=True, + timeout=300, + ) + assert proc.returncode == 0, proc.stderr[-4000:] + assert proc.stdout.strip().startswith("OK"), proc.stdout