fix(shrink): security-review findings (derived-path symlink re-validation, subprocess stderr control-strip) (v0.71.29)

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
+    # <output_dir>/model since the top-of-command check on <output_dir>).
+    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"
     )

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 <output_dir>/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