Build AMIs on temporary EC2 builders

README.md

@@ -14,20 +14,33 @@ flakes, and systemd ordering so the Nix daemon waits for `/nix`.
 
 ## Publish
 
-The smallest supported flow is:
+The normal flow builds on a temporary EC2 instance, so your local machine does
+not need to be the same architecture as the target AMI:
 
 ```bash
 cd tofu/ami-builder
-nix build .#packages.aarch64-linux.rootImage
+./publish-ami.sh --builder-instance-type t4g.large
-./publish-ami.sh --bucket s3://my-vm-import-bucket
 ```
 
-`publish-ami.sh` builds/imports the declarative NixOS root image, boots a
+`publish-ami.sh` applies this OpenTofu project to create a temporary builder
-temporary seeder instance once to create/export the ZFS `/nix` volume, creates
+instance and `/nix` volume. The builder user data installs a ZFS-capable NixOS
-the final AMI from the stopped instance, and writes `artifact.json`.
+boot generation, reboots once, seeds the ZFS `/nix` volume, installs the final
+boot generation that expects `/nix`, and powers off.
 
-The OpenTofu project in this directory owns only the final AMI registration
+After the instance is stopped, `publish-ami.sh` creates the final AMI from the
-metadata. Its state is separate from any runtime Quixos deployment state.
+stopped instance, tags the AMI and snapshots, and writes `artifact.json`.
+
+The OpenTofu state in this directory owns only temporary builder resources.
+The published AMI and its snapshots are created by the script and are not in
+OpenTofu state, so `tofu destroy` here can clean up the builder without
+deregistering the AMI.
+
+Pass extra OpenTofu variables after `--`:
+
+```bash
+./publish-ami.sh --builder-instance-type t4g.large -- \
+  -var base_nixos_ami_id=ami-...
+```
 
 ## Artifact Metadata
 

main.tf

@@ -1,45 +1,155 @@
+data "aws_vpc" "default" {
+  count   = var.vpc_id == null ? 1 : 0
+  default = true
+}
+
+data "aws_subnets" "default" {
+  count = var.subnet_id == null ? 1 : 0
+
+  filter {
+    name   = "vpc-id"
+    values = [local.vpc_id]
+  }
+}
+
+data "aws_subnet" "selected" {
+  id = local.subnet_id
+}
+
+data "aws_ami" "nixos" {
+  count       = var.base_nixos_ami_id == null ? 1 : 0
+  most_recent = true
+  owners      = var.base_nixos_ami_owners
+
+  filter {
+    name   = "name"
+    values = [local.base_nixos_ami_name_pattern]
+  }
+
+  filter {
+    name   = "architecture"
+    values = [local.aws_architecture]
+  }
+
+  filter {
+    name   = "virtualization-type"
+    values = ["hvm"]
+  }
+}
+
 locals {
-  resource_name = "${var.name}-${var.architecture}-${substr(var.input_hash, 0, 12)}"
+  vpc_id           = var.vpc_id == null ? data.aws_vpc.default[0].id : var.vpc_id
+  subnet_id        = var.subnet_id == null ? data.aws_subnets.default[0].ids[0] : var.subnet_id
+  aws_architecture = var.system == "aarch64-linux" ? "arm64" : "x86_64"
+  base_nixos_ami_name_pattern = (
+    var.base_nixos_ami_name_pattern == null
+    ? "nixos/*-${var.system}"
+    : var.base_nixos_ami_name_pattern
+  )
+  base_nixos_ami_id = var.base_nixos_ami_id == null ? data.aws_ami.nixos[0].id : var.base_nixos_ami_id
+  resource_name     = "${var.name}-builder-${local.aws_architecture}-${substr(var.input_hash, 0, 12)}"
+  nix_device_id     = replace(aws_ebs_volume.nix.id, "-", "")
+  nix_device        = "/dev/disk/by-id/nvme-Amazon_Elastic_Block_Store_${local.nix_device_id}"
   tags = {
     Name          = local.resource_name
     Project       = var.name
-    Architecture       = var.architecture
+    Role          = "ami-builder"
+    Architecture  = local.aws_architecture
+    System        = var.system
     InputHash     = var.input_hash
     BuilderGitRev = var.builder_git_rev
     FlakeLockRev  = var.flake_lock_rev
-    RootSnapshotId     = var.root_snapshot_id
-    NixSnapshotId      = var.nix_snapshot_id
-    NixosSystemClosure = var.nixos_system_closure
   }
 }
 
-resource "aws_ami" "matched" {
+resource "aws_iam_role" "builder" {
-  name                = local.resource_name
+  name = "${local.resource_name}-instance"
-  description         = "NixOS EC2 AMI with matched root and ZFS /nix snapshots"
-  architecture        = var.architecture
-  virtualization_type = "hvm"
-  root_device_name    = var.root_device_name
-  ena_support         = true
 
-  ebs_block_device {
+  assume_role_policy = jsonencode({
-    device_name           = var.root_device_name
+    Version = "2012-10-17"
-    snapshot_id           = var.root_snapshot_id
+    Statement = [
-    volume_size           = var.root_volume_size_gb
+      {
-    volume_type           = "gp3"
+        Action = "sts:AssumeRole"
-    delete_on_termination = true
+        Effect = "Allow"
+        Principal = {
+          Service = "ec2.amazonaws.com"
         }
+      }
+    ]
+  })
+}
 
-  ebs_block_device {
+resource "aws_iam_role_policy_attachment" "ssm" {
-    device_name           = var.nix_device_name
+  role       = aws_iam_role.builder.name
-    snapshot_id           = var.nix_snapshot_id
+  policy_arn = "arn:aws:iam::aws:policy/AmazonSSMManagedInstanceCore"
-    volume_size           = var.nix_volume_size_gb
+}
-    volume_type           = "gp3"
+
-    delete_on_termination = true
+resource "aws_iam_instance_profile" "builder" {
+  name = "${local.resource_name}-instance"
+  role = aws_iam_role.builder.name
+}
+
+resource "aws_security_group" "builder" {
+  name        = "${local.resource_name}-instance"
+  description = "Temporary NixOS ZFS AMI builder egress"
+  vpc_id      = local.vpc_id
+
+  egress {
+    from_port   = 0
+    to_port     = 0
+    protocol    = "-1"
+    cidr_blocks = ["0.0.0.0/0"]
   }
 
   tags = local.tags
-
+}
-  lifecycle {
+
-    prevent_destroy = true
+resource "aws_key_pair" "admin" {
-  }
+  count      = var.ssh_public_key == null ? 0 : 1
+  key_name   = "${local.resource_name}-admin"
+  public_key = var.ssh_public_key
+}
+
+resource "aws_ebs_volume" "nix" {
+  availability_zone = data.aws_subnet.selected.availability_zone
+  size              = var.nix_volume_size_gb
+  type              = "gp3"
+
+  tags = merge(local.tags, {
+    Name = "${local.resource_name}-nix"
+  })
+}
+
+resource "aws_instance" "builder" {
+  ami                         = local.base_nixos_ami_id
+  instance_type               = var.builder_instance_type
+  subnet_id                   = local.subnet_id
+  vpc_security_group_ids      = [aws_security_group.builder.id]
+  iam_instance_profile        = aws_iam_instance_profile.builder.name
+  associate_public_ip_address = true
+  key_name                    = var.ssh_public_key == null ? null : aws_key_pair.admin[0].key_name
+  user_data_replace_on_change = true
+
+  user_data = templatefile("${path.module}/user-data.sh.tftpl", {
+    builder_git_rev = var.builder_git_rev
+    flake_lock_rev  = var.flake_lock_rev
+    input_hash      = var.input_hash
+    nix_device      = local.nix_device
+    nix_pool        = var.nix_pool
+    ssh_keys        = jsonencode(compact([var.ssh_public_key]))
+    system          = var.system
+  })
+
+  root_block_device {
+    volume_size = var.root_volume_size_gb
+    volume_type = "gp3"
+  }
+
+  tags = local.tags
+}
+
+resource "aws_volume_attachment" "nix" {
+  device_name = var.nix_device_name
+  volume_id   = aws_ebs_volume.nix.id
+  instance_id = aws_instance.builder.id
 }

outputs.tf

@@ -1,9 +1,13 @@
-output "ami_id" {
+output "builder_instance_id" {
-  value = aws_ami.matched.id
+  value = aws_instance.builder.id
 }
 
 output "architecture" {
-  value = var.architecture
+  value = local.aws_architecture
+}
+
+output "system" {
+  value = var.system
 }
 
 output "input_hash" {
@@ -18,14 +22,14 @@ output "flake_lock_rev" {
   value = var.flake_lock_rev
 }
 
-output "root_snapshot_id" {
+output "nix_device_name" {
-  value = var.root_snapshot_id
+  value = var.nix_device_name
 }
 
-output "nix_snapshot_id" {
+output "nix_volume_id" {
-  value = var.nix_snapshot_id
+  value = aws_ebs_volume.nix.id
 }
 
-output "nixos_system_closure" {
+output "nix_pool" {
-  value = var.nixos_system_closure
+  value = var.nix_pool
 }

publish-ami.sh

@@ -2,70 +2,59 @@
 set -euo pipefail
 
 aws_region="${AWS_REGION:-us-west-2}"
-architecture="arm64"
+system="aarch64-linux"
-bucket=""
+builder_instance_type="t4g.large"
-subnet_id=""
-security_group_id=""
-instance_profile_name=""
-instance_type="t4g.small"
-nix_volume_size_gb="50"
 name="nixos-zfs-ec2"
+destroy_builder=false
+extra_tofu_args=()
 
 usage() {
   cat >&2 <<'EOF'
-Usage: publish-ami.sh --bucket s3://bucket[/prefix] [options]
+Usage: publish-ami.sh [options] [-- <extra tofu apply args>]
 
 Options:
-  --architecture <arm64|x86_64>       AMI architecture (default: arm64)
+  --system <aarch64-linux|x86_64-linux>
+                              Target NixOS system (default: aarch64-linux)
   --aws-region <region>       AWS region (default: AWS_REGION or us-west-2)
-  --bucket <s3-uri>                   S3 URI used by EC2 VM import
+  --builder-instance-type <t> Temporary builder EC2 type (default: t4g.large)
-  --subnet-id <subnet>                Subnet for the temporary seeder instance
-  --security-group-id <sg>            Security group for the temporary seeder instance
-  --instance-profile-name <profile>   Optional IAM instance profile for the seeder
-  --instance-type <type>              Seeder instance type (default: t4g.small)
-  --nix-volume-size-gb <gb>           /nix volume size (default: 50)
   --name <name>               AMI name prefix
+  --destroy-builder           Destroy temporary builder resources after publish
+  -h, --help                  Show this help
+
+Any arguments after `--` are passed through to `tofu apply`.
+Use that for values like:
+
+  -- -var base_nixos_ami_id=ami-... -var subnet_id=subnet-...
 EOF
 }
 
 while [ "$#" -gt 0 ]; do
   case "$1" in
-    --architecture)
+    --system)
-      architecture="$2"
+      system="$2"
       shift 2
       ;;
     --aws-region)
       aws_region="$2"
       shift 2
       ;;
-    --bucket)
+    --builder-instance-type)
-      bucket="$2"
+      builder_instance_type="$2"
-      shift 2
-      ;;
-    --subnet-id)
-      subnet_id="$2"
-      shift 2
-      ;;
-    --security-group-id)
-      security_group_id="$2"
-      shift 2
-      ;;
-    --instance-profile-name)
-      instance_profile_name="$2"
-      shift 2
-      ;;
-    --instance-type)
-      instance_type="$2"
-      shift 2
-      ;;
-    --nix-volume-size-gb)
-      nix_volume_size_gb="$2"
       shift 2
       ;;
     --name)
       name="$2"
       shift 2
       ;;
+    --destroy-builder)
+      destroy_builder=true
+      shift
+      ;;
+    --)
+      shift
+      extra_tofu_args=("$@")
+      break
+      ;;
     -h|--help)
       usage
       exit 0
@@ -78,26 +67,15 @@ while [ "$#" -gt 0 ]; do
   esac
 done
 
-if [ -z "$bucket" ]; then
+case "$system" in
-  usage
+  aarch64-linux)
-  exit 1
+    architecture="arm64"
-fi
-
-if [ -z "$instance_profile_name" ]; then
-  echo "Missing required --instance-profile-name for SSM seeding" >&2
-  usage
-  exit 1
-fi
-
-case "$architecture" in
-  arm64)
-    nix_system="aarch64-linux"
     ;;
-  x86_64)
+  x86_64-linux)
-    nix_system="x86_64-linux"
+    architecture="x86_64"
     ;;
   *)
-    echo "Unsupported architecture: $architecture" >&2
+    echo "Unsupported system: $system" >&2
     exit 1
     ;;
 esac
@@ -112,7 +90,6 @@ require() {
 require aws
 require git
 require jq
-require nix
 require sha256sum
 require tofu
 
@@ -124,250 +101,213 @@ fi
 
 builder_git_rev="$(git rev-parse HEAD)"
 flake_lock_rev="$(jq -r '.nodes.nixpkgs.locked.rev // .nodes.nixpkgs.locked.narHash // "unknown"' "$repo_root/flake.lock")"
-input_hash="$(printf '%s\n%s\n%s\n' "$builder_git_rev" "$flake_lock_rev" "$architecture" | sha256sum | awk '{ print $1 }')"
+input_hash="$(printf '%s\n%s\n%s\n' "$builder_git_rev" "$flake_lock_rev" "$system" | sha256sum | awk '{ print $1 }')"
-
+ami_name="${name}-${architecture}-${input_hash:0:12}"
-root_image="$(nix build --print-out-paths ".#packages.${nix_system}.rootImage")"
-nixos_system_closure="$(nix build --print-out-paths ".#packages.${nix_system}.systemClosure")"
-
-tmpdir="$(mktemp -d)"
-cleanup() {
-  rm -rf "$tmpdir"
-}
-trap cleanup EXIT
-
-bucket_no_scheme="${bucket#s3://}"
-bucket_name="${bucket_no_scheme%%/*}"
-bucket_prefix="${bucket_no_scheme#"$bucket_name"}"
-bucket_prefix="${bucket_prefix#/}"
-s3_key_prefix="${bucket_prefix:+$bucket_prefix/}${name}/${architecture}/${input_hash}"
-
-image_file="$(find "$root_image" -type f \( -name '*.vhd' -o -name '*.raw' \) | head -n 1)"
-if [ -z "$image_file" ]; then
-  echo "Could not locate built root image under $root_image" >&2
-  exit 1
-fi
-
-case "$image_file" in
-  *.vhd)
-    import_format="vhd"
-    ;;
-  *.raw)
-    import_format="raw"
-    ;;
-  *)
-    echo "Unsupported root image format: $image_file" >&2
-    exit 1
-    ;;
-esac
-
-aws s3 cp "$image_file" "s3://${bucket_name}/${s3_key_prefix}/root-image"
-
-container_json="$tmpdir/container.json"
-jq -n \
-  --arg bucket "$bucket_name" \
-  --arg key "${s3_key_prefix}/root-image" \
-  --arg format "$import_format" \
-  '[{Description:"nixos-zfs-ec2-root",Format:$format,UserBucket:{S3Bucket:$bucket,S3Key:$key}}]' \
-  > "$container_json"
-
-import_task_id="$(
-  aws ec2 import-image \
-    --region "$aws_region" \
-    --architecture "$architecture" \
-    --platform Linux \
-    --disk-containers "file://$container_json" \
-    --query ImportTaskId \
-    --output text
-)"
-
-while true; do
-  import_json="$(aws ec2 describe-import-image-tasks --region "$aws_region" --import-task-ids "$import_task_id")"
-  status="$(jq -r '.ImportImageTasks[0].Status' <<<"$import_json")"
-  case "$status" in
-    completed)
-      root_ami_id="$(jq -r '.ImportImageTasks[0].ImageId' <<<"$import_json")"
-      break
-      ;;
-    deleted|deleting)
-      echo "$import_json" >&2
-      exit 1
-      ;;
-    *)
-      status_message="$(jq -r '.ImportImageTasks[0].StatusMessage // ""' <<<"$import_json")"
-      echo "import-image: $status $status_message" >&2
-      sleep 20
-      ;;
-  esac
-done
-
-root_snapshot_id="$(
-  aws ec2 describe-images \
-    --region "$aws_region" \
-    --image-ids "$root_ami_id" \
-    --query 'Images[0].BlockDeviceMappings[?Ebs != `null`][0].Ebs.SnapshotId' \
-    --output text
-)"
-
-aws ec2 create-tags \
-  --region "$aws_region" \
-  --resources "$root_snapshot_id" \
-  --tags \
-    "Key=Name,Value=${name}-root-${architecture}-${input_hash:0:12}" \
-    "Key=InputHash,Value=${input_hash}" \
-    "Key=BuilderGitRev,Value=${builder_git_rev}" \
-    "Key=FlakeLockRev,Value=${flake_lock_rev}" \
-    "Key=NixosSystemClosure,Value=${nixos_system_closure}"
-
-run_args=(
-  --region "$aws_region"
-  --image-id "$root_ami_id"
-  --instance-type "$instance_type"
-  --block-device-mappings "DeviceName=/dev/sdf,Ebs={VolumeSize=${nix_volume_size_gb},VolumeType=gp3,DeleteOnTermination=true}"
-  --tag-specifications "ResourceType=instance,Tags=[{Key=Name,Value=${name}-seed-${input_hash:0:12}}]"
-  --query 'Instances[0].InstanceId'
-  --output text
-)
-
-if [ -n "$subnet_id" ]; then
-  run_args+=(--subnet-id "$subnet_id")
-fi
-if [ -n "$security_group_id" ]; then
-  run_args+=(--security-group-ids "$security_group_id")
-fi
-if [ -n "$instance_profile_name" ]; then
-  run_args+=(--iam-instance-profile "Name=$instance_profile_name")
-fi
-
-instance_id="$(aws ec2 run-instances "${run_args[@]}")"
-aws ec2 wait instance-running --region "$aws_region" --instance-ids "$instance_id"
-aws ec2 wait instance-status-ok --region "$aws_region" --instance-ids "$instance_id"
-
-cat > "$tmpdir/seed-nix.sh" <<'EOS'
-#!/usr/bin/env bash
-set -euo pipefail
-
-pool="nixos-zfs-nix"
-device=""
-for candidate in /dev/disk/by-id/nvme-Amazon_Elastic_Block_Store_*; do
-  if [ -e "$candidate" ] && ! findmnt -n -S "$candidate" >/dev/null 2>&1; then
-    device="$candidate"
-  fi
-done
-
-if [ -z "$device" ]; then
-  echo "Unable to find unattached EBS device for /nix seed" >&2
-  exit 1
-fi
-
-zpool create \
-  -f \
-  -o ashift=12 \
-  -O acltype=posixacl \
-  -O atime=off \
-  -O compression=zstd \
-  -O mountpoint=none \
-  -O xattr=sa \
-  "$pool" "$device"
-
-zfs create -o mountpoint=legacy "$pool/nix"
-mkdir -p /mnt/nixos-zfs-nix
-mount -t zfs "$pool/nix" /mnt/nixos-zfs-nix
-rsync -aHAX --numeric-ids \
-  --exclude='/var/nix/daemon-socket/*' \
-  /nix/ /mnt/nixos-zfs-nix/
-rm -f /mnt/nixos-zfs-nix/var/nix/daemon-socket/socket
-umount /mnt/nixos-zfs-nix
-zpool export "$pool"
-EOS
-
-jq -n --rawfile commands "$tmpdir/seed-nix.sh" '{commands: [$commands]}' > "$tmpdir/ssm-parameters.json"
-
-command_id="$(
-  aws ssm send-command \
-    --region "$aws_region" \
-    --instance-ids "$instance_id" \
-    --document-name AWS-RunShellScript \
-    --parameters "file://$tmpdir/ssm-parameters.json" \
-    --query Command.CommandId \
-    --output text
-)"
-
-while true; do
-  invocation_json="$(aws ssm get-command-invocation --region "$aws_region" --command-id "$command_id" --instance-id "$instance_id" 2>/dev/null || true)"
-  status="$(jq -r '.Status // "Pending"' <<<"$invocation_json")"
-  case "$status" in
-    Success)
-      break
-      ;;
-    Failed|Cancelled|TimedOut)
-      echo "$invocation_json" >&2
-      exit 1
-      ;;
-    *)
-      echo "seed /nix: $status" >&2
-      sleep 10
-      ;;
-  esac
-done
-
-aws ec2 stop-instances --region "$aws_region" --instance-ids "$instance_id" >/dev/null
-aws ec2 wait instance-stopped --region "$aws_region" --instance-ids "$instance_id"
-
-nix_volume_id="$(
-  aws ec2 describe-instances \
-    --region "$aws_region" \
-    --instance-ids "$instance_id" \
-    --query 'Reservations[0].Instances[0].BlockDeviceMappings[?DeviceName==`/dev/sdf`].Ebs.VolumeId | [0]' \
-    --output text
-)"
-
-nix_snapshot_id="$(
-  aws ec2 create-snapshot \
-    --region "$aws_region" \
-    --volume-id "$nix_volume_id" \
-    --description "${name} ${architecture} ZFS /nix ${input_hash}" \
-    --tag-specifications "ResourceType=snapshot,Tags=[{Key=Name,Value=${name}-nix-${architecture}-${input_hash:0:12}},{Key=InputHash,Value=${input_hash}},{Key=BuilderGitRev,Value=${builder_git_rev}}]" \
-    --query SnapshotId \
-    --output text
-)"
-aws ec2 wait snapshot-completed --region "$aws_region" --snapshot-ids "$nix_snapshot_id"
 
 tofu init
 tofu apply \
   -var "aws_region=$aws_region" \
   -var "name=$name" \
-  -var "architecture=$architecture" \
+  -var "system=$system" \
-  -var "root_snapshot_id=$root_snapshot_id" \
+  -var "builder_instance_type=$builder_instance_type" \
-  -var "nix_snapshot_id=$nix_snapshot_id" \
   -var "input_hash=$input_hash" \
   -var "builder_git_rev=$builder_git_rev" \
   -var "flake_lock_rev=$flake_lock_rev" \
-  -var "nixos_system_closure=$nixos_system_closure" \
+  "${extra_tofu_args[@]}"
-  -var "nix_volume_size_gb=$nix_volume_size_gb"
 
-ami_id="$(tofu output -raw ami_id)"
+builder_instance_id="$(tofu output -raw builder_instance_id)"
+nix_volume_id="$(tofu output -raw nix_volume_id)"
+nix_device_name="$(tofu output -raw nix_device_name)"
+nix_pool="$(tofu output -raw nix_pool)"
 
-aws ec2 terminate-instances --region "$aws_region" --instance-ids "$instance_id" >/dev/null
+echo "Waiting for builder instance to stop: $builder_instance_id" >&2
-aws ec2 deregister-image --region "$aws_region" --image-id "$root_ami_id" >/dev/null || true
+aws ec2 wait instance-stopped --region "$aws_region" --instance-ids "$builder_instance_id"
+
+capture_image_id="$(
+  aws ec2 create-image \
+    --region "$aws_region" \
+    --instance-id "$builder_instance_id" \
+    --name "${ami_name}-capture" \
+    --description "Temporary capture for ${ami_name}" \
+    --no-reboot \
+    --tag-specifications "ResourceType=image,Tags=[{Key=Name,Value=${ami_name}-capture},{Key=Project,Value=${name}},{Key=Role,Value=ami-capture},{Key=InputHash,Value=${input_hash}}]" \
+    --query ImageId \
+    --output text
+)"
+
+aws ec2 wait image-available --region "$aws_region" --image-ids "$capture_image_id"
+
+capture_image_json="$(aws ec2 describe-images --region "$aws_region" --image-ids "$capture_image_id")"
+root_snapshot_id="$(
+  jq -r --arg nixDevice "$nix_device_name" '
+    .Images[0].BlockDeviceMappings[]
+    | select(.DeviceName != $nixDevice and .Ebs != null)
+    | .Ebs.SnapshotId
+    ' <<<"$capture_image_json" | head -n 1
+)"
+root_device_name="$(
+  jq -r --arg nixDevice "$nix_device_name" '
+    .Images[0].BlockDeviceMappings[]
+    | select(.DeviceName != $nixDevice and .Ebs != null)
+    | .DeviceName
+  ' <<<"$capture_image_json" | head -n 1
+)"
+root_volume_size="$(
+  jq -r --arg device "$root_device_name" '
+    .Images[0].BlockDeviceMappings[]
+    | select(.DeviceName == $device)
+    | .Ebs.VolumeSize
+  ' <<<"$capture_image_json"
+)"
+nix_snapshot_id="$(
+  jq -r --arg device "$nix_device_name" '
+    .Images[0].BlockDeviceMappings[]
+    | select(.DeviceName == $device)
+    | .Ebs.SnapshotId
+  ' <<<"$capture_image_json"
+)"
+nix_volume_size="$(
+  jq -r --arg device "$nix_device_name" '
+    .Images[0].BlockDeviceMappings[]
+    | select(.DeviceName == $device)
+    | .Ebs.VolumeSize
+  ' <<<"$capture_image_json"
+)"
+
+if [ -z "$root_snapshot_id" ] || [ "$root_snapshot_id" = "null" ]; then
+  echo "Unable to determine root snapshot for $capture_image_id" >&2
+  exit 1
+fi
+
+if [ -z "$nix_snapshot_id" ] || [ "$nix_snapshot_id" = "null" ]; then
+  echo "Unable to determine /nix snapshot for $capture_image_id" >&2
+  exit 1
+fi
+
+block_device_mappings="$(
+  jq -n \
+    --arg rootDevice "$root_device_name" \
+    --arg rootSnapshot "$root_snapshot_id" \
+    --argjson rootSize "$root_volume_size" \
+    --arg nixDevice "$nix_device_name" \
+    --arg nixSnapshot "$nix_snapshot_id" \
+    --argjson nixSize "$nix_volume_size" \
+    '[
+      {
+        DeviceName: $rootDevice,
+        Ebs: {
+          SnapshotId: $rootSnapshot,
+          VolumeSize: $rootSize,
+          VolumeType: "gp3",
+          DeleteOnTermination: true
+        }
+      },
+      {
+        DeviceName: $nixDevice,
+        Ebs: {
+          SnapshotId: $nixSnapshot,
+          VolumeSize: $nixSize,
+          VolumeType: "gp3",
+          DeleteOnTermination: true
+        }
+      }
+    ]'
+)"
+
+image_id="$(
+  aws ec2 register-image \
+    --region "$aws_region" \
+    --name "$ami_name" \
+    --description "NixOS EC2 AMI with matched root and ZFS /nix snapshots" \
+    --architecture "$architecture" \
+    --virtualization-type hvm \
+    --root-device-name "$root_device_name" \
+    --ena-support \
+    --block-device-mappings "$block_device_mappings" \
+    --query ImageId \
+    --output text
+)"
+
+aws ec2 wait image-available --region "$aws_region" --image-ids "$image_id"
+
+aws ec2 create-tags \
+  --region "$aws_region" \
+  --resources "$root_snapshot_id" \
+  --tags \
+    "Key=Name,Value=${ami_name}-root" \
+    "Key=Project,Value=${name}" \
+    "Key=Architecture,Value=${architecture}" \
+    "Key=System,Value=${system}" \
+    "Key=InputHash,Value=${input_hash}" \
+    "Key=BuilderGitRev,Value=${builder_git_rev}" \
+    "Key=FlakeLockRev,Value=${flake_lock_rev}"
+
+aws ec2 create-tags \
+  --region "$aws_region" \
+  --resources "$nix_snapshot_id" "$nix_volume_id" \
+  --tags \
+    "Key=Name,Value=${ami_name}-nix" \
+    "Key=Project,Value=${name}" \
+    "Key=Architecture,Value=${architecture}" \
+    "Key=System,Value=${system}" \
+    "Key=InputHash,Value=${input_hash}" \
+    "Key=BuilderGitRev,Value=${builder_git_rev}" \
+    "Key=FlakeLockRev,Value=${flake_lock_rev}" \
+    "Key=NixPool,Value=${nix_pool}"
+
+aws ec2 create-tags \
+  --region "$aws_region" \
+  --resources "$image_id" \
+  --tags \
+    "Key=Name,Value=${ami_name}" \
+    "Key=Project,Value=${name}" \
+    "Key=Architecture,Value=${architecture}" \
+    "Key=System,Value=${system}" \
+    "Key=InputHash,Value=${input_hash}" \
+    "Key=BuilderGitRev,Value=${builder_git_rev}" \
+    "Key=FlakeLockRev,Value=${flake_lock_rev}" \
+    "Key=NixPool,Value=${nix_pool}" \
+    "Key=RootSnapshotId,Value=${root_snapshot_id}" \
+    "Key=NixSnapshotId,Value=${nix_snapshot_id}"
+
+aws ec2 deregister-image --region "$aws_region" --image-id "$capture_image_id" >/dev/null
 
 jq -n \
   --arg architecture "$architecture" \
+  --arg system "$system" \
   --arg inputHash "$input_hash" \
   --arg builderGitRev "$builder_git_rev" \
   --arg flakeLockRev "$flake_lock_rev" \
-  --arg amiId "$ami_id" \
+  --arg amiId "$image_id" \
   --arg rootSnapshotId "$root_snapshot_id" \
   --arg nixSnapshotId "$nix_snapshot_id" \
-  --arg nixosSystemClosure "$nixos_system_closure" \
+  --arg rootDeviceName "$root_device_name" \
+  --arg nixDeviceName "$nix_device_name" \
+  --arg nixPool "$nix_pool" \
   '{
     architecture: $architecture,
+    system: $system,
     inputHash: $inputHash,
     builderGitRev: $builderGitRev,
     flakeLockRev: $flakeLockRev,
     amiId: $amiId,
     rootSnapshotId: $rootSnapshotId,
     nixSnapshotId: $nixSnapshotId,
-    nixosSystemClosure: $nixosSystemClosure
+    rootDeviceName: $rootDeviceName,
+    nixDeviceName: $nixDeviceName,
+    nixPool: $nixPool
   }' > artifact.json
 
-echo "Published AMI: $ami_id"
+if [ "$destroy_builder" = true ]; then
+  tofu destroy \
+    -var "aws_region=$aws_region" \
+    -var "name=$name" \
+    -var "system=$system" \
+    -var "builder_instance_type=$builder_instance_type" \
+    -var "input_hash=$input_hash" \
+    -var "builder_git_rev=$builder_git_rev" \
+    -var "flake_lock_rev=$flake_lock_rev" \
+    "${extra_tofu_args[@]}"
+fi
+
+echo "Published AMI: $image_id"

user-data.sh.tftpl

@@ -0,0 +1,180 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+nix_device='${nix_device}'
+nix_pool='${nix_pool}'
+state_dir=/var/lib/nixos-zfs-ec2-ami
+bootstrap_marker="$state_dir/bootstrap-generation-installed"
+ready_marker="$state_dir/ami-ready"
+
+mkdir -p "$state_dir"
+
+write_bootstrap_config() {
+  mkdir -p /etc/nixos
+  cat > /etc/nixos/flake.nix <<'EOF'
+{
+  inputs = {
+    nixpkgs.url = "github:NixOS/nixpkgs/nixos-unstable";
+  };
+
+  outputs = { nixpkgs, ... }: {
+    nixosConfigurations.ami-builder-bootstrap = nixpkgs.lib.nixosSystem {
+      system = "${system}";
+      modules = [
+        ({ modulesPath, pkgs, ... }: {
+          imports = [ "$${modulesPath}/virtualisation/amazon-image.nix" ];
+
+          networking.hostName = "nixos-zfs-ec2-builder";
+          networking.hostId = "6e69787a";
+
+          nix.settings.experimental-features = [ "nix-command" "flakes" ];
+          boot.supportedFilesystems = [ "zfs" ];
+          boot.zfs.forceImportRoot = false;
+
+          environment.systemPackages = [
+            pkgs.git
+            pkgs.rsync
+            pkgs.zfs
+          ];
+
+          services.amazon-ssm-agent.enable = true;
+
+          users.users.builder-admin = {
+            isNormalUser = true;
+            extraGroups = [ "wheel" ];
+            openssh.authorizedKeys.keys = ${ssh_keys};
+          };
+          security.sudo.wheelNeedsPassword = false;
+
+          system.stateVersion = "25.05";
+        })
+      ];
+    };
+  };
+}
+EOF
+}
+
+write_final_config() {
+  mkdir -p /etc/nixos
+  cat > /etc/nixos/flake.nix <<'EOF'
+{
+  inputs = {
+    nixpkgs.url = "github:NixOS/nixpkgs/nixos-unstable";
+  };
+
+  outputs = { nixpkgs, ... }: {
+    nixosConfigurations.nixos-zfs-ec2 = nixpkgs.lib.nixosSystem {
+      system = "${system}";
+      modules = [
+        ({ modulesPath, pkgs, ... }: {
+          imports = [ "$${modulesPath}/virtualisation/amazon-image.nix" ];
+
+          networking.hostName = "nixos-zfs-ec2";
+          networking.hostId = "6e69787a";
+
+          nix.settings.experimental-features = [ "nix-command" "flakes" ];
+          boot.supportedFilesystems = [ "zfs" ];
+          boot.zfs.extraPools = [ "${nix_pool}" ];
+          boot.zfs.forceImportRoot = false;
+
+          fileSystems."/nix" = {
+            device = "${nix_pool}/nix";
+            fsType = "zfs";
+            neededForBoot = true;
+          };
+
+          systemd.sockets.nix-daemon = {
+            after = [ "nix.mount" ];
+            requires = [ "nix.mount" ];
+          };
+
+          systemd.services.nix-daemon = {
+            after = [ "nix.mount" ];
+            requires = [ "nix.mount" ];
+          };
+
+          environment.systemPackages = [
+            pkgs.git
+            pkgs.rsync
+            pkgs.zfs
+          ];
+
+          services.amazon-ssm-agent.enable = true;
+
+          users.users.builder-admin = {
+            isNormalUser = true;
+            extraGroups = [ "wheel" ];
+            openssh.authorizedKeys.keys = ${ssh_keys};
+          };
+          security.sudo.wheelNeedsPassword = false;
+
+          system.stateVersion = "25.05";
+        })
+      ];
+    };
+  };
+}
+EOF
+}
+
+if [ ! -e "$bootstrap_marker" ]; then
+  write_bootstrap_config
+  nixos-rebuild boot --flake /etc/nixos#ami-builder-bootstrap
+  touch "$bootstrap_marker"
+  systemctl reboot
+  exit 0
+fi
+
+for _ in $(seq 1 600); do
+  if [ -e "$nix_device" ]; then
+    break
+  fi
+  sleep 1
+done
+
+if [ ! -e "$nix_device" ]; then
+  echo "Nix device did not appear: $nix_device" >&2
+  exit 1
+fi
+
+if ! zpool list -H "$nix_pool" >/dev/null 2>&1; then
+  if ! zpool import "$nix_pool" >/dev/null 2>&1; then
+    zpool create \
+      -f \
+      -o ashift=12 \
+      -O acltype=posixacl \
+      -O atime=off \
+      -O compression=zstd \
+      -O mountpoint=none \
+      -O xattr=sa \
+      "$nix_pool" "$nix_device"
+
+    zfs create -o mountpoint=legacy "$nix_pool/nix"
+  fi
+fi
+
+mkdir -p /mnt/nixos-zfs-nix
+if ! mountpoint -q /mnt/nixos-zfs-nix; then
+  mount -t zfs "$nix_pool/nix" /mnt/nixos-zfs-nix
+fi
+
+write_final_config
+nixos-rebuild boot --flake /etc/nixos#nixos-zfs-ec2
+
+rsync -aHAX --numeric-ids \
+  --exclude='/var/nix/daemon-socket/*' \
+  /nix/ /mnt/nixos-zfs-nix/
+rm -f /mnt/nixos-zfs-nix/var/nix/daemon-socket/socket
+
+umount /mnt/nixos-zfs-nix
+zpool export "$nix_pool"
+
+cat > "$ready_marker" <<EOF
+architecture=${system}
+input_hash=${input_hash}
+builder_git_rev=${builder_git_rev}
+flake_lock_rev=${flake_lock_rev}
+EOF
+
+systemctl poweroff

variables.tf

@@ -7,22 +7,18 @@ variable "aws_region" {
 variable "name" {
   type        = string
   default     = "nixos-zfs-ec2"
-  description = "Name prefix for the published AMI."
+  description = "Name prefix for builder resources and the published AMI."
 }
 
-variable "architecture" {
+variable "system" {
   type        = string
-  description = "AMI architecture, such as arm64 or x86_64."
+  default     = "aarch64-linux"
-}
+  description = "Target NixOS system for the AMI."
 
-variable "root_snapshot_id" {
+  validation {
-  type        = string
+    condition     = contains(["aarch64-linux", "x86_64-linux"], var.system)
-  description = "Snapshot ID for the bootable NixOS root volume."
+    error_message = "system must be aarch64-linux or x86_64-linux."
-}
+  }
-
-variable "nix_snapshot_id" {
-  type        = string
-  description = "Snapshot ID for the preseeded ZFS /nix volume."
 }
 
 variable "input_hash" {
@@ -40,32 +36,68 @@ variable "flake_lock_rev" {
   description = "Revision or hash representing flake.lock inputs."
 }
 
-variable "nixos_system_closure" {
+variable "builder_instance_type" {
   type        = string
-  default     = ""
+  default     = "t4g.large"
-  description = "NixOS system closure path used to build the root image."
+  description = "Temporary EC2 instance type used to build and seed the AMI."
 }
 
-variable "root_device_name" {
+variable "vpc_id" {
   type        = string
-  default     = "/dev/xvda"
+  default     = null
-  description = "Root block device name used for AMI registration."
+  description = "VPC ID. Defaults to the account default VPC."
 }
 
-variable "nix_device_name" {
+variable "subnet_id" {
   type        = string
-  default     = "/dev/sdf"
+  default     = null
-  description = "AMI block device name for the /nix EBS snapshot. Runtime imports by ZFS pool identity, not this device name."
+  description = "Subnet ID. Defaults to the first subnet in the selected/default VPC."
+}
+
+variable "base_nixos_ami_id" {
+  type        = string
+  default     = null
+  description = "Base NixOS AMI used only for the temporary builder. If null, OpenTofu searches for a recent NixOS AMI."
+}
+
+variable "base_nixos_ami_owners" {
+  type        = list(string)
+  default     = ["427812963091"]
+  description = "AWS account IDs to search for base NixOS AMIs."
+}
+
+variable "base_nixos_ami_name_pattern" {
+  type        = string
+  default     = null
+  description = "Base NixOS AMI name filter."
 }
 
 variable "root_volume_size_gb" {
   type        = number
   default     = 30
-  description = "Root EBS volume size."
+  description = "Temporary builder root EBS volume size. This becomes the final AMI root snapshot size."
 }
 
 variable "nix_volume_size_gb" {
   type        = number
   default     = 50
-  description = "/nix EBS volume size."
+  description = "Temporary /nix EBS volume size. This becomes the final AMI /nix snapshot size."
+}
+
+variable "nix_pool" {
+  type        = string
+  default     = "nixos-zfs-nix"
+  description = "ZFS pool name for the AMI-provided /nix volume."
+}
+
+variable "nix_device_name" {
+  type        = string
+  default     = "/dev/sdf"
+  description = "EC2 block device name for the AMI /nix volume. NixOS imports by ZFS pool identity, not this device name."
+}
+
+variable "ssh_public_key" {
+  type        = string
+  default     = null
+  description = "Optional SSH public key for debugging the temporary builder."
 }