From bd58656bdf72ddd00a678fcb9f5b2168db31598d Mon Sep 17 00:00:00 2001 From: IshanKulkarni02 Date: Mon, 26 Oct 2026 03:12:09 +0530 Subject: [PATCH] fix(mcp): address review findings on the MCP server - Validate --run-name as a single path segment so it cannot escape strix_runs/ via `..`, an absolute path, or a leading dot (P1). - Cap read_tool_output responses at _MAX_RESULT_CHARS and seek the requested window off disk instead of loading the whole file (P2). - When spilling an oversized result fails, return a head+tail preview instead of silently dropping the tail (P2). Co-Authored-By: Claude Opus 4.8 --- strix/interface/mcp_server.py | 72 ++++++++++++++++++++++++++++------- tests/test_mcp_server.py | 44 +++++++++++++++++++++ 2 files changed, 102 insertions(+), 14 deletions(-) diff --git a/strix/interface/mcp_server.py b/strix/interface/mcp_server.py index 6ce7feb54..f81cb6a13 100644 --- a/strix/interface/mcp_server.py +++ b/strix/interface/mcp_server.py @@ -174,6 +174,26 @@ PROXY_TOOL_NAMES = frozenset( ) +def _validate_run_name(run_name: str) -> str: + """Return ``run_name`` if it is a single safe path segment, else raise. + + ``run_name`` becomes a directory under ``strix_runs/``. A value with a path + separator, a ``..`` segment, an absolute path, or a leading dot could place + run files outside that tree or resume an unrelated run, so reject anything + that is not a plain name. + """ + name = run_name.strip() + if not name: + raise ValueError("--run-name must not be empty") + if name in {".", ".."} or name.startswith("."): + raise ValueError(f"--run-name must not start with a dot: {run_name!r}") + if "/" in name or "\\" in name or Path(name).is_absolute(): + raise ValueError(f"--run-name must be a single path segment: {run_name!r}") + if any(ord(ch) < 0x20 or ord(ch) == 0x7F for ch in name): + raise ValueError("--run-name must not contain control characters") + return name + + def _init_run_state(run_name: str) -> Path: """Wire up host-side run state so the reporting, threat-model, and coverage tools persist to a run directory and survive a server restart, mirroring the @@ -335,11 +355,16 @@ def _read_spilled_output(output_id: str, offset: int, length: int) -> str: if not path.is_file(): return f"No saved output named {output_id!r}." offset = max(offset, 0) - length = length if length > 0 else _MAX_RESULT_CHARS - data = path.read_text(encoding="utf-8", errors="replace") - chunk = data[offset : offset + length] - end = offset + len(chunk) - return f"{chunk}\n\n[chars {offset}-{end} of {len(data)}; more available: {end < len(data)}]" + # Cap the window so one call can't return the whole file, and read only that + # window off disk (seek) rather than loading the whole file into memory. + length = min(length, _MAX_RESULT_CHARS) if length > 0 else _MAX_RESULT_CHARS + size = path.stat().st_size + with path.open("rb") as handle: + handle.seek(offset) + raw = handle.read(length) + chunk = raw.decode("utf-8", errors="replace") + end = offset + len(raw) + return f"{chunk}\n\n[bytes {offset}-{end} of {size}; more available: {end < size}]" async def _invoke(tool: FunctionTool, arguments: dict[str, Any], sandbox: _SandboxTools) -> str: @@ -360,13 +385,21 @@ async def _invoke(tool: FunctionTool, arguments: dict[str, Any], sandbox: _Sandb if len(text) <= _MAX_RESULT_CHARS: return text output_id = _spill_large_output(tool.name, text) - how = ( - f"; call {_READ_OUTPUT_TOOL} with output_id={output_id!r} to read the rest" - if output_id - else "" + if output_id: + preview = text[:_MAX_RESULT_CHARS] + return ( + f"{preview}\n\n[truncated to {_MAX_RESULT_CHARS} of {len(text)} chars; " + f"call {_READ_OUTPUT_TOOL} with output_id={output_id!r} to read the rest]" + ) + # The save failed (no run dir, disk full, unwritable), so there is nothing to + # page back. Show both ends rather than silently dropping the tail. + half = _MAX_RESULT_CHARS // 2 + head = text[:half] + tail = text[-half:] + return ( + f"{head}\n\n[could not save the full output ({len(text)} chars); the middle " + f"is omitted — showing the first and last {half} chars]\n\n{tail}" ) - preview = text[:_MAX_RESULT_CHARS] - return f"{preview}\n\n[truncated to {_MAX_RESULT_CHARS} of {len(text)} chars{how}]" def _server_instructions(sandbox: _SandboxTools) -> str: @@ -398,14 +431,19 @@ def _build_server(sandbox: _SandboxTools) -> Any: description=( "Read back a tool result that was truncated, using the output_id " "from the truncation notice. Returns a bounded chunk; page with " - "offset and length." + f"byte offset and length (length is capped at {_MAX_RESULT_CHARS})." ), inputSchema={ "type": "object", "properties": { "output_id": {"type": "string", "description": "Id from the truncation notice."}, "offset": {"type": "integer", "minimum": 0, "default": 0}, - "length": {"type": "integer", "minimum": 1, "default": _MAX_RESULT_CHARS}, + "length": { + "type": "integer", + "minimum": 1, + "maximum": _MAX_RESULT_CHARS, + "default": _MAX_RESULT_CHARS, + }, }, "required": ["output_id"], }, @@ -533,7 +571,13 @@ def run_mcp_server(argv: list[str]) -> int: # A unique default keeps two servers launched from the same directory from # sharing one run and overwriting each other's reports; pass --run-name to # deliberately resume or share a run. - run_name = args.run_name or f"mcp-{uuid.uuid4().hex[:8]}" + if args.run_name is not None: + try: + run_name = _validate_run_name(args.run_name) + except ValueError as exc: + parser.error(str(exc)) + else: + run_name = f"mcp-{uuid.uuid4().hex[:8]}" try: asyncio.run(_serve(local_sources, run_name)) diff --git a/tests/test_mcp_server.py b/tests/test_mcp_server.py index 78b7300e5..624ccbd42 100644 --- a/tests/test_mcp_server.py +++ b/tests/test_mcp_server.py @@ -278,6 +278,50 @@ def test_read_spilled_output_rejects_traversal() -> None: assert "Invalid output id" in mcp_server._read_spilled_output("a/b.txt", 0, 100) +def test_validate_run_name_rejects_traversal_and_separators() -> None: + # A run name becomes a directory under strix_runs/; path parts must not escape. + for bad in ["../evil", "/abs/path", "a/b", "a\\b", "..", ".", ".hidden", ""]: + with pytest.raises(ValueError, match="run-name"): + mcp_server._validate_run_name(bad) + assert mcp_server._validate_run_name(" my-run ") == "my-run" + + +@pytest.mark.asyncio +async def test_read_spilled_output_caps_length_and_pages(tmp_path: Path, monkeypatch) -> None: # type: ignore[no-untyped-def] + # An over-cap length request is clamped, and reads return only the window. + monkeypatch.chdir(tmp_path) + mcp_server._init_run_state("mcp-cap") + try: + body = "y" * (mcp_server._MAX_RESULT_CHARS * 2) + output_id = mcp_server._spill_large_output("demo", body) + assert output_id + over = mcp_server._read_spilled_output(output_id, 0, mcp_server._MAX_RESULT_CHARS * 10) + # Response body (minus the trailing notice line) never exceeds the cap. + chunk = over.split("\n\n[", 1)[0] + assert len(chunk) == mcp_server._MAX_RESULT_CHARS + assert "more available: True" in over + finally: + set_global_report_state(None) + + +@pytest.mark.asyncio +async def test_invoke_keeps_both_ends_when_save_fails() -> None: + # With no run state, the spill save fails; the tail must not be dropped. + set_global_report_state(None) + big = "A" * (mcp_server._MAX_RESULT_CHARS + 1000) + "ZZZEND" + + class _BigTool: + name = "big" + + async def on_invoke_tool(self, _ctx, _args): # type: ignore[no-untyped-def] + return big + + out = await mcp_server._invoke(_BigTool(), {}, mcp_server._SandboxTools()) # type: ignore[arg-type] + assert "could not save" in out + assert out.startswith("A") + assert out.rstrip().endswith("ZZZEND") # tail preserved + + @pytest.mark.asyncio async def test_read_tool_output_is_advertised_and_callable(tmp_path: Path, monkeypatch) -> None: # type: ignore[no-untyped-def] monkeypatch.chdir(tmp_path)