diff --git a/mempalace/entity_detector.py b/mempalace/entity_detector.py index 85b54e8..df8968d 100644 --- a/mempalace/entity_detector.py +++ b/mempalace/entity_detector.py @@ -272,6 +272,34 @@ SKIP_FILENAMES = { # ==================== CANDIDATE EXTRACTION ==================== +# Entity-candidate matching is ReDoS-prone: the English candidate pattern's +# nested alternation (i18n/en.json) backtracks catastrophically on a long +# unbroken run of printable ASCII (base64, minified JS, hashes, data URIs), +# pinning a mine on one ~5000-char window for hours (#2063). Such a run is never +# a name, so it is collapsed to a space before single-word matching. Scope: +# * ASCII-only ([!-~]): non-ASCII scripts (CJK, Cyrillic, Devanagari, +# accented Latin) are deliberately left untouched — a CJK paragraph is one +# unbroken run with no ASCII whitespace, and those locales' candidate +# patterns are bounded and never backtrack, so collapsing them would +# silently destroy their entity detection. +# * The threshold sits above the 20-char cap of the simple-name pattern +# ([A-Z][a-z]{1,19}), so no real single-word name is dropped. A CamelCase +# code identifier glued inside a long ASCII run (a path/URL/dotted name) is +# dropped as noise; whitespace-delimited proper nouns are unaffected. +_MAX_CANDIDATE_TOKEN_LEN = 24 +_LONG_ASCII_RUN_RX = re.compile(rf"[!-~]{{{_MAX_CANDIDATE_TOKEN_LEN},}}") + + +def _collapse_long_ascii_runs(text: str) -> str: + """Collapse a long unbroken run of printable ASCII to a single space (#2063). + + Applied only before single-word candidate matching — the whitespace-delimited + multi-word patterns cannot backtrack on such runs. See ``_LONG_ASCII_RUN_RX`` + for scope and rationale. + """ + return _LONG_ASCII_RUN_RX.sub(" ", text) + + def extract_candidates(text: str, languages=("en",)) -> dict: """ Extract all capitalized proper noun candidates from text. @@ -296,13 +324,15 @@ def extract_candidates(text: str, languages=("en",)) -> dict: for compound, n in compound_counts.items(): counts[compound] += n - # Single-word candidates — one pre-wrapped pattern per language + # Single-word candidates — one pre-wrapped pattern per language. + # Collapse long ASCII blobs first (base64/minified) to defuse ReDoS (#2063). + candidate_text = _collapse_long_ascii_runs(working_text) for wrapped_pat in patterns["candidate_patterns"]: try: rx = re.compile(wrapped_pat) except re.error: continue - for word in rx.findall(working_text): + for word in rx.findall(candidate_text): wl = word.lower() if wl in stopwords: continue diff --git a/mempalace/palace.py b/mempalace/palace.py index 046078b..c9c1f7d 100644 --- a/mempalace/palace.py +++ b/mempalace/palace.py @@ -26,7 +26,11 @@ from .backends import ( resolve_backend_for_palace, ) from .backends.embedding_wrapper import EmbeddingCollection -from .entity_detector import _apply_known_systems_prepass, _get_coca_filter +from .entity_detector import ( + _apply_known_systems_prepass, + _collapse_long_ascii_runs, + _get_coca_filter, +) logger = logging.getLogger("mempalace_mcp") @@ -600,9 +604,12 @@ def _candidate_entity_words(text: str) -> list: except re.error: continue _CANDIDATE_RX_CACHE = rxs + # Defuse ReDoS on long ASCII blobs before matching (#2063); see + # entity_detector._collapse_long_ascii_runs. + stripped = _collapse_long_ascii_runs(text) words = [] for rx in _CANDIDATE_RX_CACHE: - words.extend(rx.findall(text)) + words.extend(rx.findall(stripped)) return words diff --git a/tests/test_entity_detector.py b/tests/test_entity_detector.py index 77868f5..4bfb7af 100644 --- a/tests/test_entity_detector.py +++ b/tests/test_entity_detector.py @@ -3,12 +3,20 @@ import contextlib import json import os +import re +import subprocess +import sys from pathlib import Path from unittest.mock import patch +from hypothesis import given +from hypothesis import strategies as st + from mempalace.entity_detector import ( PROSE_EXTENSIONS, STOPWORDS, + _MAX_CANDIDATE_TOKEN_LEN, + _collapse_long_ascii_runs, _print_entity_list, classify_entity, confirm_entities, @@ -1059,3 +1067,82 @@ def test_zh_tw_known_limitation_inline_name_no_boundary(): result = extract_candidates(text, languages=("zh-TW",)) # Extraction is expected to miss this adversarial case. assert "朱宜振" not in result + + +# ── ReDoS guard: long ASCII blobs (#2063) ────────────────────────────── + + +def test_collapse_long_ascii_runs_scope_and_threshold(): + # A long unbroken ASCII run collapses to a single space... + assert _collapse_long_ascii_runs("A" * 24) == " " + assert _collapse_long_ascii_runs("z" * 100) == " " + # ...but a run just below the threshold is kept verbatim (boundary 23/24) + assert _collapse_long_ascii_runs("Q" * 23) == "Q" * 23 + # whitespace-delimited natural text is untouched + assert _collapse_long_ascii_runs("Alice met Bob today") == "Alice met Bob today" + # spaces and newlines bound ASCII runs + assert _collapse_long_ascii_runs("ok\n" + "B" * 40 + " end") == "ok\n end" + # non-ASCII scripts are NEVER collapsed, even as one long unbroken run — a + # CJK/Cyrillic paragraph has no ASCII whitespace (guards the zh regression). + assert _collapse_long_ascii_runs("朱" * 40) == "朱" * 40 + assert _collapse_long_ascii_runs("Ы" * 40) == "Ы" * 40 + + +@given(st.text(max_size=200)) +def test_collapse_long_ascii_runs_leaves_no_collapsible_run(s): + """Property (any input): the output never contains a collapsible ASCII run, + so the class of inputs that can drive candidate matching into catastrophic + backtracking is eliminated — not just the reported sample (#2063).""" + out = _collapse_long_ascii_runs(s) + assert re.search(rf"[!-~]{{{_MAX_CANDIDATE_TOKEN_LEN},}}", out) is None + + +def test_extract_candidates_drops_overlong_ascii_blob_and_keeps_names(): + """A long unbroken ASCII run (base64/minified/hash) is out-of-domain: it must + not surface as an entity candidate, nor trigger the catastrophic backtracking + that pinned mines for hours (#2063). Normal names around it are still + detected.""" + longtok = "Aa" + "Bb" * 30 # 62-char unbroken ASCII run + text = (longtok + " ") * 3 + "Lantern ships Lantern plus Lantern here." + result = extract_candidates(text) + assert longtok not in result + assert "Lantern" in result + + +def test_extract_candidates_preserves_cjk_without_ascii_whitespace(): + """CJK text has no ASCII whitespace, so a whole Chinese paragraph is one + unbroken run. The mitigation collapses only long ASCII runs, so CJK entity + detection keeps working (#2063 regression guard for zh-CN/zh-TW, whose + candidate path has no multi-word fallback).""" + # 朱宜振 appears 3x, each flanked by full-width (non-ASCII) punctuation. + text = "朱宜振:這個方案沒問題。朱宜振,你負責前端。朱宜振,好,我來處理。今天大家都同意。" + result = extract_candidates(text, languages=("zh-TW",)) + assert "朱宜振" in result + + +def test_entity_extraction_no_redos_on_adversarial_ascii_run(): + """The real catastrophic trigger — a Cap+lower prefix, a long pure-uppercase + ASCII run, then a word char blocking the trailing word boundary — must + complete fast, not hang (#2063). Run in a subprocess with a hard timeout so + a regression fails fast instead of pinning CI (no pytest-timeout available).""" + code = ( + "from mempalace.entity_detector import extract_candidates\n" + "from mempalace.palace import _candidate_entity_words\n" + "payload = 'Aa' + 'B' * 400 + '0'\n" + "extract_candidates(payload)\n" + "_candidate_entity_words(payload)\n" + "print('OK')\n" + ) + try: + result = subprocess.run( + [sys.executable, "-c", code], + capture_output=True, + text=True, + timeout=30, + ) + except subprocess.TimeoutExpired: + raise AssertionError( + "entity extraction hung on an adversarial ASCII run — ReDoS regression (#2063)" + ) + assert result.returncode == 0, result.stderr + assert "OK" in result.stdout diff --git a/tests/test_palace.py b/tests/test_palace.py index a07e971..33f8d4a 100644 --- a/tests/test_palace.py +++ b/tests/test_palace.py @@ -6,6 +6,7 @@ from _chroma_palace_helper import make_minimal_chroma_sqlite from mempalace.backends import CollectionNotInitializedError, PalaceNotFoundError from mempalace.palace import ( + _candidate_entity_words, _metadata_matches_extract_mode, _open_collection_or_explain, backend_requires_single_writer, @@ -255,3 +256,13 @@ def test_open_collection_or_explain_distinguishes_collection_subclass(tmp_path, assert result is None assert any("initialized but empty" in line for line in lines) assert not any("No palace found" in line for line in lines) + + +def test_candidate_entity_words_drops_overlong_blob(): + """#2063: a long unbroken ASCII run must be collapsed before matching so the + candidate patterns cannot backtrack catastrophically; such runs are never + entity names. Normal names are still returned.""" + longtok = "Aa" + "Bb" * 30 # 62-char unbroken ASCII run + words = _candidate_entity_words(longtok + " and Lantern") + assert longtok not in words + assert "Lantern" in words