diff --git a/bot_bottle/backend/firecracker/infra_artifact.py b/bot_bottle/backend/firecracker/infra_artifact.py index d8af7ee..0108cf8 100644 --- a/bot_bottle/backend/firecracker/infra_artifact.py +++ b/bot_bottle/backend/firecracker/infra_artifact.py @@ -57,24 +57,34 @@ def local_build_requested() -> bool: return os.environ.get("BOT_BOTTLE_INFRA_BUILD", "").strip().lower() == "local" -def infra_artifact_version(init_script: str) -> str: +def infra_artifact_version(init_script: str, *, repo_root: Path = _REPO_ROOT) -> str: """Content hash (16 hex) of everything baked into the infra rootfs: the - whole shipped `bot_bottle` package (the infra image `COPY`s it wholesale), - the three fixed Dockerfiles, and the guest init. Deterministic across the - publish host and the launch host when both run the same checkout, so the - tag the launch host pulls is exactly the tag publish produced.""" + whole shipped `bot_bottle` package, the three fixed Dockerfiles, and the + guest init. Deterministic across the publish host and the launch host when + both run the same checkout, so the tag the launch host pulls is exactly the + tag publish produced. + + The package is `COPY bot_bottle /app/bot_bottle`'d wholesale into the image, + so hash *every* regular file under it — not just `*.py`. Non-Python inputs + (e.g. `egress_entrypoint.sh`, `netpool.defaults.env`) are baked in too, and + a change to one must bump the version or a launch host could boot a stale + rootfs whose code differs from its checkout. `__pycache__`/`.pyc` are the + only exclusions — build artifacts, never copied.""" h = hashlib.sha256() h.update(f"format={_ARTIFACT_FORMAT}\n".encode()) - pkg = _REPO_ROOT / "bot_bottle" - for path in sorted(pkg.rglob("*.py")): - if "__pycache__" in path.parts: + pkg = repo_root / "bot_bottle" + for path in sorted(pkg.rglob("*")): + if not path.is_file(): continue - h.update(str(path.relative_to(_REPO_ROOT)).encode()) + if "__pycache__" in path.parts or path.suffix == ".pyc": + continue + h.update(str(path.relative_to(repo_root)).encode()) + h.update(b"\0") h.update(path.read_bytes()) for name in _DOCKERFILES: - p = _REPO_ROOT / name h.update(name.encode()) - h.update(p.read_bytes()) + h.update(b"\0") + h.update((repo_root / name).read_bytes()) h.update(b"init\0") h.update(init_script.encode()) return h.hexdigest()[:16] diff --git a/bot_bottle/backend/firecracker/publish_infra.py b/bot_bottle/backend/firecracker/publish_infra.py index ae67b23..2fbc28a 100644 --- a/bot_bottle/backend/firecracker/publish_infra.py +++ b/bot_bottle/backend/firecracker/publish_infra.py @@ -47,8 +47,22 @@ def _sha256(path: Path) -> str: return h.hexdigest() -def _put(url: str, body: bytes, token: str) -> None: - req = urllib.request.Request(url, data=body, method="PUT") +def _put(url: str, body: "bytes | Path", token: str) -> None: + """PUT `body` (raw bytes, or a Path streamed from disk) to `url`. The rootfs + is hundreds of MB, so it is passed as a Path and streamed — `urlopen` reads + the open file in blocks rather than materializing it in memory (with an + explicit Content-Length, which Gitea requires and which also stops urllib + from `len()`-ing a non-bytes body).""" + handle = None + if isinstance(body, Path): + length = body.stat().st_size + handle = open(body, "rb") + data: object = handle + else: + length = len(body) + data = body + req = urllib.request.Request(url, data=data, method="PUT") # type: ignore[arg-type] + req.add_header("Content-Length", str(length)) if token: req.add_header("Authorization", f"token {token}") req.add_header("Content-Type", "application/octet-stream") @@ -64,6 +78,9 @@ def _put(url: str, body: bytes, token: str) -> None: raise SystemExit(f"upload failed (HTTP {e.code}): {url}\n{e.read().decode(errors='replace')}") except urllib.error.URLError as e: raise SystemExit(f"registry unreachable: {url} ({e.reason})") + finally: + if handle is not None: + handle.close() def _delete(url: str, token: str) -> None: @@ -127,8 +144,8 @@ def main(argv: list[str] | None = None) -> int: if args.force: _delete(gz_url, token) _delete(sha_url, token) - _put(gz_url, gz.read_bytes(), token) - _put(sha_url, sha.read_bytes(), token) + _put(gz_url, gz, token) # streamed from disk (hundreds of MB) + _put(sha_url, sha.read_bytes(), token) # tiny, in-memory is fine print(f"published infra rootfs {version}") return 0 diff --git a/tests/unit/test_infra_artifact.py b/tests/unit/test_infra_artifact.py index 22dd50c..6d84cb9 100644 --- a/tests/unit/test_infra_artifact.py +++ b/tests/unit/test_infra_artifact.py @@ -79,6 +79,44 @@ class TestVersion(unittest.TestCase): ia.infra_artifact_version("a"), ia.infra_artifact_version("b")) +class TestVersionInputs(unittest.TestCase): + """The hash must cover *every* file baked into the rootfs, not just `*.py` + (`COPY bot_bottle` is wholesale) — else a non-Python change (e.g. the egress + entrypoint shell script) leaves the version unchanged and a launch host + boots a rootfs whose code differs from its checkout.""" + + def _fake_repo(self, root: Path) -> None: + pkg = root / "bot_bottle" + pkg.mkdir() + (pkg / "app.py").write_text("print('hi')\n") + (pkg / "egress_entrypoint.sh").write_text("#!/bin/sh\nexec mitmdump\n") + (pkg / "netpool.defaults.env").write_text("FOO=1\n") + for name in ("Dockerfile.orchestrator", "Dockerfile.gateway", "Dockerfile.infra"): + (root / name).write_text(f"FROM scratch # {name}\n") + + def test_non_python_file_change_bumps_version(self) -> None: + with tempfile.TemporaryDirectory() as d: + root = Path(d) + self._fake_repo(root) + before = ia.infra_artifact_version("init", repo_root=root) + (root / "bot_bottle" / "egress_entrypoint.sh").write_text( + "#!/bin/sh\nexec mitmdump --different\n") + after = ia.infra_artifact_version("init", repo_root=root) + self.assertNotEqual(before, after) + + def test_pyc_and_pycache_ignored(self) -> None: + with tempfile.TemporaryDirectory() as d: + root = Path(d) + self._fake_repo(root) + before = ia.infra_artifact_version("init", repo_root=root) + cache = root / "bot_bottle" / "__pycache__" + cache.mkdir() + (cache / "app.cpython-312.pyc").write_bytes(b"\x00bytecode") + (root / "bot_bottle" / "app.pyc").write_bytes(b"\x00bytecode") + after = ia.infra_artifact_version("init", repo_root=root) + self.assertEqual(before, after) + + class TestEnsureArtifact(_CacheMixin): def test_downloads_verifies_and_caches(self) -> None: version = "deadbeef00000000" diff --git a/tests/unit/test_publish_infra.py b/tests/unit/test_publish_infra.py new file mode 100644 index 0000000..f1228ae --- /dev/null +++ b/tests/unit/test_publish_infra.py @@ -0,0 +1,62 @@ +"""Unit: the infra-artifact publisher's upload path (PRD 0069 Stage 2). + +The rootfs is hundreds of MB, so `_put` must stream it from disk rather than +read it into memory. Network is mocked; no Docker, no real build. +""" + +from __future__ import annotations + +import tempfile +import unittest +import urllib.request +from pathlib import Path +from unittest import mock + +from bot_bottle.backend.firecracker import publish_infra as pub + + +class _Resp: + status = 201 + + def __enter__(self) -> "_Resp": + return self + + def __exit__(self, *a: object) -> bool: + return False + + +class TestPut(unittest.TestCase): + def test_streams_file_body_with_content_length(self) -> None: + captured: list[urllib.request.Request] = [] + + def fake_urlopen(req: urllib.request.Request, *a: object, **k: object) -> _Resp: + captured.append(req) + return _Resp() + + with tempfile.TemporaryDirectory() as d: + f = Path(d) / "rootfs.ext4.gz" + payload = b"x" * 4096 + f.write_bytes(payload) + with mock.patch.object(pub.urllib.request, "urlopen", fake_urlopen): + pub._put("https://reg/pkg", f, token="t") + + req = captured[0] + # Body is the open file object (streamed), never the bytes in memory. + self.assertTrue(hasattr(req.data, "read")) + self.assertNotIsInstance(req.data, (bytes, bytearray)) + self.assertEqual(str(len(payload)), req.get_header("Content-length")) + + def test_small_bytes_body_still_works(self) -> None: + captured: list[urllib.request.Request] = [] + + def fake_urlopen(req: urllib.request.Request, *a: object, **k: object) -> _Resp: + captured.append(req) + return _Resp() + + with mock.patch.object(pub.urllib.request, "urlopen", fake_urlopen): + pub._put("https://reg/sha", b"abc123 rootfs\n", token="") + self.assertEqual(b"abc123 rootfs\n", captured[0].data) + + +if __name__ == "__main__": + unittest.main()