From 104ad84b6cd2e860b65f5e4980371afbf83204c6 Mon Sep 17 00:00:00 2001 From: Alpamys Date: Sun, 5 Jul 2026 11:16:19 +0500 Subject: [PATCH] fix(shrink): security-review findings (derived-path symlink re-validation, subprocess stderr control-strip) (v0.71.29) --- src/soup_cli/commands/shrink.py | 23 ++++++++++++++++- tests/test_v07129.py | 46 ++++++++++++++++++++++++++++++--- 2 files changed, 65 insertions(+), 4 deletions(-) diff --git a/src/soup_cli/commands/shrink.py b/src/soup_cli/commands/shrink.py index 1da1d3c..9aef5b8 100644 --- a/src/soup_cli/commands/shrink.py +++ b/src/soup_cli/commands/shrink.py @@ -52,6 +52,17 @@ _MAX_HEAL_EPOCHS = 100 _PPL_MAX_LENGTH = 512 +# Strip C0 / ESC / DEL before subprocess- or model-derived text hits the +# terminal (rich.markup.escape only neutralises [...] markup, not raw ESC bytes +# — mirrors commands/data_doctor.py::_for_terminal, v0.71.27). +_CONTROL_STRIP_TABLE = {i: None for i in range(0x20) if i not in (0x09, 0x0A, 0x0D)} +_CONTROL_STRIP_TABLE[0x7F] = None + + +def _for_terminal(text: str) -> str: + return text.translate(_CONTROL_STRIP_TABLE) + + # --------------------------------------------------------------------------- # Path + data helpers # --------------------------------------------------------------------------- @@ -369,6 +380,9 @@ def _shrink_impl( prune_model_layers(mdl, chosen.start, chosen.block_size) out_root = Path(output_dir) model_out = out_root / "model" + # Re-validate the DERIVED write path (a symlink could have been planted at + # /model since the top-of-command check on ). + enforce_under_cwd_and_no_symlink(str(model_out), "output model dir") model_out.mkdir(parents=True, exist_ok=True) mdl.save_pretrained(str(model_out)) tokenizer.save_pretrained(str(model_out)) @@ -381,6 +395,7 @@ def _shrink_impl( healed = False if heal is not None: adapter_dir = out_root / "heal_adapter" + enforce_under_cwd_and_no_symlink(str(adapter_dir), "heal adapter dir") console.print( f"[dim]Healing (distill original -> pruned, ~{heal_steps} steps) ...[/]" ) @@ -572,7 +587,10 @@ def _run_heal( f"heal distill exceeded {_HEAL_TIMEOUT_SECONDS}s timeout" ) from exc if result.returncode != 0: - tail = (result.stderr or b"").decode("utf-8", "replace")[-500:] + # Strip control bytes: the child's stderr is attacker-influenceable and + # reaches the terminal via the friendly error handler (escape() does not + # neutralise raw ESC/OSC sequences). + tail = _for_terminal((result.stderr or b"").decode("utf-8", "replace")[-500:]) raise RuntimeError(f"heal distill failed (rc={result.returncode}): {tail}") _fuse_adapter(base_dir=pruned_dir, adapter_dir=out_dir, trc=trc) @@ -594,6 +612,9 @@ def _fuse_adapter(*, base_dir: str, adapter_dir: str, trc: bool = False) -> None from peft import PeftModel from transformers import AutoModelForCausalLM, AutoTokenizer + # Re-validate the swap target right before mutating it (the heal subprocess + # ran for potentially hours — re-close the TOCTOU window on base_dir). + enforce_under_cwd_and_no_symlink(base_dir, "fuse base dir") base = AutoModelForCausalLM.from_pretrained( base_dir, trust_remote_code=trc, torch_dtype="auto" ) diff --git a/tests/test_v07129.py b/tests/test_v07129.py index 9617d5d..5d7563c 100644 --- a/tests/test_v07129.py +++ b/tests/test_v07129.py @@ -689,14 +689,15 @@ class TestReviewFixes: assert "loss == loss" not in src assert "math.isnan(loss)" in src - def test_fuse_adapter_produces_dense_model(self, tmp_path): - """_fuse_adapter merges a LoRA adapter back into the base in place, so - the shipped dir is a single dense model (no adapter_config.json).""" + def test_fuse_adapter_produces_dense_model(self, tmp_path, monkeypatch): + """_fuse_adapter merges a LoRA adapter back into the base (atomic swap), + so the shipped dir is a single dense model (no adapter_config.json).""" from peft import LoraConfig, TaskType, get_peft_model from transformers import AutoModelForCausalLM, AutoTokenizer from soup_cli.commands.shrink import _fuse_adapter + monkeypatch.chdir(tmp_path) # base_dir must stay under cwd (containment) base_dir = tmp_path / "base" tok = AutoTokenizer.from_pretrained("HuggingFaceTB/SmolLM2-135M-Instruct") _tiny_llama(4, vocab_size=len(tok)).save_pretrained(str(base_dir)) @@ -718,6 +719,45 @@ class TestReviewFixes: fused = AutoModelForCausalLM.from_pretrained(str(base_dir)) assert fused.config.num_hidden_layers == 4 + def test_output_model_symlink_rejected(self, tmp_path, monkeypatch): + """A symlink planted at /model is rejected before any write + (derived-path TOCTOU guard). POSIX-only (needs os.symlink).""" + import os as _os + + if not hasattr(_os, "symlink"): + pytest.skip("no os.symlink") + from typer.testing import CliRunner + + from soup_cli.cli import app + + monkeypatch.chdir(tmp_path) + model_dir = _write_tiny_model(tmp_path / "src_sym", layers=6) + calib = tmp_path / "calib.jsonl" + calib.write_text('{"text":"the quick brown fox"}\n', encoding="utf-8") + out_dir = tmp_path / "shrunk_sym" + out_dir.mkdir() + escape_target = tmp_path / "escape_target" + escape_target.mkdir() + try: + _os.symlink(str(escape_target), str(out_dir / "model"), + target_is_directory=True) + except (OSError, NotImplementedError): + pytest.skip("symlink creation not permitted on this platform") + r = CliRunner().invoke( + app, + ["shrink", "--model", model_dir, "--drop-layers", "2", + "--calib", "calib.jsonl", "--device", "cpu", + "--output-dir", str(out_dir), "--tolerance", "5.0"], + ) + assert r.exit_code != 0 + + def test_for_terminal_strips_control_bytes(self): + from soup_cli.commands.shrink import _for_terminal + + assert _for_terminal("a\x1b]0;evilbc") == "a]0;evilbc" + # tab / LF / CR preserved. + assert _for_terminal("a\tb\nc\rd") == "a\tb\nc\rd" + def test_dont_ship_exit_code_2(self, tmp_path, monkeypatch): """A genuine perplexity regression past tolerance exits 2 (DON'T SHIP).""" from typer.testing import CliRunner