correction: Bound alignment cache growth

src/codex32/correction.py

@@ -23,7 +23,7 @@
 
 from collections.abc import Callable, Sequence
 from dataclasses import dataclass, replace
-from functools import cache
+from functools import cache, lru_cache
 from math import comb
 from time import monotonic
 from typing import Literal
@@ -250,6 +250,10 @@ class _Spec:
     1023,
 )
 
+# Keep every target length from one generic unknown-length search hot without
+# letting arbitrary HRPs create process-lifetime state.
+_ALIGNMENT_CACHE_SIZE = 11
+
 
 def _spec_for_checksum(checksum: _Checksum) -> _Spec:
     if checksum is _CODEX32:
@@ -276,7 +280,7 @@ def _pack_syndromes(values: tuple[int, ...]) -> int:
     return sum(value << (10 * index) for index, value in enumerate(values))
 
 
-@cache
+@lru_cache(maxsize=_ALIGNMENT_CACHE_SIZE)
 def _syndrome_alignment(spec: _Spec, hrp: str, length: int) -> tuple[int, tuple[tuple[int, ...], ...]]:
     base = _pack_syndromes(_syndromes(spec, _residue(spec, hrp, [0] * length), target=True))
     powers = _poly_powers(spec.generator, length)

tests/test_alignment.py

@@ -6,6 +6,7 @@
 
 from codex32._alignment import _IncrementalSyndromes, _View
 from codex32.correction import (
+    _ALIGNMENT_CACHE_SIZE,
     _LONG_SPEC,
     _SHORT_SPEC,
     _pack_syndromes,
@@ -15,6 +16,45 @@
 )
 
 
+def test_syndrome_alignment_cache_is_bounded_for_untrusted_hrps():
+    _syndrome_alignment.cache_clear()
+
+    for index in range(_ALIGNMENT_CACHE_SIZE + 5):
+        _syndrome_alignment(_SHORT_SPEC, f"attacker{index}", 45)
+
+    info = _syndrome_alignment.cache_info()
+    assert info.maxsize == _ALIGNMENT_CACHE_SIZE
+    assert info.currsize == _ALIGNMENT_CACHE_SIZE
+
+
+def test_syndrome_alignment_cache_fits_one_unknown_length_search():
+    observed = 50
+    targets = sorted({observed + delta for delta in (*range(-4, 5), -8, 8)})
+    _syndrome_alignment.cache_clear()
+
+    first = [_syndrome_alignment(_SHORT_SPEC, "generic", target) for target in targets]
+    before = _syndrome_alignment.cache_info()
+    second = [_syndrome_alignment(_SHORT_SPEC, "generic", target) for target in targets]
+    after = _syndrome_alignment.cache_info()
+
+    assert len(targets) == _ALIGNMENT_CACHE_SIZE
+    assert all(left is right for left, right in zip(first, second, strict=True))
+    assert after.misses == before.misses
+    assert after.hits == before.hits + len(targets)
+
+
+def test_syndrome_alignment_cache_still_reuses_recent_entries():
+    _syndrome_alignment.cache_clear()
+    first = _syndrome_alignment(_SHORT_SPEC, "ms", 45)
+    before = _syndrome_alignment.cache_info()
+
+    second = _syndrome_alignment(_SHORT_SPEC, "ms", 45)
+
+    after = _syndrome_alignment.cache_info()
+    assert first is second
+    assert after.hits == before.hits + 1
+
+
 @pytest.mark.parametrize(
     ("hrp", "length", "spec"),
     (