Files
turnstone/tests/test_env_scrub.py
T
Patrick Buckley c339615e39 Bound search tool output against pathological inputs (#473)
* Bound search tool output against pathological inputs

Replaces the per-line truncation with a fully bounded pipeline so the
search tool can no longer overflow the LLM context — or OOM the parent —
on minified bundles, multi-GB JSONL records, or huge result sets.

Backend:
- Prefer ripgrep when on PATH; grep is the fallback. Detection is
  cached via functools.cache.
- ripgrep flags do most of the bounding natively: --max-columns 1024
  + --max-columns-preview, --max-filesize 10M, --max-count 100,
  --no-config, --no-messages, plus negative globs for the same
  noisy directories grep has been excluding.
- ripgrep added to the Dockerfile.

Streaming subprocess (_search_capture):
- subprocess.Popen with a streaming, byte-capped stdout read (4 MB).
  Defends against single-line files (training data, minified bundles)
  that would have OOM'd the previous subprocess.run capture.
- threading.Timer watchdog enforces tool_timeout even when the
  pipe read is blocked in the kernel — proc.wait(timeout=…) alone
  was insufficient because the read sat ahead of it.
- Stderr drained in a daemon thread to avoid pipe-deadlock when the
  child writes to stderr while we're still reading stdout. Cap on
  captured stderr keeps a hostile child from growing the buffer.

Tier-based formatter (_format_search_results):
- Tier 1: full path:line:content output, stream-emitted with a
  running-cost short-circuit so we never materialize past the budget.
- Tier 2: K samples per file with overflow notes; K is computed
  analytically from budget / file_count / avg-line-length so we hit
  the right ladder rung in a single pass.
- Tier 3: per-file counts only, also budget-bounded with a tail line
  reporting the omitted files. Sorted by descending count.
- Total output budget (32 KB) is well under tool_truncation, so the
  head+tail _truncate_output strategy never silently drops middle
  files in a search result.

Argument injection fix:
- The ripgrep arg list was missing the `--` separator that the grep
  branch already had. With auto_approve on the search tool, that was
  exploitable: path='--pre=COMMAND' would have made ripgrep run the
  script as a per-file preprocessor and surface its stdout. Added
  `--` and a regression test.

State-machine cleanup in _exec_search:
- rc < 0 (signal-killed by something other than us) now surfaces a
  dedicated 'killed by signal N' message instead of being parsed as
  success.
- capped + zero parsed records (e.g. one multi-MB line with no \n)
  now returns a dedicated byte-cap message instead of the malformed-
  output message that previously masked the real cause.
- _report_tool_result descriptions now match the returned payload
  (no more 'no matches' tag on a 'malformed' payload).

Defence-in-depth on env scrub:
- RIPGREP_CONFIG_PATH, GIT_CONFIG, GIT_CONFIG_GLOBAL, GIT_CONFIG_SYSTEM
  added to _EXPLICIT_SCRUB. We pass --no-config on the rg CLI today,
  but if a future caller forgets the flag, an attacker who can set
  one of these env vars could plant a config containing --pre=… and
  recreate the same RCE shape.

Tests:
- TestSearchLineTruncation rewritten to mock _search_capture instead
  of subprocess.run (the previous tests passed ChatSession kwargs
  that no longer satisfy the constructor).
- TestSearchBackendSelection covers rg/grep detection and arg
  construction, including the --pre flag-injection regression.
- TestSearchOutputBudget exercises Tier 1/2/3 directly.
- TestSearchCaptureStreaming spawns real Python subprocess writers
  to exercise the byte-cap trim, mega-line-no-newline edge case, the
  watchdog timeout when the child writes nothing, and the stderr
  drain under load.
- test_env_scrub picks up the new tool-config keys.

* Address Copilot review on #473

- Budget the Tier 2/3 header up front so the formatter's emission stays
  strictly within _SEARCH_OUTPUT_BUDGET. Previously the fit checks only
  counted body bytes, letting the final string overflow by ~120 chars
  (header + separator) and triggering _truncate_output's head+tail
  dropout — exactly the shape this code was trying to avoid.
- Restore the (5, 3, 1) ladder in Tier 2: the analytical K from perf-2
  is kept as a starting estimate, but if that K's actual emission
  doesn't fit (the estimate ignores the header and overweights shared-
  path compression) we step down through the ladder before falling
  through to Tier 3. The previous one-shot K could collapse to counts-
  only when 3/file or 1/file would have fit.
- Only normalise rc to 0 in the capped-output path when rc < 0 (our
  SIGKILL). There's a narrow race where the child can exit naturally
  between our read and our kill; preserving a non-negative rc means
  rg's rc=2 ('matches found but some files had errors') no longer
  silently turns into a clean success when the byte cap also fires.
- Clarify _MAX_SEARCH_LINE_LENGTH doc: the cap applies to the content
  portion (after path:lineno:), not the whole emitted line.
- Add explanatory comments on the two intentional `except Exception:
  pass` blocks in _search_capture (stderr drain, pipe close in the
  cleanup finally) so static analysis and future readers can see the
  silence is deliberate.
- Tighten the budget tests: now assert strict `<= _SEARCH_OUTPUT_BUDGET`
  instead of the +512-char slack that was masking the header overflow.
- New regression tests:
  - Tier 2 ladder step-down (K=5 over budget, K=3 fits, no Tier 3 fall-through)
  - capped + rc=2 surfaces stderr instead of being normalised to success
  - capped + rc<0 (our SIGKILL) flows through as a partial-result success

* chore(search): post-review cleanup

Follow-up to the Copilot-review fixes in 39d2aa2 — these are all small
quality items (no behaviour change, no new tests).

- q-1: collapse the Tier 2 candidates filter to a single expression.
  Drops the redundant inner ``max(estimated_k, 1)`` and the unreachable
  ``if not candidates`` branch (the ladder ends in 1 and ``estimated_k``
  is already floored at 1, so the comprehension always yields ≥ ``[1]``).
  ``or [...]`` is kept as defence against future ladder changes.
- q-2: update _format_search_results docstring to match the new ladder
  semantics (analytical seed → step down through (5, 3, 1) from the
  highest rung ≤ the estimate). The previous wording suggested every
  Tier 2 attempt started at 5.
- q-3: combine the two ``from turnstone.core.session import ...``
  statements in test_tier2_steps_down_ladder_before_falling_to_tier3
  into a single top-of-function import (matches the surrounding tests).
- q-4: shorten the explanatory comments on the two best-effort cleanup
  paths in _search_capture to one line each. Both sites now read with
  the same shape ("# best-effort: pipe may be torn down by ...").
- q-5: trim the _MAX_SEARCH_LINE_LENGTH comment from 7 lines back to 3.
  Keeps the load-bearing semantic (cap is on the content portion only)
  and the pathological-line defence; drops the paths-aren't-bounded
  parenthetical, which was background reading rather than WHY.
2026-05-04 00:46:43 -07:00

160 lines
5.6 KiB
Python

"""Tests for turnstone.core.env — subprocess environment scrubbing."""
from __future__ import annotations
import os
from unittest.mock import patch
from turnstone.core.env import _is_safe, _is_secret, scrubbed_env
class TestIsSecret:
def test_explicit_scrub_list(self):
assert _is_secret("OPENAI_API_KEY") is True
assert _is_secret("ANTHROPIC_API_KEY") is True
assert _is_secret("TURNSTONE_JWT_SECRET") is True
assert _is_secret("AWS_SECRET_ACCESS_KEY") is True
def test_tool_config_paths_scrubbed(self):
"""Tool-config env vars whose target files load executable
directives must be scrubbed even though they don't match a
secret-suffix pattern. Defence-in-depth alongside on-CLI
``--no-config`` for ripgrep and friends."""
assert _is_secret("RIPGREP_CONFIG_PATH") is True
assert _is_secret("GIT_CONFIG") is True
assert _is_secret("GIT_CONFIG_GLOBAL") is True
assert _is_secret("GIT_CONFIG_SYSTEM") is True
def test_suffix_matching(self):
assert _is_secret("MY_CUSTOM_API_KEY") is True
assert _is_secret("DB_PASSWORD") is True
assert _is_secret("AUTH_TOKEN") is True
assert _is_secret("SERVICE_CREDENTIAL") is True
assert _is_secret("GCP_CREDENTIALS") is True
def test_safe_vars_not_secret(self):
assert _is_secret("PATH") is False
assert _is_secret("HOME") is False
assert _is_secret("LANG") is False
def test_no_false_positives_on_substring(self):
"""Suffix matching avoids false positives like MONKEYTYPE."""
assert _is_secret("MONKEYTYPE") is False
assert _is_secret("KEYBOARD_LAYOUT") is False
assert _is_secret("PYTHONPATH") is False
assert _is_secret("EDITOR") is False
assert _is_secret("GOPATH") is False
class TestIsSafe:
def test_safe_names(self):
assert _is_safe("PATH") is True
assert _is_safe("HOME") is True
assert _is_safe("TERM") is True
assert _is_safe("MANWIDTH") is True
def test_safe_prefixes(self):
assert _is_safe("LC_ALL") is True
assert _is_safe("LC_CTYPE") is True
assert _is_safe("XDG_RUNTIME_DIR") is True
def test_non_safe_names(self):
assert _is_safe("OPENAI_API_KEY") is False
assert _is_safe("CUSTOM_VAR") is False
class TestScrubbedEnv:
def test_strips_api_keys(self):
fake_env = {
"PATH": "/usr/bin",
"HOME": "/home/user",
"OPENAI_API_KEY": "sk-secret",
"ANTHROPIC_API_KEY": "ant-secret",
"CUSTOM_VAR": "safe_value",
}
with patch.dict(os.environ, fake_env, clear=True):
result = scrubbed_env()
assert result["PATH"] == "/usr/bin"
assert result["HOME"] == "/home/user"
assert result["CUSTOM_VAR"] == "safe_value"
assert "OPENAI_API_KEY" not in result
assert "ANTHROPIC_API_KEY" not in result
def test_strips_pattern_matched_secrets(self):
fake_env = {
"PATH": "/usr/bin",
"MY_SERVICE_TOKEN": "tok-123",
"DB_PASSWORD": "pass123",
}
with patch.dict(os.environ, fake_env, clear=True):
result = scrubbed_env()
assert "MY_SERVICE_TOKEN" not in result
assert "DB_PASSWORD" not in result
def test_extra_vars_merged(self):
fake_env = {"PATH": "/usr/bin"}
with patch.dict(os.environ, fake_env, clear=True):
result = scrubbed_env(extra={"MANWIDTH": "80"})
assert result["MANWIDTH"] == "80"
assert result["PATH"] == "/usr/bin"
def test_passthrough_overrides_scrub(self):
fake_env = {
"PATH": "/usr/bin",
"OPENAI_API_KEY": "sk-needed",
}
with patch.dict(os.environ, fake_env, clear=True):
result = scrubbed_env(passthrough=["OPENAI_API_KEY"])
assert result["OPENAI_API_KEY"] == "sk-needed"
def test_preserves_locale_vars(self):
fake_env = {
"PATH": "/usr/bin",
"LC_ALL": "en_US.UTF-8",
"LC_CTYPE": "en_US.UTF-8",
}
with patch.dict(os.environ, fake_env, clear=True):
result = scrubbed_env()
assert result["LC_ALL"] == "en_US.UTF-8"
assert result["LC_CTYPE"] == "en_US.UTF-8"
def test_preserves_unknown_non_secret_vars(self):
fake_env = {
"PATH": "/usr/bin",
"PYTHONPATH": "/opt/lib",
"GOPATH": "/home/user/go",
}
with patch.dict(os.environ, fake_env, clear=True):
result = scrubbed_env()
assert result["PYTHONPATH"] == "/opt/lib"
assert result["GOPATH"] == "/home/user/go"
def test_extra_can_reintroduce_scrubbed_var(self):
"""extra= intentionally overrides scrubbing (operator-controlled)."""
fake_env = {"PATH": "/usr/bin", "OPENAI_API_KEY": "sk-original"}
with patch.dict(os.environ, fake_env, clear=True):
result = scrubbed_env(extra={"OPENAI_API_KEY": "sk-injected"})
assert result["OPENAI_API_KEY"] == "sk-injected"
def test_less_prefix_does_not_leak_secrets(self):
"""LESS pager vars are safe but LESS_SECRET_TOKEN is not."""
fake_env = {
"PATH": "/usr/bin",
"LESS": "-R",
"LESSOPEN": "| lesspipe %s",
"LESS_SECRET_TOKEN": "tok-secret",
}
with patch.dict(os.environ, fake_env, clear=True):
result = scrubbed_env()
assert result["LESS"] == "-R"
assert result["LESSOPEN"] == "| lesspipe %s"
assert "LESS_SECRET_TOKEN" not in result