Add OCI registry publish support

This commit is contained in:
Thales Maciel 2026-03-08 18:27:32 -03:00
parent f6d3bf0e90
commit 6406f673c1
3 changed files with 817 additions and 3 deletions

View file

@ -3,18 +3,22 @@
from __future__ import annotations
import argparse
import base64
import hashlib
import io
import json
import os
import shutil
import subprocess
import tarfile
import urllib.error
import urllib.parse
import urllib.request
import uuid
from dataclasses import dataclass
from datetime import UTC, datetime
from pathlib import Path
from typing import Any
from typing import Any, cast
from pyro_mcp.runtime import DEFAULT_PLATFORM
from pyro_mcp.vm_environments import get_environment
@ -31,6 +35,8 @@ OCI_IMAGE_LAYER_MEDIA_TYPE = "application/vnd.oci.image.layer.v1.tar"
OCI_IMAGE_INDEX_MEDIA_TYPE = "application/vnd.oci.image.index.v1+json"
OCI_LAYOUT_VERSION = "1.0.0"
VALIDATION_READ_LIMIT = 1024 * 1024
DEFAULT_OCI_USERNAME_ENV = "OCI_REGISTRY_USERNAME"
DEFAULT_OCI_PASSWORD_ENV = "OCI_REGISTRY_PASSWORD"
@dataclass(frozen=True)
@ -250,6 +256,340 @@ def _write_tar_blob_from_bytes(
raise
def _normalize_headers(headers: Any) -> dict[str, str]:
return {str(key).lower(): str(value) for key, value in headers.items()}
def _basic_auth_header(username: str | None, password: str | None) -> str | None:
if username is None or password is None:
return None
token = base64.b64encode(f"{username}:{password}".encode("utf-8")).decode("ascii")
return f"Basic {token}"
def _parse_authenticate_parameters(raw: str) -> dict[str, str]:
params: dict[str, str] = {}
for segment in raw.split(","):
if "=" not in segment:
continue
key, value = segment.split("=", 1)
params[key.strip()] = value.strip().strip('"')
return params
def _registry_url(registry: str, repository: str, suffix: str) -> str:
return f"https://{registry}/v2/{repository}/{suffix}"
def _registry_credentials(
*,
username_env: str = DEFAULT_OCI_USERNAME_ENV,
password_env: str = DEFAULT_OCI_PASSWORD_ENV,
) -> tuple[str | None, str | None]:
username = os.environ.get(username_env)
password = os.environ.get(password_env)
return (
username if isinstance(username, str) and username != "" else None,
password if isinstance(password, str) and password != "" else None,
)
def _rewind_body(data: object) -> None:
if hasattr(data, "seek"):
data.seek(0)
def _read_response_body(response: Any) -> bytes:
if hasattr(response, "read"):
return bytes(response.read())
return b""
def _fetch_registry_token(
authenticate: str,
*,
repository: str,
scope_actions: str,
username: str | None,
password: str | None,
) -> str:
if not authenticate.startswith("Bearer "):
raise RuntimeError("unsupported OCI registry authentication scheme")
params = _parse_authenticate_parameters(authenticate[len("Bearer ") :])
realm = params.get("realm")
if realm is None:
raise RuntimeError("OCI auth challenge did not include a token realm")
query = {
"service": params.get("service", ""),
"scope": params.get("scope", f"repository:{repository}:{scope_actions}"),
}
headers: dict[str, str] = {}
basic_auth = _basic_auth_header(username, password)
if basic_auth is not None:
headers["Authorization"] = basic_auth
request = urllib.request.Request(
f"{realm}?{urllib.parse.urlencode(query)}",
headers=headers,
method="GET",
)
with urllib.request.urlopen(request, timeout=90) as response: # noqa: S310
payload = json.loads(_read_response_body(response).decode("utf-8"))
if not isinstance(payload, dict):
raise RuntimeError("OCI auth token response was not a JSON object")
raw_token = payload.get("token") or payload.get("access_token")
if not isinstance(raw_token, str) or raw_token == "":
raise RuntimeError("OCI auth token response did not include a bearer token")
return raw_token
def _request_registry(
url: str,
*,
method: str,
repository: str,
scope_actions: str,
headers: dict[str, str] | None = None,
data: object = None,
allow_statuses: tuple[int, ...] = (),
username: str | None = None,
password: str | None = None,
) -> tuple[int, bytes, dict[str, str]]:
request_headers = dict(headers or {})
request = urllib.request.Request(
url,
data=cast(Any, data),
headers=request_headers,
method=method,
)
try:
with urllib.request.urlopen(request, timeout=90) as response: # noqa: S310
return (
response.status,
_read_response_body(response),
_normalize_headers(response.headers),
)
except urllib.error.HTTPError as exc:
if exc.code == 401:
authenticate = exc.headers.get("WWW-Authenticate")
if authenticate is None:
raise RuntimeError("OCI registry denied access without an auth challenge") from exc
token = _fetch_registry_token(
authenticate,
repository=repository,
scope_actions=scope_actions,
username=username,
password=password,
)
authenticated_headers = {**request_headers, "Authorization": f"Bearer {token}"}
_rewind_body(data)
retry = urllib.request.Request(
url,
data=cast(Any, data),
headers=authenticated_headers,
method=method,
)
try:
with urllib.request.urlopen(retry, timeout=90) as response: # noqa: S310
return (
response.status,
_read_response_body(response),
_normalize_headers(response.headers),
)
except urllib.error.HTTPError as retry_exc:
if retry_exc.code in allow_statuses:
return (
retry_exc.code,
_read_response_body(retry_exc),
_normalize_headers(retry_exc.headers),
)
raise RuntimeError(f"registry request failed for {url}: {retry_exc}") from retry_exc
if exc.code in allow_statuses:
return exc.code, _read_response_body(exc), _normalize_headers(exc.headers)
raise RuntimeError(f"registry request failed for {url}: {exc}") from exc
def _load_oci_layout_manifest(
layout_dir: Path,
) -> tuple[dict[str, Any], bytes, dict[str, Any], str]:
index_path = layout_dir / "index.json"
if not index_path.exists():
raise RuntimeError(f"OCI layout index not found: {index_path}")
index_payload = json.loads(index_path.read_text(encoding="utf-8"))
manifests = index_payload.get("manifests")
if not isinstance(manifests, list) or len(manifests) != 1:
raise RuntimeError("OCI layout must contain exactly one manifest descriptor")
descriptor = manifests[0]
if not isinstance(descriptor, dict):
raise RuntimeError("OCI layout manifest descriptor is malformed")
raw_digest = descriptor.get("digest")
if not isinstance(raw_digest, str):
raise RuntimeError("OCI layout manifest descriptor is missing a digest")
manifest_path = _blob_path(layout_dir / "blobs", raw_digest)
manifest_bytes = manifest_path.read_bytes()
manifest = json.loads(manifest_bytes.decode("utf-8"))
if not isinstance(manifest, dict):
raise RuntimeError("OCI layout manifest payload is malformed")
media_type = descriptor.get("mediaType")
if not isinstance(media_type, str) or media_type == "":
media_type = str(manifest.get("mediaType") or OCI_IMAGE_MANIFEST_MEDIA_TYPE)
return descriptor, manifest_bytes, manifest, media_type
def _blob_exists(
*,
registry: str,
repository: str,
digest: str,
username: str | None,
password: str | None,
) -> bool:
status, _, _ = _request_registry(
_registry_url(registry, repository, f"blobs/{digest}"),
method="HEAD",
repository=repository,
scope_actions="pull,push",
allow_statuses=(404,),
username=username,
password=password,
)
return status == 200
def _upload_blob(
*,
registry: str,
repository: str,
digest: str,
blob_path: Path,
username: str | None,
password: str | None,
) -> dict[str, Any]:
if _blob_exists(
registry=registry,
repository=repository,
digest=digest,
username=username,
password=password,
):
return {"digest": digest, "size": blob_path.stat().st_size, "uploaded": False}
status, _, headers = _request_registry(
_registry_url(registry, repository, "blobs/uploads/"),
method="POST",
repository=repository,
scope_actions="pull,push",
headers={"Content-Length": "0"},
allow_statuses=(202,),
username=username,
password=password,
)
if status != 202:
raise RuntimeError(f"unexpected registry status when starting blob upload: {status}")
location = headers.get("location")
if location is None:
raise RuntimeError("registry did not return a blob upload location")
upload_url = urllib.parse.urljoin(f"https://{registry}", location)
separator = "&" if "?" in upload_url else "?"
upload_url = f"{upload_url}{separator}{urllib.parse.urlencode({'digest': digest})}"
size = blob_path.stat().st_size
with blob_path.open("rb") as blob_fp:
status, _, _ = _request_registry(
upload_url,
method="PUT",
repository=repository,
scope_actions="pull,push",
headers={
"Content-Length": str(size),
"Content-Type": "application/octet-stream",
},
data=blob_fp,
allow_statuses=(201,),
username=username,
password=password,
)
if status != 201:
raise RuntimeError(f"unexpected registry status when uploading blob: {status}")
return {"digest": digest, "size": size, "uploaded": True}
def publish_environment_oci_layout(
*,
environment: str,
layout_root: Path,
registry: str | None = None,
repository: str | None = None,
reference: str | None = None,
username_env: str = DEFAULT_OCI_USERNAME_ENV,
password_env: str = DEFAULT_OCI_PASSWORD_ENV,
) -> dict[str, Any]:
spec = get_environment(environment)
resolved_registry = registry or spec.oci_registry
resolved_repository = repository or spec.oci_repository
resolved_reference = reference or spec.oci_reference or spec.version
if resolved_registry is None or resolved_repository is None:
raise RuntimeError(f"environment {environment!r} does not define an OCI registry target")
layout_dir = layout_root / _environment_slug(environment)
descriptor, manifest_bytes, manifest, media_type = _load_oci_layout_manifest(layout_dir)
username, password = _registry_credentials(
username_env=username_env,
password_env=password_env,
)
config = manifest.get("config")
layers = manifest.get("layers")
if not isinstance(config, dict):
raise RuntimeError("OCI layout manifest is missing a config descriptor")
if not isinstance(layers, list):
raise RuntimeError("OCI layout manifest is missing layers")
uploaded_blobs: list[dict[str, Any]] = []
descriptors = [config, *layers]
for entry in descriptors:
if not isinstance(entry, dict):
raise RuntimeError("OCI layout descriptor is malformed")
raw_digest = entry.get("digest")
if not isinstance(raw_digest, str):
raise RuntimeError("OCI layout descriptor is missing a digest")
blob_path = _blob_path(layout_dir / "blobs", raw_digest)
uploaded_blobs.append(
_upload_blob(
registry=resolved_registry,
repository=resolved_repository,
digest=raw_digest,
blob_path=blob_path,
username=username,
password=password,
)
)
status, _, headers = _request_registry(
_registry_url(resolved_registry, resolved_repository, f"manifests/{resolved_reference}"),
method="PUT",
repository=resolved_repository,
scope_actions="pull,push",
headers={
"Content-Length": str(len(manifest_bytes)),
"Content-Type": media_type,
},
data=manifest_bytes,
allow_statuses=(201,),
username=username,
password=password,
)
if status != 201:
raise RuntimeError(f"unexpected registry status when publishing manifest: {status}")
return {
"environment": spec.name,
"layout_dir": str(layout_dir),
"registry": resolved_registry,
"repository": resolved_repository,
"reference": resolved_reference,
"manifest_digest": headers.get("docker-content-digest", str(descriptor["digest"])),
"uploaded_blobs": uploaded_blobs,
}
def validate_sources(paths: RuntimeBuildPaths, lock: RuntimeBuildLock) -> None:
firecracker_source = _resolved_source_path(paths, lock.binaries["firecracker"])
jailer_source = _resolved_source_path(paths, lock.binaries["jailer"])
@ -761,6 +1101,7 @@ def _build_parser() -> argparse.ArgumentParser: # pragma: no cover - CLI wiring
"build-rootfs",
"materialize",
"export-environment-oci",
"publish-environment-oci",
"stage-binaries",
"stage-kernel",
"stage-rootfs",
@ -778,7 +1119,12 @@ def _build_parser() -> argparse.ArgumentParser: # pragma: no cover - CLI wiring
parser.add_argument("--materialized-dir", default=str(DEFAULT_RUNTIME_MATERIALIZED_DIR))
parser.add_argument("--environment")
parser.add_argument("--output-dir", default=str(DEFAULT_RUNTIME_OCI_LAYOUT_DIR))
parser.add_argument("--layout-root", default=str(DEFAULT_RUNTIME_OCI_LAYOUT_DIR))
parser.add_argument("--registry")
parser.add_argument("--repository")
parser.add_argument("--reference")
parser.add_argument("--username-env", default=DEFAULT_OCI_USERNAME_ENV)
parser.add_argument("--password-env", default=DEFAULT_OCI_PASSWORD_ENV)
return parser
@ -791,14 +1137,16 @@ def main() -> None: # pragma: no cover - CLI wiring
materialized_dir=Path(args.materialized_dir),
platform=args.platform,
)
lock = _load_lock(paths)
if args.command == "fetch-binaries":
lock = _load_lock(paths)
materialize_binaries(paths, lock)
return
if args.command == "build-kernel":
lock = _load_lock(paths)
materialize_kernel(paths, lock)
return
if args.command == "build-rootfs":
lock = _load_lock(paths)
materialize_rootfs(paths, lock)
return
if args.command == "materialize":
@ -815,30 +1163,50 @@ def main() -> None: # pragma: no cover - CLI wiring
)
print(json.dumps(result, indent=2, sort_keys=True))
return
if args.command == "publish-environment-oci":
if not isinstance(args.environment, str) or args.environment == "":
raise RuntimeError("--environment is required for publish-environment-oci")
result = publish_environment_oci_layout(
environment=args.environment,
layout_root=Path(args.layout_root),
registry=args.registry,
repository=args.repository,
reference=args.reference,
username_env=args.username_env,
password_env=args.password_env,
)
print(json.dumps(result, indent=2, sort_keys=True))
return
if args.command == "bundle":
build_bundle(paths, sync=True)
return
if args.command == "stage-binaries":
lock = _load_lock(paths)
paths.build_platform_root.mkdir(parents=True, exist_ok=True)
_copy_notice(paths)
stage_binaries(paths, lock)
return
if args.command == "stage-kernel":
lock = _load_lock(paths)
paths.build_platform_root.mkdir(parents=True, exist_ok=True)
stage_kernel(paths, lock)
return
if args.command == "stage-rootfs":
lock = _load_lock(paths)
paths.build_platform_root.mkdir(parents=True, exist_ok=True)
stage_rootfs(paths, lock)
return
if args.command == "stage-agent":
lock = _load_lock(paths)
paths.build_platform_root.mkdir(parents=True, exist_ok=True)
stage_agent(paths, lock)
return
if args.command == "validate":
lock = _load_lock(paths)
validate_sources(paths, lock)
return
if args.command == "manifest":
lock = _load_lock(paths)
generate_manifest(paths, lock)
return
if args.command == "sync":