atlas-iac/testing/tests/test_hermes_multiarch_combine.py

"""Safety tests for the fail-closed multi-arch manifest-list combiner."""

from __future__ import annotations

import hashlib
import importlib.util
import io
import json
import sys
from pathlib import Path

import pytest


ROOT = Path(__file__).resolve().parents[2]
SCRIPT = ROOT / "ci/scripts/hermes_multiarch_combine.py"
REVISION = "a" * 40
BUILD = "17"
DESTINATION = f"registry.bstein.dev/bstein/hermes-agent:git-{REVISION}-build-{BUILD}"
DOCKER_MANIFEST_LIST = "application/vnd.docker.distribution.manifest.list.v2+json"
DOCKER_MANIFEST = "application/vnd.docker.distribution.manifest.v2+json"
IMAGE_CONFIG = "application/vnd.docker.container.image.v1+json"


def _load():
    spec = importlib.util.spec_from_file_location("hermes_multiarch_combine", SCRIPT)
    assert spec is not None and spec.loader is not None
    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)
    return module


def _sha(body: bytes) -> str:
    return "sha256:" + hashlib.sha256(body).hexdigest()


class Response(io.BytesIO):
    """Minimal context-managed urllib response."""

    def __init__(self, status: int, body: bytes = b"", headers=None):
        super().__init__(body)
        self.status = status
        self.headers = headers or {}

    def __enter__(self):
        return self

    def __exit__(self, *_args):
        self.close()


def _leaf(architecture: str, *, os_name: str = "linux"):
    """Build a self-consistent config+manifest whose bytes hash to real digests."""
    config_body = json.dumps(
        {"architecture": architecture, "os": os_name}, sort_keys=True
    ).encode()
    config_digest = _sha(config_body)
    manifest_body = json.dumps(
        {
            "schemaVersion": 2,
            "mediaType": DOCKER_MANIFEST,
            "config": {
                "mediaType": IMAGE_CONFIG,
                "digest": config_digest,
                "size": len(config_body),
            },
            "layers": [],
        },
        sort_keys=True,
    ).encode()
    return {
        "arch": architecture,
        "config_body": config_body,
        "config_digest": config_digest,
        "manifest_body": manifest_body,
        "digest": _sha(manifest_body),
        "content_type": DOCKER_MANIFEST,
    }


def _descriptor(module, leaf):
    return {
        "mediaType": leaf["content_type"],
        "size": len(leaf["manifest_body"]),
        "digest": leaf["digest"],
        "platform": {"architecture": leaf["arch"], "os": "linux"},
    }


def _index_digest(module, amd64, arm64) -> str:
    body = module._manifest_list_bytes(
        [_descriptor(module, amd64), _descriptor(module, arm64)]
    )
    return _sha(body)


class Registry:
    """Route the combiner's deterministic request sequence by method and path."""

    def __init__(self, module, amd64, arm64, *, head_status=404, existing="") -> None:
        self.module = module
        self.leaves = {amd64["arch"]: amd64, arm64["arch"]: arm64}
        self.index_digest = _index_digest(module, amd64, arm64)
        self.head_status = head_status
        self.existing = existing
        self.put_body = None
        self.calls = []

    def __call__(self, request, timeout):
        method = request.method
        url = request.full_url
        self.calls.append((method, url))
        for arch, leaf in self.leaves.items():
            if url.endswith(f"-{arch}"):
                return Response(
                    200,
                    leaf["manifest_body"],
                    {
                        "Docker-Content-Digest": leaf["digest"],
                        "Content-Type": leaf["content_type"],
                    },
                )
            # Blob URLs percent-encode the ``sha256:`` colon; match the raw hex.
            if url.endswith(leaf["config_digest"].split(":", 1)[1]):
                return Response(200, leaf["config_body"], {})
        # Final index tag (no arch suffix, ends with the build tag).
        if method == "HEAD":
            return Response(
                self.head_status,
                headers={"Docker-Content-Digest": self.existing},
            )
        if method == "PUT":
            self.put_body = request.data
            return Response(201, headers={"Docker-Content-Digest": self.index_digest})
        return Response(
            200,
            self.module._manifest_list_bytes(
                [
                    _descriptor(self.module, self.leaves["amd64"]),
                    _descriptor(self.module, self.leaves["arm64"]),
                ]
            ),
            {
                "Docker-Content-Digest": self.index_digest,
                "Content-Type": DOCKER_MANIFEST_LIST,
            },
        )


def _combine(module, registry, arch_digests=None):
    if arch_digests is None:
        arch_digests = {
            "amd64": registry.leaves["amd64"]["digest"],
            "arm64": registry.leaves["arm64"]["digest"],
        }
    return module.combine_multiarch_index(
        destination=DESTINATION,
        arch_digests=arch_digests,
        username="robot",
        password="private",
        opener=registry,
    )


def test_combines_two_verified_leaves_into_one_index() -> None:
    """Both native leaves are re-read, arch-proven, and published as a list."""
    module = _load()
    amd64, arm64 = _leaf("amd64"), _leaf("arm64")
    registry = Registry(module, amd64, arm64)
    result = _combine(module, registry)
    assert result["result"] == "published"
    assert result["index_digest"] == registry.index_digest
    assert result["amd64_digest"] == amd64["digest"]
    assert result["arm64_digest"] == arm64["digest"]
    # The published bytes are exactly what we hashed for the index digest.
    assert _sha(registry.put_body) == registry.index_digest
    assert "private" not in json.dumps(result)


def test_idempotent_when_index_already_matches() -> None:
    """A replay is accepted only when the existing index digest is identical."""
    module = _load()
    amd64, arm64 = _leaf("amd64"), _leaf("arm64")
    index_digest = _index_digest(module, amd64, arm64)
    registry = Registry(
        module, amd64, arm64, head_status=200, existing=index_digest
    )
    result = _combine(module, registry)
    assert result["result"] == "already-present"
    assert registry.put_body is None


def test_existing_index_with_other_digest_fails_closed() -> None:
    """An occupied final tag with a different index digest never republishes."""
    module = _load()
    amd64, arm64 = _leaf("amd64"), _leaf("arm64")
    registry = Registry(
        module, amd64, arm64, head_status=200, existing="sha256:" + "c" * 64
    )
    with pytest.raises(RuntimeError, match="another index digest"):
        _combine(module, registry)


def test_rejects_leaf_digest_that_disagrees_with_evidence() -> None:
    """A leaf whose registry digest is not the Kaniko evidence fails closed."""
    module = _load()
    amd64, arm64 = _leaf("amd64"), _leaf("arm64")
    registry = Registry(module, amd64, arm64)
    with pytest.raises(RuntimeError, match="amd64 leaf digest"):
        _combine(
            module,
            registry,
            arch_digests={"amd64": "sha256:" + "d" * 64, "arm64": arm64["digest"]},
        )


def test_rejects_leaf_whose_config_architecture_is_wrong() -> None:
    """A leaf that claims the wrong architecture in its config fails closed."""
    module = _load()
    # Build an "amd64" candidate tag whose config actually says arm64.
    swapped = _leaf("arm64")
    swapped["arch"] = "amd64"  # served under the amd64 tag, but arm64 inside
    arm64 = _leaf("arm64")
    registry = Registry(module, swapped, arm64)
    with pytest.raises(RuntimeError, match="architecture"):
        _combine(
            module,
            registry,
            arch_digests={"amd64": swapped["digest"], "arm64": arm64["digest"]},
        )


def test_rejects_leaf_served_as_a_manifest_list() -> None:
    """A per-arch leaf must be a single image, never itself an index/list."""
    module = _load()
    amd64, arm64 = _leaf("amd64"), _leaf("arm64")
    amd64["content_type"] = DOCKER_MANIFEST_LIST
    registry = Registry(module, amd64, arm64)
    with pytest.raises(RuntimeError, match="single-image manifest"):
        _combine(module, registry)


def test_requires_both_architectures() -> None:
    """The combiner refuses any arch set other than exactly arm64 and amd64."""
    module = _load()
    amd64, arm64 = _leaf("amd64"), _leaf("arm64")
    registry = Registry(module, amd64, arm64)
    with pytest.raises(ValueError, match="arm64 and amd64"):
        module.combine_multiarch_index(
            destination=DESTINATION,
            arch_digests={"amd64": amd64["digest"]},
            username="robot",
            password="private",
            opener=registry,
        )


def test_rejects_destination_with_arch_suffix() -> None:
    """Only the arch-less final tag is a valid combine destination."""
    module = _load()
    amd64, arm64 = _leaf("amd64"), _leaf("arm64")
    with pytest.raises(ValueError, match="invalid multi-arch destination"):
        module.combine_multiarch_index(
            destination=f"{DESTINATION}-arm64",
            arch_digests={"amd64": amd64["digest"], "arm64": arm64["digest"]},
            username="robot",
            password="private",
            opener=Registry(module, amd64, arm64),
        )


def test_cli_binds_evidence_files_and_writes_index_outputs(
    tmp_path: Path, monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str]
) -> None:
    """The command reads both Kaniko pairs and emits arch-less index evidence."""
    module = _load()
    amd64, arm64 = _leaf("amd64"), _leaf("arm64")
    registry = Registry(module, amd64, arm64)
    files = {}
    for arch, leaf in (("amd64", amd64), ("arm64", arm64)):
        digest_file = tmp_path / f"{arch}.digest"
        image_file = tmp_path / f"{arch}.image"
        digest_file.write_text(leaf["digest"] + "\n", encoding="utf-8")
        image_file.write_text(
            f"{DESTINATION}-{arch}@{leaf['digest']}\n", encoding="utf-8"
        )
        files[arch] = (digest_file, image_file)
    out_digest = tmp_path / "index.digest"
    out_image = tmp_path / "index.image"
    monkeypatch.setenv("HARBOR_USER", "robot")
    monkeypatch.setenv("HARBOR_PASSWORD", "private")
    captured = {}

    def fake_combine(*, destination, arch_digests, username, password):
        captured["destination"] = destination
        captured["arch_digests"] = dict(arch_digests)
        return {"index_digest": registry.index_digest, "result": "published"}

    monkeypatch.setattr(module, "combine_multiarch_index", fake_combine)
    monkeypatch.setattr(
        sys,
        "argv",
        [
            "hermes_multiarch_combine.py",
            "--destination",
            DESTINATION,
            "--source-revision",
            REVISION,
            "--build-number",
            BUILD,
            "--arm64-digest-file",
            str(files["arm64"][0]),
            "--arm64-image-file",
            str(files["arm64"][1]),
            "--amd64-digest-file",
            str(files["amd64"][0]),
            "--amd64-image-file",
            str(files["amd64"][1]),
            "--digest-file",
            str(out_digest),
            "--image-file",
            str(out_image),
        ],
    )
    assert module.main() == 0
    assert out_digest.read_text(encoding="utf-8").strip() == registry.index_digest
    assert (
        out_image.read_text(encoding="utf-8").strip()
        == f"{DESTINATION}@{registry.index_digest}"
    )
    assert json.loads(capsys.readouterr().out)["result"] == "published"
    # main() must bind each arch's two evidence files to the correct leaf digest.
    assert captured["destination"] == DESTINATION
    assert captured["arch_digests"] == {
        "amd64": amd64["digest"],
        "arm64": arm64["digest"],
    }


def test_cli_rejects_mismatched_per_arch_evidence_pair(
    tmp_path: Path, monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str]
) -> None:
    """A digest/image evidence file mismatch is reported as a JSON error."""
    module = _load()
    amd64, arm64 = _leaf("amd64"), _leaf("arm64")
    files = {}
    for arch, leaf in (("amd64", amd64), ("arm64", arm64)):
        digest_file = tmp_path / f"{arch}.digest"
        image_file = tmp_path / f"{arch}.image"
        digest_file.write_text(leaf["digest"] + "\n", encoding="utf-8")
        # amd64 image file points at the wrong digest.
        recorded = "sha256:" + "e" * 64 if arch == "amd64" else leaf["digest"]
        image_file.write_text(
            f"{DESTINATION}-{arch}@{recorded}\n", encoding="utf-8"
        )
        files[arch] = (digest_file, image_file)
    monkeypatch.setenv("HARBOR_USER", "robot")
    monkeypatch.setenv("HARBOR_PASSWORD", "private")
    monkeypatch.setattr(
        sys,
        "argv",
        [
            "hermes_multiarch_combine.py",
            "--destination",
            DESTINATION,
            "--source-revision",
            REVISION,
            "--build-number",
            BUILD,
            "--arm64-digest-file",
            str(files["arm64"][0]),
            "--arm64-image-file",
            str(files["arm64"][1]),
            "--amd64-digest-file",
            str(files["amd64"][0]),
            "--amd64-image-file",
            str(files["amd64"][1]),
            "--digest-file",
            str(tmp_path / "index.digest"),
            "--image-file",
            str(tmp_path / "index.image"),
        ],
    )
    assert module.main() == 1
    assert "does not match" in json.loads(capsys.readouterr().out)["error"]


# The combiner serves both Hermes multi-arch images; only the repository name in
# the destination differs. hermes-webui must route to its own registry paths.
WEBUI_DESTINATION = (
    f"registry.bstein.dev/bstein/hermes-webui:git-{REVISION}-build-{BUILD}"
)


def test_destination_pattern_captures_both_components() -> None:
    """The fail-closed pattern accepts exactly hermes-agent and hermes-webui."""
    module = _load()
    agent = module.DESTINATION_PATTERN.fullmatch(DESTINATION)
    webui = module.DESTINATION_PATTERN.fullmatch(WEBUI_DESTINATION)
    assert agent is not None and agent.group("component") == "hermes-agent"
    assert webui is not None and webui.group("component") == "hermes-webui"
    # A third, unexpected repository is still rejected.
    assert (
        module.DESTINATION_PATTERN.fullmatch(
            f"registry.bstein.dev/bstein/hermes-other:git-{REVISION}-build-{BUILD}"
        )
        is None
    )


def test_combines_webui_leaves_under_the_webui_component() -> None:
    """A hermes-webui destination re-reads leaves from the webui registry paths."""
    module = _load()
    amd64, arm64 = _leaf("amd64"), _leaf("arm64")
    registry = Registry(module, amd64, arm64)
    result = module.combine_multiarch_index(
        destination=WEBUI_DESTINATION,
        arch_digests={"amd64": amd64["digest"], "arm64": arm64["digest"]},
        username="robot",
        password="private",
        opener=registry,
    )
    assert result["component"] == "hermes-webui"
    assert result["result"] == "published"
    # Every registry call the combiner made must be scoped to the webui repo, and
    # never leak into the agent repo.
    assert registry.calls, "combiner made no registry calls"
    for _method, url in registry.calls:
        assert "/v2/bstein/hermes-webui/" in url
        assert "hermes-agent" not in url