"""Regression tests for ReDoS in agent_loop's `...` stripping.
CodeQL flagged `py/polynomial-redos` on the lazy `.*?` pattern
used in `src/agent_loop.py` (one compiled `_THINK_RE`, one inline copy). It is
applied with `re.sub` over a whole model response. When the closing delimiter
is missing, the engine rescans to end-of-string from every `` opener ->
O(n^2) on attacker-influenced input (prompt injection via tool output /
retrieved content echoed back by the model).
The fix replaces the regex with `_strip_think_blocks`, a forward-only linear
scan that is byte-for-byte equivalent to the original
`re.sub(r'.*?', '', text, flags=DOTALL|IGNORECASE)`.
These tests pin BOTH halves:
* output is identical to the reference regex for legitimate inputs, and
* pathological "many openers, no closer" input completes promptly.
"""
import re
import time
from src.agent_loop import _strip_think_blocks
# The exact pattern this fix replaces. Used only as an equivalence oracle on
# well-formed inputs (never on the adversarial one, where it is the slow path).
_REFERENCE_RE = re.compile(r".*?", re.DOTALL | re.IGNORECASE)
def _reference(text: str) -> str:
return _REFERENCE_RE.sub("", text or "")
# Loose ceiling: the linear helper finishes in well under 100ms; the vulnerable
# regex took seconds-to-tens-of-seconds on the same input.
_BUDGET_S = 4.0
# -- equivalence with the original regex -------------------------------------
EQUIV_CASES = [
"",
"no tags here at all",
"hiddenvisible",
"beforecotafter",
"aonebtwoc",
"only",
"tail",
"anestedrest", # lazy stops at first closer
"leadingorphanx", # orphan closer is NOT stripped
"trailingno closer for this one", # dangling opener kept verbatim
"CASE UP mix x", # case-insensitive
"multi\nline\na\nb\nc\nkeep", # DOTALL across newlines
"not matched by narrow regex", # only literal
"space-in-tag not matched", # literal tag only
]
def test_strip_think_blocks_matches_reference_regex():
for case in EQUIV_CASES:
assert _strip_think_blocks(case) == _reference(case), repr(case)
def test_empty_and_none_safe():
assert _strip_think_blocks("") == ""
assert _strip_think_blocks(None) in (None, "")
# -- ReDoS bound -------------------------------------------------------------
def test_many_openers_no_closer_is_linear():
# Attacker echoes thousands of "" with no closer. The lazy regex
# rescans to EOS from each opener (O(n^2)); the helper scans once.
hostile = "" * 60_000 + "x"
start = time.perf_counter()
out = _strip_think_blocks(hostile)
elapsed = time.perf_counter() - start
# No closer anywhere -> nothing is stripped, input returned intact.
assert out == hostile
assert elapsed < _BUDGET_S, f"took {elapsed:.2f}s (expected linear)"
def test_openers_then_one_far_closer_is_linear():
hostile = "" * 60_000 + "" + "tail"
start = time.perf_counter()
out = _strip_think_blocks(hostile)
elapsed = time.perf_counter() - start
# First opener pairs with the single closer; lazy match spans to it.
assert out == "tail"
assert elapsed < _BUDGET_S, f"took {elapsed:.2f}s (expected linear)"