diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 0dd12c7..2f4ee4c 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -79,6 +79,26 @@ jobs: with: context: . push: false + load: true tags: local-llm-router:${{ github.sha }} cache-from: type=gha cache-to: type=gha,mode=max + - name: Scan the image for known vulnerabilities + # A release image may not carry a critical or high vulnerability that + # has a fix. This job gates the merge, and only merges are released. + uses: aquasecurity/trivy-action@v0.36.0 + with: + image-ref: local-llm-router:${{ github.sha }} + scanners: vuln + severity: CRITICAL,HIGH + ignore-unfixed: true + exit-code: "1" + - name: Start the image and wait for it to answer + run: | + docker run -d --name smoke -p 8000:8000 -e ROUTER_API_KEYS=smoke local-llm-router:${{ github.sha }} + for attempt in $(seq 1 20); do + if curl -fsS http://127.0.0.1:8000/healthz; then exit 0; fi + sleep 1 + done + docker logs smoke + exit 1 diff --git a/Dockerfile b/Dockerfile index 6bd0c24..21681d5 100644 --- a/Dockerfile +++ b/Dockerfile @@ -11,9 +11,18 @@ ENV PYTHONDONTWRITEBYTECODE=1 \ PYTHONUNBUFFERED=1 \ ROUTER_ENVIRONMENT=production -RUN useradd --create-home --uid 10001 appuser +# Take the distribution's security fixes; the base image tag lags behind them. +RUN apt-get update \ + && apt-get upgrade -y --no-install-recommends \ + && rm -rf /var/lib/apt/lists/* \ + && useradd --create-home --uid 10001 appuser COPY --from=builder /wheels /wheels -RUN python -m pip install --no-cache-dir /wheels/* && rm -rf /wheels +# The running service never installs anything, so the installer and the +# libraries it vendors are removed rather than patched. +RUN python -m pip install --no-cache-dir /wheels/* \ + && rm -rf /wheels \ + && python -m pip uninstall -y pip \ + && rm -rf /usr/local/lib/python3.13/ensurepip WORKDIR /app COPY config ./config diff --git a/README.md b/README.md index 0a37414..a860cb9 100644 --- a/README.md +++ b/README.md @@ -344,6 +344,40 @@ step, which stays disabled until a destination is configured. The tests run agai on SQLite; the server deployment in [`deploy/kubernetes/mlflow.yaml`](deploy/kubernetes/mlflow.yaml) has not been run on a cluster. +## Artifact scanning + +Nothing is released without being scanned. + +**Container.** CI scans the release image with [Trivy](https://trivy.dev/) and fails on any +critical or high vulnerability that has a fix, then starts the image and waits for `/healthz`. +The image takes the distribution's security updates at build time and ships without `pip`, so +the installer and the libraries it vendors are not there to be vulnerable. + +**Models and adapters.** Loading weights can run code: a pickle checkpoint executes what it +names, and a repository that ships Python asks the loader to import it. `llm_router.artifact_scan` +inspects an artifact directory without loading anything from it. + +```bash +python -m llm_router.artifact_scan ./artifacts/general-local --subject general-local +python -m llm_router.artifact_scan ./artifacts/claims-lora --kind adapter --expect sha256:... +``` + +| Rule | Refused | +|---|---| +| `pickle-format` | Any pickle-capable file (`.bin`, `.pt`, `.ckpt`, `.pkl`, ...), or a pickle under any other name. The finding names what it would import. | +| `executable-code` | Python, shell, or native libraries inside the artifact. | +| `remote-code` | A configuration with `auto_map` or `trust_remote_code`, or one that cannot be read. | +| `invalid-safetensors` | A header that does not account for exactly the bytes in the file. | +| `no-weights` | An artifact with no safetensors weights. | +| `symlink` | A link pointing outside the artifact. | +| `adapter-config` | An adapter that is not LoRA, or whose rank the engine would not load. | +| `checksum-mismatch` | A digest other than the one the catalog records. | + +The digest covers every file's path and content, so it is the value to record as a checksum in +a deployment revision. The command exits 1 on any finding. It checks formats and provenance; it +does not judge what a model has learned. The catalog's checksums are placeholders, so no real +artifact has been scanned against them yet. + ## Failure behaviour | Condition | What the gateway does | diff --git a/src/llm_router/artifact_scan.py b/src/llm_router/artifact_scan.py new file mode 100644 index 0000000..b2c1784 --- /dev/null +++ b/src/llm_router/artifact_scan.py @@ -0,0 +1,315 @@ +"""Model and adapter artifact scanning before release (section 14). + +Loading model weights can run code. A pickle-based checkpoint executes +whatever it names when it is opened, and a repository that ships its own +Python asks the loader to import it. This module inspects an artifact +directory without loading anything from it and refuses both, along with +weights whose safetensors header does not describe the bytes that follow and +an artifact whose digest is not the one the catalog recorded. +""" + +import hashlib +import io +import json +import pickletools +import struct +import zipfile +from collections.abc import Iterator, Sequence +from enum import StrEnum +from pathlib import Path +from typing import Any + +from pydantic import BaseModel + +# Formats that deserialize through pickle, or can. +PICKLE_SUFFIXES = frozenset( + {".bin", ".pt", ".pth", ".pkl", ".pickle", ".ckpt", ".joblib", ".npy", ".npz"} +) +CODE_SUFFIXES = frozenset({".py", ".pyc", ".sh", ".so", ".dll", ".dylib", ".exe"}) +CONFIG_FILES = frozenset({"config.json", "adapter_config.json", "tokenizer_config.json"}) +MAX_HEADER_BYTES = 100 * 1024 * 1024 +# The serving configuration loads adapters up to this rank and no further. +MAX_LORA_RANK = 32 +SAFETENSORS_DTYPES = frozenset( + {"BOOL", "U8", "I8", "I16", "U16", "F16", "BF16", "I32", "U32", "F32", "F64", "I64", "U64"} +) + + +class ArtifactKind(StrEnum): + MODEL = "model" + ADAPTER = "adapter" + + +class Rule(StrEnum): + PICKLE = "pickle-format" + CODE = "executable-code" + REMOTE_CODE = "remote-code" + SAFETENSORS = "invalid-safetensors" + NO_WEIGHTS = "no-weights" + SYMLINK = "symlink" + ADAPTER_CONFIG = "adapter-config" + CHECKSUM = "checksum-mismatch" + + +class Finding(BaseModel): + rule: Rule + path: str + detail: str + + +class ScanReport(BaseModel): + root: str + kind: ArtifactKind + digest: str + files: int + findings: tuple[Finding, ...] + + @property + def passed(self) -> bool: + return not self.findings + + +def _files(root: Path) -> Iterator[Path]: + yield from sorted(path for path in root.rglob("*") if path.is_file() or path.is_symlink()) + + +def artifact_digest(root: Path) -> str: + """One digest for a whole artifact: every file's path and content, in order.""" + + outer = hashlib.sha256() + for path in _files(root): + if path.is_symlink(): + continue + inner = hashlib.sha256() + with path.open("rb") as handle: + for block in iter(lambda: handle.read(1 << 20), b""): + inner.update(block) + outer.update(f"{path.relative_to(root).as_posix()}\0{inner.hexdigest()}\n".encode()) + return f"sha256:{outer.hexdigest()}" + + +def pickle_imports(data: bytes) -> tuple[str, ...]: + """The callables a pickle would import, read from its opcodes without running it.""" + + imports: list[str] = [] + strings: list[str] = [] + try: + for opcode, argument, _ in pickletools.genops(io.BytesIO(data)): + if opcode.name == "GLOBAL" and isinstance(argument, str): + imports.append(argument.replace(" ", ".")) + elif opcode.name == "STACK_GLOBAL" and len(strings) >= 2: + imports.append(f"{strings[-2]}.{strings[-1]}") + elif isinstance(argument, str): + strings.append(argument) + except Exception: + # A truncated or hostile stream; what was read so far is still reported. + pass + return tuple(dict.fromkeys(imports)) + + +def _pickle_payloads(path: Path) -> list[bytes]: + """Pickle streams in a file, whatever it is named.""" + + with path.open("rb") as handle: + head = handle.read(4) + if head[:2] == b"PK": + try: + with zipfile.ZipFile(path) as archive: + return [archive.read(name) for name in archive.namelist() if name.endswith(".pkl")] + except zipfile.BadZipFile: + return [] + if len(head) >= 2 and head[0] == 0x80 and 2 <= head[1] <= 5: + return [path.read_bytes()] + return [] + + +def _safetensors_problem(path: Path) -> str | None: + """Why a safetensors file is not what its header claims, or None if it is sound.""" + + size = path.stat().st_size + with path.open("rb") as handle: + prefix = handle.read(8) + if len(prefix) < 8: + return "shorter than the 8-byte header length" + (header_length,) = struct.unpack(" MAX_HEADER_BYTES or header_length > size - 8: + return f"header length {header_length} does not fit the file" + raw = handle.read(header_length) + try: + header = json.loads(raw) + except ValueError: + return "header is not JSON" + if not isinstance(header, dict): + return "header is not an object" + data_size = size - 8 - header_length + spans: list[tuple[int, int]] = [] + for name, tensor in header.items(): + if name == "__metadata__": + continue + if not isinstance(tensor, dict) or tensor.get("dtype") not in SAFETENSORS_DTYPES: + return f"tensor {name} has no known dtype" + offsets = tensor.get("data_offsets") + if ( + not isinstance(offsets, list) + or len(offsets) != 2 + or not all(isinstance(item, int) for item in offsets) + or not 0 <= offsets[0] <= offsets[1] <= data_size + ): + return f"tensor {name} points outside the file" + spans.append((offsets[0], offsets[1])) + position = 0 + for begin, end in sorted(spans): + if begin != position: + return "tensor data has a gap or an overlap" + position = end + if position != data_size: + return "file holds bytes no tensor accounts for" + return None + + +def _config_findings(path: Path, relative: str, kind: ArtifactKind) -> list[Finding]: + try: + config: Any = json.loads(path.read_text(encoding="utf-8")) + except (ValueError, UnicodeDecodeError): + return [Finding(rule=Rule.REMOTE_CODE, path=relative, detail="configuration is not JSON")] + if not isinstance(config, dict): + return [ + Finding(rule=Rule.REMOTE_CODE, path=relative, detail="configuration is not an object") + ] + findings: list[Finding] = [] + if "auto_map" in config or config.get("trust_remote_code"): + findings.append( + Finding( + rule=Rule.REMOTE_CODE, + path=relative, + detail="asks the loader to import code shipped with the artifact", + ) + ) + if path.name == "adapter_config.json" and kind is ArtifactKind.ADAPTER: + if str(config.get("peft_type", "")).upper() != "LORA": + findings.append( + Finding( + rule=Rule.ADAPTER_CONFIG, + path=relative, + detail=f"peft_type is {config.get('peft_type')!r}, only LORA is served", + ) + ) + rank = config.get("r") + if not isinstance(rank, int) or not 1 <= rank <= MAX_LORA_RANK: + findings.append( + Finding( + rule=Rule.ADAPTER_CONFIG, + path=relative, + detail=f"rank {rank!r} is outside 1..{MAX_LORA_RANK}", + ) + ) + return findings + + +def scan( + root: Path, kind: ArtifactKind = ArtifactKind.MODEL, expected_digest: str | None = None +) -> ScanReport: + """Inspect an artifact directory and report everything that blocks its release.""" + + resolved = root.resolve() + findings: list[Finding] = [] + weights = 0 + count = 0 + for path in _files(root): + count += 1 + relative = path.relative_to(root).as_posix() + if path.is_symlink(): + if not path.resolve().is_relative_to(resolved): + findings.append( + Finding(rule=Rule.SYMLINK, path=relative, detail="links outside the artifact") + ) + continue + suffix = path.suffix.lower() + if suffix in CODE_SUFFIXES: + findings.append( + Finding(rule=Rule.CODE, path=relative, detail="artifacts may not ship code") + ) + continue + payloads = _pickle_payloads(path) + if payloads or suffix in PICKLE_SUFFIXES: + imports = [name for payload in payloads for name in pickle_imports(payload)] + detail = "weights must be safetensors; this format runs code when loaded" + if imports: + detail += f" (imports {', '.join(sorted(set(imports))[:8])})" + findings.append(Finding(rule=Rule.PICKLE, path=relative, detail=detail)) + continue + if suffix == ".safetensors": + weights += 1 + problem = _safetensors_problem(path) + if problem: + findings.append(Finding(rule=Rule.SAFETENSORS, path=relative, detail=problem)) + elif path.name in CONFIG_FILES: + findings.extend(_config_findings(path, relative, kind)) + if not weights: + findings.append(Finding(rule=Rule.NO_WEIGHTS, path=".", detail="no safetensors weights")) + if kind is ArtifactKind.ADAPTER and not (root / "adapter_config.json").is_file(): + findings.append( + Finding(rule=Rule.ADAPTER_CONFIG, path=".", detail="adapter_config.json is missing") + ) + digest = artifact_digest(root) + if expected_digest and digest != expected_digest: + findings.append( + Finding( + rule=Rule.CHECKSUM, + path=".", + detail=f"digest is {digest}, the catalog records {expected_digest}", + ) + ) + return ScanReport( + root=root.as_posix(), kind=kind, digest=digest, files=count, findings=tuple(findings) + ) + + +def main(argv: Sequence[str] | None = None) -> int: + """Scan a model or adapter directory; exit 1 if anything blocks its release. + + With --subject the digest is checked against the checksum the catalog + records for that model or adapter. With --expect it is checked against + the digest given. + """ + + import argparse + + from llm_router.governance import SubjectKind, governed_versions + from llm_router.registry import load_registry + + parser = argparse.ArgumentParser(description=main.__doc__) + parser.add_argument("path") + parser.add_argument("--kind", choices=[item.value for item in ArtifactKind]) + parser.add_argument("--subject", help="catalog model or adapter this artifact is") + parser.add_argument("--expect", help="digest the artifact must have") + parser.add_argument("--catalog", default="config/registry.yaml") + arguments = parser.parse_args(argv) + + root = Path(arguments.path) + if not root.is_dir(): + parser.error(f"{root} is not a directory") + kind = ArtifactKind(arguments.kind) if arguments.kind else ArtifactKind.MODEL + expected = arguments.expect + if arguments.subject: + subject = next( + ( + item + for item in governed_versions(load_registry(arguments.catalog)) + if item.name == arguments.subject + ), + None, + ) + if subject is None: + parser.error(f"{arguments.subject} is not in the catalog") + if not arguments.kind and subject.kind is SubjectKind.ADAPTER: + kind = ArtifactKind.ADAPTER + expected = expected or subject.checksum + + report = scan(root, kind, expected) + print(json.dumps({**report.model_dump(mode="json"), "passed": report.passed}, indent=2)) + return 0 if report.passed else 1 + + +if __name__ == "__main__": # pragma: no cover - command-line entry point + raise SystemExit(main()) diff --git a/tests/unit/test_artifact_scan.py b/tests/unit/test_artifact_scan.py new file mode 100644 index 0000000..f424f6d --- /dev/null +++ b/tests/unit/test_artifact_scan.py @@ -0,0 +1,237 @@ +import json +import os +import pickle +import struct +import zipfile +from pathlib import Path +from typing import Any + +import pytest + +from llm_router.artifact_scan import ( + ArtifactKind, + Rule, + artifact_digest, + main, + pickle_imports, + scan, +) + + +def safetensors(tensors: dict[str, tuple[int, int]], data_size: int, **extra: Any) -> bytes: + header: dict[str, Any] = {"__metadata__": {"format": "pt"}, **extra} + for name, (begin, end) in tensors.items(): + header[name] = {"dtype": "F32", "shape": [(end - begin) // 4], "data_offsets": [begin, end]} + encoded = json.dumps(header).encode() + return struct.pack(" Path: + root.mkdir(parents=True, exist_ok=True) + (root / "model.safetensors").write_bytes(safetensors({"a": (0, 16), "b": (16, 48)}, 48)) + (root / "config.json").write_text(json.dumps({"model_type": "llama"}), encoding="utf-8") + (root / "tokenizer.json").write_text("{}", encoding="utf-8") + return root + + +def write_adapter(root: Path, **config: Any) -> Path: + root.mkdir(parents=True, exist_ok=True) + (root / "adapter_model.safetensors").write_bytes(safetensors({"lora_A": (0, 32)}, 32)) + (root / "adapter_config.json").write_text( + json.dumps({"peft_type": "LORA", "r": 16, **config}), encoding="utf-8" + ) + return root + + +class Hostile: + """Pickles to a call of os.system. It is only ever serialized, never loaded.""" + + def __reduce__(self) -> tuple[Any, tuple[str]]: + return (os.system, ("echo should-never-run",)) + + +def rules(root: Path, kind: ArtifactKind = ArtifactKind.MODEL) -> list[Rule]: + return [finding.rule for finding in scan(root, kind).findings] + + +def test_a_sound_model_and_a_sound_adapter_pass(tmp_path: Path) -> None: + model = scan(write_model(tmp_path / "model")) + adapter = scan(write_adapter(tmp_path / "adapter"), ArtifactKind.ADAPTER) + + assert model.passed and model.files == 3 + assert adapter.passed + assert model.digest.startswith("sha256:") and model.digest != adapter.digest + + +def test_pickle_weights_are_refused_and_what_they_would_import_is_named(tmp_path: Path) -> None: + root = write_model(tmp_path) + (root / "pytorch_model.bin").write_bytes(pickle.dumps(Hostile())) + + report = scan(root) + + finding = next(item for item in report.findings if item.rule is Rule.PICKLE) + assert finding.path == "pytorch_model.bin" + assert "runs code when loaded" in finding.detail + assert "system" in finding.detail + + +def test_a_pickle_is_found_whatever_the_file_is_called(tmp_path: Path) -> None: + root = write_model(tmp_path) + (root / "notes.txt").write_bytes(pickle.dumps(Hostile())) + with zipfile.ZipFile(root / "weights.dat", "w") as archive: + archive.writestr("archive/data.pkl", pickle.dumps(Hostile(), protocol=2)) + (root / "broken.zip.txt").write_bytes(b"PK not really a zip") + + found = {item.path for item in scan(root).findings if item.rule is Rule.PICKLE} + + assert found == {"notes.txt", "weights.dat"} + + +def test_pickle_imports_are_read_without_running_the_pickle() -> None: + for protocol in (0, 2, 4): + imports = pickle_imports(pickle.dumps(Hostile(), protocol=protocol)) + assert any(name.endswith(".system") for name in imports), protocol + assert pickle_imports(pickle.dumps({"weights": [1, 2, 3]})) == () + # A truncated stream still yields what came before the cut. + assert pickle_imports(pickle.dumps(Hostile(), protocol=0)[:-3]) + + +def test_an_artifact_may_not_ship_code_or_ask_for_it_to_be_imported(tmp_path: Path) -> None: + root = write_model(tmp_path) + (root / "modeling_custom.py").write_text("import os\n", encoding="utf-8") + (root / "config.json").write_text( + json.dumps({"auto_map": {"AutoModel": "modeling_custom.Model"}}), encoding="utf-8" + ) + (root / "tokenizer_config.json").write_text( + json.dumps({"trust_remote_code": True}), encoding="utf-8" + ) + + assert sorted(rules(root)) == [Rule.CODE, Rule.REMOTE_CODE, Rule.REMOTE_CODE] + + +@pytest.mark.parametrize("content", ["not json", "[]"]) +def test_an_unreadable_configuration_is_not_waved_through(tmp_path: Path, content: str) -> None: + root = write_model(tmp_path) + (root / "config.json").write_text(content, encoding="utf-8") + + assert rules(root) == [Rule.REMOTE_CODE] + + +@pytest.mark.parametrize( + ("content", "problem"), + [ + (b"\x01\x02", "shorter than the 8-byte header length"), + (struct.pack(" None: + root = write_model(tmp_path) + (root / "model.safetensors").write_bytes(content) + + findings = scan(root).findings + + assert [item.rule for item in findings] == [Rule.SAFETENSORS] + assert problem in findings[0].detail + + +def test_an_artifact_with_no_weights_is_refused(tmp_path: Path) -> None: + (tmp_path / "config.json").write_text("{}", encoding="utf-8") + + assert rules(tmp_path) == [Rule.NO_WEIGHTS] + + +@pytest.mark.parametrize( + ("config", "detail"), + [({"peft_type": "PROMPT_TUNING"}, "only LORA is served"), ({"r": 64}, "outside 1..32")], +) +def test_an_adapter_the_engine_could_not_serve_is_refused( + tmp_path: Path, config: dict[str, Any], detail: str +) -> None: + report = scan(write_adapter(tmp_path, **config), ArtifactKind.ADAPTER) + + assert [item.rule for item in report.findings] == [Rule.ADAPTER_CONFIG] + assert detail in report.findings[0].detail + + +def test_an_adapter_without_its_configuration_is_refused(tmp_path: Path) -> None: + write_adapter(tmp_path) + (tmp_path / "adapter_config.json").unlink() + + assert rules(tmp_path, ArtifactKind.ADAPTER) == [Rule.ADAPTER_CONFIG] + + +def test_a_link_out_of_the_artifact_is_refused(tmp_path: Path) -> None: + root = write_model(tmp_path / "model") + outside = tmp_path / "secret.txt" + outside.write_text("not part of the artifact", encoding="utf-8") + try: + (root / "leak.txt").symlink_to(outside) + (root / "alias.json").symlink_to(root / "tokenizer.json") + except OSError: + pytest.skip("this platform does not allow creating symlinks") + + findings = scan(root).findings + + assert [(item.rule, item.path) for item in findings] == [(Rule.SYMLINK, "leak.txt")] + + +def test_the_digest_follows_content_and_paths_and_gates_the_release(tmp_path: Path) -> None: + first = write_model(tmp_path / "first") + second = write_model(tmp_path / "second") + assert artifact_digest(first) == artifact_digest(second) + + recorded = artifact_digest(first) + assert scan(first, expected_digest=recorded).passed + + (second / "tokenizer.json").write_text('{"changed": true}', encoding="utf-8") + assert artifact_digest(second) != recorded + tampered = scan(second, expected_digest=recorded) + assert [item.rule for item in tampered.findings] == [Rule.CHECKSUM] + + (first / "tokenizer.json").rename(first / "renamed.json") + assert artifact_digest(first) != recorded + + +def test_the_command_line_passes_a_sound_artifact_and_fails_an_unsafe_one( + tmp_path: Path, capsys: pytest.CaptureFixture[str] +) -> None: + root = write_model(tmp_path / "model") + assert main([str(root)]) == 0 + assert json.loads(capsys.readouterr().out)["passed"] is True + + (root / "pytorch_model.bin").write_bytes(pickle.dumps({"weights": 1})) + assert main([str(root), "--expect", artifact_digest(root)]) == 1 + printed = json.loads(capsys.readouterr().out) + assert printed["passed"] is False + assert [item["rule"] for item in printed["findings"]] == ["pickle-format"] + + +def test_a_catalog_subject_supplies_the_kind_and_the_recorded_checksum( + tmp_path: Path, capsys: pytest.CaptureFixture[str] +) -> None: + adapter = write_adapter(tmp_path / "adapter") + + # The catalog records a placeholder checksum, which no real artifact has. + assert main([str(adapter), "--subject", "claims-extraction-lora"]) == 1 + printed = json.loads(capsys.readouterr().out) + assert printed["kind"] == "adapter" + assert [item["rule"] for item in printed["findings"]] == ["checksum-mismatch"] + assert "sha256:mock-claims" in printed["findings"][0]["detail"] + + # A subject with no recorded checksum is scanned without a digest gate. + assert main([str(adapter), "--subject", "support-classification-lora"]) == 0 + + +@pytest.mark.parametrize("arguments", [["missing-directory"], [".", "--subject", "nobody"]]) +def test_the_command_line_refuses_what_it_cannot_scan(arguments: list[str]) -> None: + with pytest.raises(SystemExit) as raised: + main(arguments) + assert raised.value.code == 2