AMI upload, the production way

docs/user/aws.md

@@ -10,11 +10,11 @@ file format but an object in one EC2 region of one account: it refers to an EBS
 and records how instances boot it (architecture, boot mode, TPM, network driver). tine builds everything
 that is a function of the image, and uploads it as a separate `buck run` step that uses your AWS login.
 
-## The AMI bundle
+## The AMI
 
-Every bootable [image.disk](images.md#imagedisk) has an `[ami]` subtarget: the raw disk and the
-parameters to register it with. `buck build` it to get the bundle, `buck run` it to
-[import](#importing-an-ami) it.
+Every bootable [image.disk](images.md#imagedisk) has an `[ami]` subtarget. `buck build` it to build the
+disk and an `<IMAGE_ID>_<IMAGE_VERSION>_<arch>.ami.json` with its registration parameters.
+`buck run` it to [import](#importing-an-ami) the disk.
 
 Nothing about the disk is specific to EC2: every initrd supports
 [systemd IMDS](https://www.freedesktop.org/software/systemd/man/devel/systemd-imds.html) for passing
@@ -32,10 +32,10 @@ The AMI is named after the disk's published name: `<IMAGE_ID>_<IMAGE_VERSION>_<a
 EC2 does not accept in names (such as the `^` in versions) get replaced by `-`. The unchanged name is the
 AMI's description.
 
-With a `secure_boot_key`, the bundle carries a UEFI variable store that enrolls the key's certificate as PK,
-KEK and db, also available as `[ami][uefi-data]`. An instance then enforces Secure Boot from its first boot
-on and trusts only this key; `examples/image-secureboot` has such a disk. Without a key, an instance boots
-in setup mode with Secure Boot disabled.
+With a `secure_boot_key`, the parameters carry a UEFI variable store that enrolls the key's certificate as
+PK, KEK and db. An instance then enforces Secure Boot from its first boot on and trusts only this key;
+`examples/image-secureboot` has such a disk. Without a key, an instance boots in setup mode with Secure
+Boot disabled.
 
 ## Logging in
 
@@ -57,6 +57,20 @@ The import needs these IAM permissions:
 
 ## Importing an AMI
 
+Disk partitions often have a lot of empty space, e.g. an ESP sized for future updates. Only the blocks
+holding data are uploaded. This only works with unencrypted snapshots, so the import refuses to run while
+[EBS encryption by default](https://docs.aws.amazon.com/ebs/latest/userguide/encryption-by-default.html)
+is enabled for the region.
+
+If the import fails while uploading, EC2 cancels the incomplete snapshot after an hour. If it fails while
+registering, the completed snapshot stays; the import printed its ID, delete it with
+`aws ec2 delete-snapshot`.
+
+### Development mode
+
+This is the convenient way for testing: Building and uploading are a single command, but the entire build
+has access to your AWS login from above. Don't do this on untrusted projects!
+
 Run the disk's `[ami]` subtarget:
 
 ```sh
@@ -72,14 +86,30 @@ This prints the new AMI ID. Options:
 - `--jobs`: the number of blocks uploaded in parallel, 16 by default.
 - `--no-tpm`: register the AMI without a NitroTPM, for arm64 bare metal instance types.
 
-Disk partitions often have a lot of empty space, e.g. an ESP sized for future updates. Only the blocks
-holding data are uploaded. This only works with unencrypted snapshots, so the import refuses to run while
-[EBS encryption by default](https://docs.aws.amazon.com/ebs/latest/userguide/encryption-by-default.html)
-is enabled for the region.
+### Importing from a release
 
-If the import fails while uploading, EC2 cancels the incomplete snapshot after an hour. If it fails while
-registering, the completed snapshot stays; the import printed its ID, delete it with
-`aws ec2 delete-snapshot`.
+For production, building and uploading should happen as two separate steps, with user and credential
+isolation. Add the AMI parameters to your release artifacts:
+
+```Starlark
+image.artifacts(
+    name = "release",
+    targets = [":image[raw.zst]", ":image[ami]", ...],
+)
+```
+
+Build the `image.artifacts` target in the build step, without AWS credentials.
+
+In the upload step, log into AWS and upload the AMI from these two files, with the same options as above.
+Do that from a trusted checkout such as your project's `main`, not an untrusted development branch.
+Verify files you did not just build yourself, e.g. downloaded from a release, against the release's
+signed checksums first.
+
+```sh
+tine buck run tine//image:ami-import -- /dir/image_1_x86-64.raw.zst /dir/image_1_x86-64.ami.json
+```
+
+This decompresses the disk on the fly during upload.
 
 ## Running an instance with credentials
 

docs/user/images.md

@@ -642,8 +642,8 @@ identity when copied out of the build, exactly matching the UKI's `<image_id>_<v
 The `[qcow2]`, `[raw.zst]` and `[vmdk]` subtargets re-encode the raw disk into a compact qcow2, a compressed
 raw, or a compressed (streamOptimized) VMDK for VMware and cloud importers on demand, each publishing one
 `image.DiskConversionInfo`; the encodings are reachable only through those
-subtargets, because one result cannot carry the same provider type twice. `[ami]` is the raw disk with the
-parameters an Amazon Machine Image is registered with, see [aws.md](aws.md). The target has no aggregate
+subtargets, because one result cannot carry the same provider type twice. `[ami]` publishes the parameters
+an Amazon Machine Image of the raw disk is registered with, see [aws.md](aws.md). The target has no aggregate
 bootable-image provider: it returns `image.ImageInfo`, `image.DdiInfo`, `image.InitrdInfo`,
 `image.ImageDirectoryInfo`, `image.UkiInfo`, and `image.PublishedInfo` independently. A disk that does not
 boot has no `[uki]`, `[initrd]`, `[boot]` or `[ami]` subtargets, returns no `image.InitrdInfo` or

examples/image-secureboot/BUCK

@@ -135,10 +135,12 @@ image.disk(
 # out trusting it. See docs/user/aws.md.
 box.sh_test(
     name = "uefi-data-test",
-    args = ["$(location :image[ami][uefi-data])", "$(location {}[cert])".format(_SIGNING)],
+    args = ["$(location :image[ami])", "$(location {}[cert])".format(_SIGNING)],
     box = "@tine//tools:aws.box",
     script = """
-virt-fw-vars --input "$1" --output-json /dev/stdout | python3 -c '
+uefi_data=$(mktemp)
+python3 -c 'import json, sys; print(json.load(open(sys.argv[1]))["UefiData"])' "$1" > "$uefi_data"
+virt-fw-vars --input "$uefi_data" --output-json /dev/stdout | python3 -c '
 import json, ssl, sys
 cert = ssl.PEM_cert_to_DER_cert(open(sys.argv[1]).read())
 variables = {v["name"]: bytes.fromhex(v["data"]) for v in json.load(sys.stdin)["variables"]}

examples/image/BUCK

@@ -308,23 +308,20 @@ done
     args = ["$(location :boot-demo.fedora)", "$(location :boot-demo.fedora[qcow2])", "$(location :boot-demo.fedora[vmdk])"],
 )
 
-# The AMI bundle carries the disk, and registers it for the image's architecture with a root volume of its
-# virtual_size.
+# The AMI parameters register the disk for the image's architecture with a root volume of its virtual_size.
 box.sh_test(
     name = "ami-test",
     labels = ["image"],
     box = "@tine//tools:dev.box",
     script = """
-cmp "$1/disk.raw" "$2"
 python3 -c '
 import json, sys
 register = json.load(open(sys.argv[1]))
 assert register["Architecture"] == sys.argv[2], register
 assert register["BlockDeviceMappings"][0]["Ebs"]["VolumeSize"] == 2, register
-' "$1/register.json" "$3"
+' "$@"
 """,
-    args = ["$(location :boot-demo.fedora[ami])", "$(location :boot-demo.fedora)"]
-    + architecture.select({name: [spellings.ec2] for name, spellings in ARCHITECTURES.items()}),
+    args = ["$(location :boot-demo.fedora[ami])"] + architecture.select({name: [spellings.ec2] for name, spellings in ARCHITECTURES.items()}),
 )
 
 # data-disk has no ESP partition, so image.disk builds it without systemd-boot, UKIs or an initrd.

image/ami.bzl

@@ -5,6 +5,7 @@
 
 load("//box:runtime.bzl", "BoxInfo", "box_run")
 load("//platforms:architecture.bzl", "ARCHITECTURES")
+load(":publish.bzl", "PublishedInfo")
 
 # What `RegisterImage` accepts in an image name besides letters and digits.
 _NAME_PUNCTUATION = "()[] ./-'@_"
@@ -40,13 +41,13 @@ def declare_ami(
     secure_boot_certificate: Artifact | None,
     virtual_size: str | None,
 ) -> list[Provider]:
-    """Declare the AMI bundle of a bootable disk, for its `[ami]` subtarget.
+    """Declare the AMI RegisterImage parameters of a bootable disk, for its `[ami]` subtarget.
 
-    The bundle holds the disk and the `RegisterImage` parameters, which `aws ec2 register-image
-    --cli-input-json` takes as they are once the import adds the snapshot. Running the subtarget
-    imports the bundle with `importer`. With a `secure_boot_certificate`, the bundle also holds the
-    UEFI variable store that enrolls it, also as the subtarget's `[uefi-data]`; `box` provides the
-    tool that writes it.
+    `aws ec2 register-image --cli-input-json` takes them as they are once the import adds the snapshot.
+    The subtarget publishes them as `<basename>.ami.json`, so that a release carries everything an import
+    needs besides the disk. Running the subtarget imports the disk with `importer`. With a
+    `secure_boot_certificate`, the parameters include the UEFI variable store that enrolls it; `box`
+    provides the tool that writes it.
     """
     register = {
         "Architecture": ARCHITECTURES[arch].ec2,
@@ -72,21 +73,45 @@ def declare_ami(
     if virtual_size != None:
         register["BlockDeviceMappings"][0]["Ebs"]["VolumeSize"] = _gib(virtual_size)
 
-    contents = {
-        # Links, not copies: the disk is sparse, as its partitions are mostly empty.
-        "disk.raw": disk,
-        "register.json": actions.write_json("ami-register.json", register),
-    }
-    sub_targets = {}
-    if secure_boot_certificate != None:
-        contents["uefi-data"] = _uefi_data(actions, box, secure_boot_certificate)
-        sub_targets["uefi-data"] = [DefaultInfo(default_output = contents["uefi-data"])]
-    bundle = actions.symlinked_dir("ami", contents)
+    name = "{}.ami.json".format(basename)
+    if secure_boot_certificate == None:
+        parameters = actions.write_json(name, register)
+    else:
+        parameters = actions.declare_output(name)
+        actions.dynamic_output_new(
+            _with_uefi_data(
+                out = parameters.as_output(),
+                register = register,
+                uefi_data = _uefi_data(actions, box, secure_boot_certificate),
+            ),
+        )
     return [
-        DefaultInfo(default_output = bundle, sub_targets = sub_targets),
-        RunInfo(args = cmd_args(importer, bundle)),
+        # Building an AMI builds the disk, as for any other format; only the parameters are published,
+        # as a release carries the disk already.
+        DefaultInfo(default_output = parameters, other_outputs = [disk]),
+        PublishedInfo(artifacts = {name: parameters}),
+        RunInfo(args = cmd_args(importer, disk, parameters)),
     ]
 
+def _with_uefi_data_impl(
+    actions: AnalysisActions,
+    out: OutputArtifact,
+    register: dict[str, typing.Any],
+    uefi_data: ArtifactValue,
+) -> list[Provider]:
+    # The store is base64 already, which is what RegisterImage takes.
+    actions.write_json(out, register | {"UefiData": uefi_data.read_string().strip()})
+    return []
+
+_with_uefi_data = dynamic_actions(
+    impl = _with_uefi_data_impl,
+    attrs = {
+        "out": dynattrs.output(),
+        "register": dynattrs.value(dict[str, typing.Any]),
+        "uefi_data": dynattrs.artifact_value(),
+    },
+)
+
 def _uefi_data(actions: AnalysisActions, box: BoxInfo, certificate: Artifact) -> Artifact:
     """The instance's initial UEFI variable store, which trusts `certificate` as PK, KEK, and db.
 

image/ami_import.py

@@ -1,14 +1,15 @@
 # SPDX-FileCopyrightText: Amutable GmbH <https://amutable.com/>
 # SPDX-License-Identifier: MPL-2.0
 
-"""Import an AMI bundle into EC2: upload its disk as an EBS snapshot, then register the image.
+"""Import a disk into EC2: upload it as an EBS snapshot, then register it with its AMI parameters.
 
 Every AWS call goes through the `aws` CLI, so this uses whatever login and region the CLI is
 configured with, and never handles credentials itself.
 """
 
 import argparse
 import base64
+import compression.zstd
 import concurrent.futures
 import errno
 import hashlib
@@ -24,13 +25,30 @@
 # The only block size the EBS direct APIs take.
 BLOCK_SIZE = 512 * 1024
 GIB = 1024**3
+# Biggest possible zstd frame header, for reading the decompressed size.
+ZSTD_FRAME_HEADER_MAX = 18
 
 
-def data_blocks(disk: Path) -> Iterator[tuple[int, bytes]]:
-    """Yield each block of `disk` that holds anything but zeros, with its index.
+def zstd_size(disk: Path) -> int | None:
+    """The decompressed size of a zstd-compressed `disk`, or None for an uncompressed one."""
+    with disk.open("rb") as f:
+        header = f.read(ZSTD_FRAME_HEADER_MAX)
+    try:
+        size = compression.zstd.get_frame_info(header).decompressed_size
+    except compression.zstd.ZstdError:
+        return None
+    if size is None:
+        sys.exit(f"ami-import: {disk} does not record its decompressed size")
+    return size
 
-    Holes are never read, so a sparse disk costs its allocated size rather than its apparent one.
-    """
+
+def disk_size(disk: Path) -> int:
+    """The size of `disk`, decompressed if it is zstd-compressed."""
+    size = zstd_size(disk)
+    return disk.stat().st_size if size is None else size
+
+
+def _allocated_blocks(disk: Path) -> Iterator[tuple[int, bytes]]:
     with disk.open("rb") as f:
         fd = f.fileno()
         size = os.fstat(fd).st_size
@@ -45,12 +63,34 @@ def data_blocks(disk: Path) -> Iterator[tuple[int, bytes]]:
                 raise
             end = os.lseek(fd, start, os.SEEK_HOLE)
             for index in range(start // BLOCK_SIZE, (end + BLOCK_SIZE - 1) // BLOCK_SIZE):
-                block = os.pread(fd, BLOCK_SIZE, index * BLOCK_SIZE).ljust(BLOCK_SIZE, b"\0")
-                if block.count(0) != BLOCK_SIZE:
-                    yield index, block
+                yield index, os.pread(fd, BLOCK_SIZE, index * BLOCK_SIZE).ljust(BLOCK_SIZE, b"\0")
             offset = (end + BLOCK_SIZE - 1) // BLOCK_SIZE * BLOCK_SIZE
 
 
+def _decompressed_blocks(disk: Path, size: int) -> Iterator[tuple[int, bytes]]:
+    with compression.zstd.open(disk) as f:
+        index = 0
+        while block := f.read(BLOCK_SIZE):
+            yield index, block.ljust(BLOCK_SIZE, b"\0")
+            index += 1
+        # The snapshot was started with the size the first frame records.
+        if f.tell() != size:
+            sys.exit(f"ami-import: {disk} decompresses to {f.tell()} bytes, not the {size} its header says")
+
+
+def data_blocks(disk: Path) -> Iterator[tuple[int, bytes]]:
+    """Yield each block of `disk` that holds anything but zeros, with its index.
+
+    A zstd-compressed disk is decompressed on the fly. Holes of an uncompressed one are never read,
+    so a sparse disk costs its allocated size rather than its apparent one.
+    """
+    size = zstd_size(disk)
+    blocks = _allocated_blocks(disk) if size is None else _decompressed_blocks(disk, size)
+    for index, block in blocks:
+        if block.count(0) != BLOCK_SIZE:
+            yield index, block
+
+
 def checksum(block: bytes) -> str:
     """The per-block checksum `PutSnapshotBlock` requires: base64 of its SHA256."""
     return base64.b64encode(hashlib.sha256(block).digest()).decode()
@@ -83,7 +123,7 @@ def upload(aws: Aws, disk: Path, name: str, jobs: int) -> str:
     if aws("ec2", "get-ebs-encryption-by-default")["EbsEncryptionByDefault"]:
         sys.exit("ami-import: EBS encryption by default is enabled for this account and region; disable it")
 
-    size = (disk.stat().st_size + GIB - 1) // GIB
+    size = (disk_size(disk) + GIB - 1) // GIB
     snapshot = str(
         aws(
             "ebs",
@@ -137,18 +177,13 @@ def put(index: int, block: bytes, scratch: Path) -> None:
     return snapshot
 
 
-def register(aws: Aws, bundle: Path, name: str, snapshot: str, tpm: bool) -> str:
-    """Register the image from the bundle's parameters, booting `snapshot`, and return its ID."""
-    parameters = json.loads((bundle / "register.json").read_text())
+def register(aws: Aws, parameters: dict[str, Any], name: str, snapshot: str, tpm: bool) -> str:
+    """Register the image with `parameters`, booting `snapshot`, and return its ID."""
     if not tpm:
         del parameters["TpmSupport"]
     parameters["Name"] = name
     parameters["BlockDeviceMappings"][0]["Ebs"]["SnapshotId"] = snapshot
     parameters["TagSpecifications"] = [{"ResourceType": "image", "Tags": [{"Key": "Name", "Value": name}]}]
-    # Only a disk signed for Secure Boot has one, see `_uefi_data` in ami.bzl.
-    uefi_data = bundle / "uefi-data"
-    if uefi_data.exists():
-        parameters["UefiData"] = uefi_data.read_text().strip()
     with tempfile.NamedTemporaryFile("w", suffix=".json") as f:
         json.dump(parameters, f)
         f.flush()
@@ -157,8 +192,9 @@ def register(aws: Aws, bundle: Path, name: str, snapshot: str, tpm: bool) -> str
 
 def main() -> None:
     parser = argparse.ArgumentParser(description=__doc__)
-    parser.add_argument("bundle", type=Path, help="the AMI bundle to import")
-    parser.add_argument("--name", help="the image name, instead of the bundle's; EC2 requires it unique")
+    parser.add_argument("disk", type=Path, help="the .raw[.zst] disk to upload")
+    parser.add_argument("parameters", type=Path, help="the disk's RegisterImage parameters (.ami.json)")
+    parser.add_argument("--name", help="the image name, instead of the parameters'; EC2 requires it unique")
     parser.add_argument("--region", help="the region to import into, instead of the CLI's configured one")
     parser.add_argument("--jobs", type=int, default=16, help="blocks to upload in parallel")
     parser.add_argument(
@@ -170,13 +206,14 @@ def main() -> None:
     args = parser.parse_args()
 
     aws = Aws(args.region)
-    name = args.name or json.loads((args.bundle / "register.json").read_text())["Name"]
+    parameters = json.loads(args.parameters.read_text())
+    name = args.name or parameters["Name"]
     # Names are unique per account and region; find out before spending the upload.
     if aws("ec2", "describe-images", "--owners=self", f"--filters=Name=name,Values={name}")["Images"]:
         sys.exit(f"ami-import: an image named {name!r} exists already; pass another --name")
 
-    snapshot = upload(aws, args.bundle / "disk.raw", name, args.jobs)
-    print(register(aws, args.bundle, name, snapshot, args.tpm))
+    snapshot = upload(aws, args.disk, name, args.jobs)
+    print(register(aws, parameters, name, snapshot, args.tpm))
 
 
 if __name__ == "__main__":