@@ -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,28 @@ 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.
+
+```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
 
 
@@ -9,6 +9,7 @@
 
 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",
@@ -152,7 +192,7 @@ def register(aws: Aws, parameters: dict[str, Any], name: str, snapshot: str, tpm
 
 def main() -> None:
     parser = argparse.ArgumentParser(description=__doc__)
-    parser.add_argument("disk", type=Path, help="the .raw disk to upload")
+    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")
 
@@ -7,10 +7,12 @@
 """
 
 import base64
+import compression.zstd
 import hashlib
 import tempfile
 import unittest
 from pathlib import Path
+from typing import override
 
 import ami_import
 from ami_import import BLOCK_SIZE
@@ -56,6 +58,37 @@ def test_pads_a_short_last_block(self) -> None:
     def test_empty_disk_has_no_blocks(self) -> None:
         self.assertEqual(list(ami_import.data_blocks(self.disk(3 * BLOCK_SIZE, {}))), [])
 
+    def test_size(self) -> None:
+        self.assertEqual(ami_import.disk_size(self.disk(BLOCK_SIZE + 10, {3: b"x"})), BLOCK_SIZE + 10)
+
+
+class CompressedDataBlocksTest(DataBlocksTest):
+    """The same disks, zstd-compressed as a release's `.raw.zst`, but still named `.raw`.
+
+    The name must not matter, only the content.
+    """
+
+    @override
+    def disk(self, size: int, writes: dict[int, bytes]) -> Path:
+        disk = super().disk(size, writes)
+        # as `zstd` writes it from a file, with the decompressed size in the frame header
+        disk.write_bytes(compression.zstd.compress(disk.read_bytes()))
+        return disk
+
+    def test_more_than_the_header_says(self) -> None:
+        disk = self.disk(BLOCK_SIZE, {0: b"x"})
+        with disk.open("ab") as f:
+            f.write(compression.zstd.compress(b"y"))
+        with self.assertRaisesRegex(SystemExit, "decompresses to 524289 bytes, not the 524288"):
+            list(ami_import.data_blocks(disk))
+
+    def test_no_size_in_header(self) -> None:
+        disk = self.disk(BLOCK_SIZE, {})
+        with compression.zstd.ZstdFile(disk, "w") as f:
+            f.write(b"x")
+        with self.assertRaisesRegex(SystemExit, "does not record its decompressed size"):
+            ami_import.disk_size(disk)
+
 
 class ChecksumTest(unittest.TestCase):
     def test_is_base64_sha256(self) -> None: