242 lines
8.9 KiB
Python
242 lines
8.9 KiB
Python
"""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
|