diff --git a/cron/lifecycle_guard.py b/cron/lifecycle_guard.py index fb4c18e91c..e4ab23cd4f 100644 --- a/cron/lifecycle_guard.py +++ b/cron/lifecycle_guard.py @@ -291,9 +291,100 @@ _SHELL_EXECUTABLES = frozenset({"sh", "bash", "dash", "ksh", "zsh"}) _SHELL_OPTIONS_WITH_VALUES = frozenset({"-O", "+O", "-o", "+o"}) _MAX_REFERENCED_SCRIPT_BYTES = 1024 * 1024 _MAX_REFERENCED_SCRIPT_DEPTH = 8 +# Whole-walk work limits (#78398). The per-file byte cap and recursion depth +# above bound one read, not the walk: a command can reference arbitrarily many +# files, and the pure-Python shlex lexer is expensive on thousands of short +# lines (one lexer per line) and quadratic on one enormous token. In +# production that unbounded breadth held the GIL for minutes on every gateway +# terminal call. The budget is shared across one complete walk and charged +# BEFORE any text reaches shlex. +# +# Exhaustion is fail-closed, matching the existing contract for one oversized +# file: an unscanned referenced script could hide a lifecycle command. The +# limits are sized well above any real wrapper graph (4x the per-file cap, +# 16k lines, 1024 distinct scripts, 64 remote reads) so legitimate commands +# never reach them — an exhausted walk is logged at WARNING so an operator +# can tell it apart from a genuine lifecycle-command block. +# +# Local file reads are microseconds and every reference already costs a line +# of budget, so the path cap is generous; a remote read is a backend +# roundtrip, so it gets its own much tighter cap. +_MAX_LIFECYCLE_SCAN_BYTES = 4 * _MAX_REFERENCED_SCRIPT_BYTES +_MAX_LIFECYCLE_SCAN_LINES = 16384 +_MAX_LIFECYCLE_SCAN_LINE_BYTES = 256 * 1024 +_MAX_LIFECYCLE_SCAN_PATHS = 1024 +_MAX_LIFECYCLE_SCAN_REMOTE_READS = 64 _CONTROL_CHARS = frozenset(";&|()") +class _LifecycleScanBudget: + """Shared work budget for one complete referenced-script walk.""" + + __slots__ = ( + "bytes_remaining", + "lines_remaining", + "paths_remaining", + "remote_reads_remaining", + ) + + def __init__(self) -> None: + # Read the module constants at construction so tests (and operators) + # can lower them without defaults capturing stale values at import. + self.bytes_remaining = _MAX_LIFECYCLE_SCAN_BYTES + self.lines_remaining = _MAX_LIFECYCLE_SCAN_LINES + self.paths_remaining = _MAX_LIFECYCLE_SCAN_PATHS + self.remote_reads_remaining = _MAX_LIFECYCLE_SCAN_REMOTE_READS + + def charge_text(self, text: str) -> bool: + """Charge *text* before tokenization; False when it does not fit.""" + # UTF-8 is at least one byte per code point, so the character count + # is a free lower bound — skip the encode for obviously-oversized input. + if len(text) > self.bytes_remaining: + return False + encoded = len(text.encode("utf-8", errors="replace")) + if encoded > self.bytes_remaining: + return False + lines = text.count("\n") + 1 + if lines > self.lines_remaining: + return False + # One huge token is the quadratic shlex case; bound the longest + # physical line (bytes >= chars, so a char check is sufficient to + # reject and the encode is only needed on the boundary). + longest = max((len(line) for line in text.split("\n")), default=0) + if longest > _MAX_LIFECYCLE_SCAN_LINE_BYTES: + return False + self.bytes_remaining -= encoded + self.lines_remaining -= lines + return True + + def charge_path(self) -> bool: + """Charge one unique referenced path before any local/remote read.""" + if self.paths_remaining <= 0: + return False + self.paths_remaining -= 1 + return True + + def charge_remote_read(self) -> bool: + """Charge one remote-backend read (a network roundtrip each).""" + if self.remote_reads_remaining <= 0: + return False + self.remote_reads_remaining -= 1 + return True + + +def lifecycle_scan_root_within_budget(text: str) -> bool: + """Whether *text* may safely enter an optional tokenizer pass. + + Used by ``tools/terminal_tool.py`` to gate its launchctl-specific pre-scan + (which tokenizes with shlex). ``False`` is not a verdict: callers must + still run the full guard, which fails closed for an over-budget root. + """ + try: + return _LifecycleScanBudget().charge_text(text) + except Exception: + return False + + # Directory names that sit directly under a `Library` path component and # mark a FileProvider-backed subtree: `Mobile Documents` is iCloud Drive; # `CloudStorage` hosts every third-party FileProvider domain (Dropbox, @@ -909,9 +1000,14 @@ def _has_binary_magic(data: bytes) -> bool: return data.startswith(_BINARY_MAGICS) -def _read_referenced_script(path: Path) -> tuple[Optional[str], bool]: +def _read_referenced_script( + path: Path, *, max_bytes: Optional[int] = None +) -> tuple[Optional[str], bool]: """Return ``(text, unsafe)`` using bounded, regular-file-only reads. + ``max_bytes`` lowers the per-file cap to what the calling walk can still + afford (never raises it above ``_MAX_REFERENCED_SCRIPT_BYTES``). + This is the shared choke point for every local script read the guard performs (the terminal walk in ``_contains_unsafe_gateway_action`` AND the cron-script scan in ``_read_script_for_scanning``), so the @@ -921,6 +1017,9 @@ def _read_referenced_script(path: Path) -> tuple[Optional[str], bool]: (#88052). The lexical check covers direct cloud paths; the resolved check covers local launchers that are symlinks into a cloud subtree. """ + byte_limit = _MAX_REFERENCED_SCRIPT_BYTES + if max_bytes is not None: + byte_limit = min(byte_limit, max(0, int(max_bytes))) if _is_cloud_placeholder_path(path): return None, True try: @@ -964,12 +1063,14 @@ def _read_referenced_script(path: Path) -> tuple[Optional[str], bool]: data = os.read(descriptor, _BINARY_SNIFF_BYTES) if data.startswith(_BINARY_MAGIC_PREFIXES): return None, False + # A regular file whose size already exceeds the cap fails closed + # without reading it (the walk budget can be far below 1 MiB). + if metadata.st_size > byte_limit: + return None, True # Read the remainder (bounded). Loop because os.read may return # short for non-regular-file-backed descriptors. - while len(data) <= _MAX_REFERENCED_SCRIPT_BYTES: - chunk = os.read( - descriptor, _MAX_REFERENCED_SCRIPT_BYTES + 1 - len(data) - ) + while len(data) <= byte_limit: + chunk = os.read(descriptor, byte_limit + 1 - len(data)) if not chunk: break data += chunk @@ -992,14 +1093,16 @@ def _read_referenced_script(path: Path) -> tuple[Optional[str], bool]: # Check the size BEFORE stripping: stripping shrinks the buffer, so doing it # first would let an oversized file slip under the threshold and skip this # fail-closed branch. - if len(data) > _MAX_REFERENCED_SCRIPT_BYTES: + if len(data) > byte_limit: return None, True if b"\x00" in data: data = data.replace(b"\x00", b"") return data.decode("utf-8", errors="replace"), False -def _sanitize_remote_script_text(text: Optional[str]) -> tuple[Optional[str], bool]: +def _sanitize_remote_script_text( + text: Optional[str], *, max_bytes: Optional[int] = None +) -> tuple[Optional[str], bool]: """Apply the local-read contract to text from a ``read_remote_script`` callback. The recursion boundary must not trust its callbacks: any backend (SSH, @@ -1020,19 +1123,39 @@ def _sanitize_remote_script_text(text: Optional[str]) -> tuple[Optional[str], bo return None, False if "\x00" in text: return None, False - if len(text.encode("utf-8", errors="replace")) > _MAX_REFERENCED_SCRIPT_BYTES: + byte_limit = _MAX_REFERENCED_SCRIPT_BYTES + if max_bytes is not None: + byte_limit = min(byte_limit, max(0, int(max_bytes))) + if len(text) > byte_limit: + return None, True # chars <= bytes: over the cap without encoding + if len(text.encode("utf-8", errors="replace")) > byte_limit: return None, True return text, False +def _budget_exhausted(what: str, depth: int) -> bool: + logger.warning( + "lifecycle guard scan budget exhausted (%s at depth %d); " + "failing closed — see _MAX_LIFECYCLE_SCAN_* in cron/lifecycle_guard.py", + what, depth, + ) + return True + + def _contains_unsafe_gateway_action( command: str, *, cwd: Optional[str], depth: int, visited: set[Path], + budget: _LifecycleScanBudget, read_remote_script: Optional[_ReadRemoteScriptFn] = None, ) -> bool: + # Charge BEFORE _direct_lifecycle_scan: every scan in it (including the + # wrapper-prefix lifecycle detector) tokenizes with shlex, so checking + # afterwards would keep the CPU spike. + if not budget.charge_text(command): + return _budget_exhausted("text", depth) if _direct_lifecycle_scan(command): return True if depth >= _MAX_REFERENCED_SCRIPT_DEPTH: @@ -1044,6 +1167,7 @@ def _contains_unsafe_gateway_action( cwd=cwd, depth=depth + 1, visited=visited, + budget=budget, read_remote_script=read_remote_script, ): return True @@ -1068,17 +1192,26 @@ def _contains_unsafe_gateway_action( return True if resolved in visited: continue + if not budget.charge_path(): + return _budget_exhausted("paths", depth) visited.add(resolved) - script_text, unsafe = _read_referenced_script(script_path) + # Never read more than the walk can still afford to tokenize; a file + # larger than the remainder fails closed exactly like an oversized one. + script_text, unsafe = _read_referenced_script( + script_path, max_bytes=budget.bytes_remaining + ) if unsafe: return True if script_text is None and read_remote_script is not None: # Local path missing; try the remote backend if one is available. + if not budget.charge_remote_read(): + return _budget_exhausted("remote reads", depth) # The callback's output crosses the same trust boundary as a # local read — sanitize it identically before it enters the # recursion (binary skip + size fail-closed). script_text, unsafe = _sanitize_remote_script_text( - read_remote_script(str(script_path)) + read_remote_script(str(script_path)), + max_bytes=budget.bytes_remaining, ) if unsafe: return True @@ -1092,6 +1225,7 @@ def _contains_unsafe_gateway_action( cwd=script_dir, depth=depth + 1, visited=visited, + budget=budget, read_remote_script=read_remote_script, ): return True @@ -1126,6 +1260,7 @@ def contains_gateway_lifecycle_command_or_referenced_script( cwd=cwd, depth=0, visited=set(), + budget=_LifecycleScanBudget(), read_remote_script=read_remote_script, ) except Exception: @@ -1254,7 +1389,12 @@ def check_gateway_lifecycle( # `hermes gateway restart` embedded in a .py script is still # blocked. Non-regular/oversized script files still fail closed # via the lifecycle-shaped sentinel in _read_script_for_scanning. - unsafe = _lifecycle_command_scan_with_data_exemption(combined) + # The data-exemption masker tokenizes the text with shlex, so it is + # charged against the same walk budget as the shell path (#78398). + if not _LifecycleScanBudget().charge_text(combined): + unsafe = _budget_exhausted("text", 0) + else: + unsafe = _lifecycle_command_scan_with_data_exemption(combined) else: script_dir = _resolve_script_directory(script) if script else None unsafe = contains_gateway_lifecycle_command_or_referenced_script( diff --git a/tests/cron/test_lifecycle_guard_budget.py b/tests/cron/test_lifecycle_guard_budget.py new file mode 100644 index 0000000000..181461744f --- /dev/null +++ b/tests/cron/test_lifecycle_guard_budget.py @@ -0,0 +1,241 @@ +"""Whole-walk work budget for the gateway lifecycle guard (#78398). + +The per-file byte cap and recursion depth bound one read, not the walk. These +tests pin the shared budget that bounds the whole referenced-script walk and +is charged *before* any text reaches ``shlex``. + +Budget constants are monkeypatched to tiny values so the tests are fast and +deterministic; ``_LifecycleScanBudget`` reads them at construction time. +""" + +from __future__ import annotations + +import pytest + +import cron.lifecycle_guard as lifecycle_guard + +guard = lifecycle_guard.contains_gateway_lifecycle_command_or_referenced_script + + +def _explode(*_args, **_kwargs): + raise AssertionError("over-budget text reached shlex") + + +# --- root command (depth 0) ----------------------------------------------- + + +def test_root_byte_limit_allows_exact_and_rejects_plus_one(monkeypatch): + monkeypatch.setattr(lifecycle_guard, "_MAX_LIFECYCLE_SCAN_BYTES", 8) + monkeypatch.setattr(lifecycle_guard, "_MAX_LIFECYCLE_SCAN_LINE_BYTES", 8) + + assert guard("x" * 8) is False + + monkeypatch.setattr(lifecycle_guard.shlex, "shlex", _explode) + assert guard("x" * 9) is True + + +def test_root_line_limit_allows_exact_and_rejects_plus_one(monkeypatch): + monkeypatch.setattr(lifecycle_guard, "_MAX_LIFECYCLE_SCAN_LINES", 2) + + assert guard("one\ntwo") is False + assert guard("one\ntwo\nthree") is True + + +def test_single_giant_line_rejected_before_shlex(monkeypatch): + """One enormous token is the quadratic shlex case.""" + monkeypatch.setattr(lifecycle_guard, "_MAX_LIFECYCLE_SCAN_LINE_BYTES", 8) + monkeypatch.setattr(lifecycle_guard.shlex, "shlex", _explode) + + assert guard("xxxxxxxxx\necho ok") is True + + +def test_root_budget_counts_utf8_bytes(monkeypatch): + monkeypatch.setattr(lifecycle_guard, "_MAX_LIFECYCLE_SCAN_BYTES", 4) + monkeypatch.setattr(lifecycle_guard, "_MAX_LIFECYCLE_SCAN_LINE_BYTES", 4) + + assert guard("éé") is False + assert guard("ééé") is True + + +def test_exhaustion_is_logged_at_warning(monkeypatch, caplog): + monkeypatch.setattr(lifecycle_guard, "_MAX_LIFECYCLE_SCAN_BYTES", 4) + monkeypatch.setattr(lifecycle_guard, "_MAX_LIFECYCLE_SCAN_LINE_BYTES", 4) + + with caplog.at_level("WARNING", logger=lifecycle_guard.logger.name): + assert guard("echo hello") is True + assert "budget exhausted" in caplog.text + + +def test_lifecycle_scan_root_within_budget_is_not_a_verdict(monkeypatch): + monkeypatch.setattr(lifecycle_guard, "_MAX_LIFECYCLE_SCAN_BYTES", 8) + monkeypatch.setattr(lifecycle_guard, "_MAX_LIFECYCLE_SCAN_LINE_BYTES", 8) + + assert lifecycle_guard.lifecycle_scan_root_within_budget("x" * 8) is True + assert lifecycle_guard.lifecycle_scan_root_within_budget("x" * 9) is False + + +# --- referenced-script walk ------------------------------------------------ + + +def test_unique_path_budget_bounds_reads_and_fails_closed(monkeypatch, tmp_path): + monkeypatch.setattr(lifecycle_guard, "_MAX_LIFECYCLE_SCAN_PATHS", 2) + for i in range(3): + (tmp_path / f"s{i}.sh").write_text("echo ok\n", encoding="utf-8") + + two = " && ".join(f"bash {tmp_path}/s{i}.sh" for i in range(2)) + three = " && ".join(f"bash {tmp_path}/s{i}.sh" for i in range(3)) + + assert guard(two) is False + assert guard(three) is True + + +def test_repeated_path_does_not_spend_unique_path_budget(monkeypatch, tmp_path): + monkeypatch.setattr(lifecycle_guard, "_MAX_LIFECYCLE_SCAN_PATHS", 1) + script = tmp_path / "s.sh" + script.write_text("echo ok\n", encoding="utf-8") + + assert guard(f"bash {script} && bash {script} && sh {script}") is False + + +def test_remote_read_budget_charged_before_remote_read(monkeypatch): + monkeypatch.setattr(lifecycle_guard, "_MAX_LIFECYCLE_SCAN_REMOTE_READS", 1) + reads: list[str] = [] + + def remote(path: str): + reads.append(path) + return "echo ok\n" + + assert ( + guard( + "bash /remote/a.sh && bash /remote/b.sh", + read_remote_script=remote, + ) + is True + ) + assert reads == ["/remote/a.sh"] + + +def test_cumulative_text_budget_bounds_recursive_scan(monkeypatch, tmp_path): + """Two scripts individually under the per-file cap exceed the walk cap. + + Relative references keep the root command short so the budget arithmetic + is about the scripts, not the tmp_path length.""" + monkeypatch.setattr(lifecycle_guard, "_MAX_LIFECYCLE_SCAN_BYTES", 48) + (tmp_path / "a.sh").write_text("echo " + "a" * 10 + "\n", encoding="utf-8") # 16 + (tmp_path / "b.sh").write_text("echo " + "b" * 10 + "\n", encoding="utf-8") # 16 + cwd = str(tmp_path) + + # 9 (root) + 16 fits in 48; 19 (root) + 16 + 16 does not → fail closed. + assert guard("bash a.sh", cwd=cwd) is False + assert guard("bash a.sh;bash b.sh", cwd=cwd) is True + + +def test_referenced_read_is_capped_at_remaining_budget(monkeypatch, tmp_path): + """A file bigger than what the walk can still afford is never read whole: + the read helper receives the remaining budget as its cap.""" + monkeypatch.setattr(lifecycle_guard, "_MAX_LIFECYCLE_SCAN_BYTES", 64) + (tmp_path / "big.sh").write_text("echo " + "x" * 200 + "\n", encoding="utf-8") + + caps: list = [] + original = lifecycle_guard._read_referenced_script + + def spy(path, *, max_bytes=None): + caps.append(max_bytes) + return original(path, max_bytes=max_bytes) + + monkeypatch.setattr(lifecycle_guard, "_read_referenced_script", spy) + + root = "bash big.sh" + assert guard(root, cwd=str(tmp_path)) is True + assert caps == [64 - len(root)] + + +def test_remote_script_sanitizer_honours_remaining_budget(): + text, unsafe = lifecycle_guard._sanitize_remote_script_text( + "echo ok\n", max_bytes=4 + ) + assert (text, unsafe) == (None, True) + text, unsafe = lifecycle_guard._sanitize_remote_script_text( + "echo ok\n", max_bytes=8 + ) + assert (text, unsafe) == ("echo ok\n", False) + + +def test_line_budget_fails_closed_before_tokenizing_every_line( + monkeypatch, tmp_path +): + monkeypatch.setattr(lifecycle_guard, "_MAX_LIFECYCLE_SCAN_LINES", 4) + script = tmp_path / "many.sh" + script.write_text("echo ok\n" * 10, encoding="utf-8") + + lexers = 0 + real_shlex = lifecycle_guard.shlex.shlex + + def counting(*args, **kwargs): + nonlocal lexers + lexers += 1 + return real_shlex(*args, **kwargs) + + monkeypatch.setattr(lifecycle_guard.shlex, "shlex", counting) + root = f"bash {script}" + assert guard(root) is True + # Only the one-line root was tokenized (a handful of lexers across the + # direct scans); the 10-line script never was. + assert 0 < lexers < 10 + + +# --- scheduler entry point -------------------------------------------------- + + +def test_check_gateway_lifecycle_shell_script_budget(monkeypatch, tmp_path): + monkeypatch.setattr(lifecycle_guard, "_MAX_LIFECYCLE_SCAN_BYTES", 8) + monkeypatch.setattr(lifecycle_guard, "_MAX_LIFECYCLE_SCAN_LINE_BYTES", 8) + script = tmp_path / "long-line.sh" + + script.write_text("x" * 7, encoding="utf-8") + lifecycle_guard.check_gateway_lifecycle("", str(script)) + + script.write_text("x" * 9, encoding="utf-8") + with pytest.raises(lifecycle_guard.GatewayLifecycleBlocked): + lifecycle_guard.check_gateway_lifecycle("", str(script)) + + +def test_check_gateway_lifecycle_python_path_charges_masker(monkeypatch, tmp_path): + """The .py branch's data-exemption masker tokenizes too, so it is budgeted + and fails closed before shlex on an over-budget line.""" + monkeypatch.setattr(lifecycle_guard, "_MAX_LIFECYCLE_SCAN_LINE_BYTES", 16) + + small = tmp_path / "small.py" + small.write_text("x = 1\n", encoding="utf-8") + lifecycle_guard.check_gateway_lifecycle("run report", str(small)) + + monkeypatch.setattr(lifecycle_guard.shlex, "shlex", _explode) + long_line = tmp_path / "long.py" + long_line.write_text("x = 1\n" + "y" * 40 + "\n", encoding="utf-8") + with pytest.raises(lifecycle_guard.GatewayLifecycleBlocked): + lifecycle_guard.check_gateway_lifecycle("run report", str(long_line)) + + +# --- no regression on realistic benign graphs ------------------------------ + + +def test_default_budget_admits_a_wide_benign_wrapper_graph(tmp_path): + """Issue #78398's shape: one wrapper invoking 200 small legitimate scripts + must still be allowed under the DEFAULT limits (an earlier fail-closed + attempt with a 64-path cap blocked exactly this).""" + children = [] + for i in range(200): + child = tmp_path / f"c{i}.sh" + child.write_text("echo step && ls -la /tmp\n" * 20, encoding="utf-8") + children.append(child) + hub = tmp_path / "hub.sh" + hub.write_text("".join(f"bash {c}\n" for c in children), encoding="utf-8") + + assert guard(f"bash {hub}") is False + + # ...and a lifecycle command hidden behind the 200 benign scripts is still + # found: the budget bounds work, it does not stop the walk early. + evil = tmp_path / "evil.sh" + evil.write_text("hermes gateway restart\n", encoding="utf-8") + hub.write_text(hub.read_text() + f"bash {evil}\n", encoding="utf-8") + assert guard(f"bash {hub}") is True diff --git a/tests/hermes_cli/test_gateway_restart_loop.py b/tests/hermes_cli/test_gateway_restart_loop.py index a005b9c2e8..6df661bc54 100644 --- a/tests/hermes_cli/test_gateway_restart_loop.py +++ b/tests/hermes_cli/test_gateway_restart_loop.py @@ -587,6 +587,33 @@ class TestTerminalToolGatewayLifecycleGuard: assert result["exit_code"] == 1 assert "KeepAlive" in result["error"] + def test_oversized_root_skips_launchctl_prescan_and_fails_closed( + self, monkeypatch + ): + """#78398: an over-budget root must never reach shlex — not even via + the launchctl pre-scan that runs before the full guard.""" + import cron.lifecycle_guard as lifecycle_guard + import tools.terminal_tool as tt + + self._patch_env(monkeypatch, self._make_fake_env(), inside_gateway=True) + monkeypatch.setattr( + lifecycle_guard, "_MAX_LIFECYCLE_SCAN_BYTES", 8, raising=False + ) + monkeypatch.setattr( + lifecycle_guard, "_MAX_LIFECYCLE_SCAN_LINE_BYTES", 8, raising=False + ) + + def explode_if_tokenized(*args, **kwargs): + raise AssertionError("over-budget root reached shlex") + + monkeypatch.setattr(lifecycle_guard.shlex, "shlex", explode_if_tokenized) + + result = json.loads(tt.terminal_tool(command="x" * 9)) + + assert result["exit_code"] == 1 + assert "command or referenced script" in result["error"] + assert "KeepAlive" not in result["error"] + @pytest.mark.parametrize("command", [ # Neutral, non-hermes label: label-independent detection is the point # (#62891 second reproduction used `ai.hermes.svc-reload-tmp`). diff --git a/tools/terminal_tool.py b/tools/terminal_tool.py index a5e2d2bfc0..eb375fd8d4 100644 --- a/tools/terminal_tool.py +++ b/tools/terminal_tool.py @@ -3146,8 +3146,15 @@ def terminal_tool( _MAX_REFERENCED_SCRIPT_BYTES, contains_gateway_lifecycle_command_or_referenced_script, contains_launchctl_submit_command, + lifecycle_scan_root_within_budget, ) - if contains_launchctl_submit_command(command): + # Keep the specific launchctl diagnostic when this optional + # pre-scan fits the budget. The full fail-closed guard below still + # runs when it does not, so oversized roots never reach shlex here. + if ( + lifecycle_scan_root_within_budget(command) + and contains_launchctl_submit_command(command) + ): return json.dumps({ "output": "", "exit_code": 1,