diff --git a/src/codex32/correction.py b/src/codex32/correction.py index 4a1a561..c30ce7c 100644 --- a/src/codex32/correction.py +++ b/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) diff --git a/tests/test_alignment.py b/tests/test_alignment.py index 8543fa2..b67b8ac 100644 --- a/tests/test_alignment.py +++ b/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"), (