feat: AMI generation with UEFI, systemd repart and NitroTPM

flake.nix

@@ -70,13 +70,11 @@
       # NOTE: Github Actions doesn't support kvm on arm64 builds
       nixosConfigurations.x86_64-linux = nixpkgs.lib.nixosSystem {
         modules = [
+          ./modules/amazon-ami.nix
           (
             { modulesPath, ... }:
             {
-              imports = [ "${modulesPath}/virtualisation/amazon-image.nix" ];
-              image.modules.amazon = {
-                amazonImage.format = "raw";
-              };
+              amazon.ami.enable = true;
               nixpkgs.hostPlatform = "x86_64-linux";
               system.stateVersion = "26.05";
             }
@@ -92,7 +90,7 @@
           formatting = treefmtEval.${system}.config.build.check self;
         })
         // {
-          x86_64-linux.system = self.nixosConfigurations.x86_64-linux.config.system.build.images.amazon;
+          x86_64-linux.system = self.nixosConfigurations.x86_64-linux.config.system.build.amazon-image;
         };
 
       devShells = genAttrs supportedSystems (system: {

modules/amazon-ami.nix

@@ -0,0 +1,128 @@
+{
+  config,
+  lib,
+  pkgs,
+  modulesPath,
+  ...
+}:
+let
+  cfg = config.amazon.ami;
+in
+{
+  imports = [
+    "${modulesPath}/image/repart.nix"
+    "${modulesPath}/virtualisation/amazon-image.nix"
+  ];
+
+  options.amazon.ami = {
+    enable = lib.mkEnableOption "Amazon AMI generation using systemd-repart";
+
+    name = lib.mkOption {
+      type = lib.types.str;
+      default = "nixos-${config.system.nixos.label}-${pkgs.stdenv.hostPlatform.system}";
+      description = "The name of the AMI.";
+    };
+
+    description = lib.mkOption {
+      type = lib.types.str;
+      default = "NixOS ${config.system.nixos.label} ${pkgs.stdenv.hostPlatform.system}";
+      description = "The description of the AMI.";
+    };
+
+    tpmSupport = lib.mkOption {
+      type = lib.types.bool;
+      default = false;
+      description = "Enable TPM 2.0 support for UEFI x86_64 images.";
+    };
+
+    enaSupport = lib.mkOption {
+      type = lib.types.bool;
+      default = true;
+      description = "Enable ENA support.";
+    };
+
+    imdsSupport = lib.mkOption {
+      type = lib.types.enum [ "v2.0" "v1.0" ];
+      default = "v2.0";
+      description = "IMDS version support.";
+    };
+
+    public = lib.mkOption {
+      type = lib.types.bool;
+      default = false;
+      description = "Whether to make the AMI public.";
+    };
+
+    attestation = {
+      enable = lib.mkEnableOption "NitroTPM Attestation support";
+    };
+  };
+
+  config = lib.mkIf cfg.enable {
+    # drop BIOS support and only build UEFI images
+    amazonImage.format = "raw";
+
+    boot.loader.grub.enable = false;
+    boot.loader.systemd-boot.enable = true;
+    boot.loader.efi.canTouchEfiVariables = true;
+
+    # systemd-repart for image generation
+    image.repart = {
+      enable = true;
+      name = cfg.name;
+      partitions = {
+        "esp" = {
+          repartConfig = {
+            Type = "esp";
+            Format = "vfat";
+            SizeMinBytes = "128M";
+          };
+        };
+        "root" = {
+          repartConfig = {
+            Type = "root";
+            Format = "ext4";
+            Minimize = "best";
+          };
+        };
+      };
+    };
+
+    # EC2 metadata and instance attestation bits
+    boot.initrd.systemd.enable = true;
+
+    environment.systemPackages = lib.optionals cfg.attestation.enable [
+      pkgs.nitrotpm-tools
+      pkgs.tpm2-tools
+    ];
+
+    system.build.ami-info = pkgs.writeText "ami-info.json" (builtins.toJSON {
+      name = cfg.name;
+      description = cfg.description;
+      tpmSupport = cfg.tpmSupport;
+      enaSupport = cfg.enaSupport;
+      imdsSupport = cfg.imdsSupport;
+      public = cfg.public;
+      architecture = if pkgs.stdenv.hostPlatform.isAarch64 then "arm64" else "x86_64";
+      boot_mode = "uefi";
+    });
+
+    system.build.amazon-image = pkgs.stdenv.mkDerivation {
+      name = "amazon-image-${cfg.name}";
+
+      nativeBuildInputs = [ pkgs.nitrotpm-tools ];
+
+      buildCommand = ''
+        mkdir -p $out
+        IMAGE=$(ls ${config.system.build.image}/*.raw)
+        ln -s $IMAGE $out/image.raw
+        ln -s ${config.system.build.ami-info} $out/ami-info.json
+
+        ${lib.optionalString cfg.attestation.enable ''
+          echo "computing NitroTPM PCRs..."
+          nitro-tpm-pcr-compute --image $IMAGE > $out/tpm_pcr.json
+        ''}
+      '';
+    };
+  };
+}

upload-ami/src/upload_ami/upload_ami.py

@@ -25,7 +25,13 @@ class ImageInfo(TypedDict):
     label: str
     system: str
     boot_mode: BootModeValuesType
-    format: str
+    format: str | None
+    name: str | None
+    description: str | None
+    tpm_support: bool | None
+    ena_support: bool | None
+    imds_support: Literal["v1.0", "v2.0"] | None
+    public: bool | None
 
 
 def upload_to_s3_if_not_exists(
@@ -202,10 +208,11 @@ def register_image_if_not_exists(
                     },
                 }
             ],
+            "Description": image_info.get("description") or f"NixOS {image_name}",
             "RootDeviceName": "/dev/xvda",
             "VirtualizationType": "hvm",
-            "EnaSupport": True,
-            "ImdsSupport": "v2.0",
+            "EnaSupport": image_info.get("ena_support", True),
+            "ImdsSupport": image_info.get("imds_support", "v2.0"),
             "SriovNetSupport": "simple",
             "TagSpecifications": [
                 {
@@ -219,7 +226,7 @@ def register_image_if_not_exists(
         }
 
         if (
-            enable_tpm
+            (enable_tpm or image_info.get("tpm_support"))
             and architecture == "x86_64"
             and image_info["boot_mode"] == "uefi"
         ):
@@ -375,7 +382,9 @@ def upload_ami(
     image_file = Path(image_info["file"])
     label = image_info["label"]
     system = image_info["system"]
-    image_name = prefix + label + "-" + system + ("." + run_id if run_id else "")
+    image_name = image_info.get("name") or (
+        prefix + label + "-" + system + ("." + run_id if run_id else "")
+    )
 
     image_format = image_info.get("format") or "VHD"
     if ebs_direct:
@@ -390,8 +399,14 @@ def upload_ami(
             s3, ec2, s3_bucket, image_name, image_file, image_format, import_role_name
         )
 
+    is_public = public or image_info.get("public", False)
     image_id = register_image_if_not_exists(
-        ec2, image_name, image_info, snapshot_id, public, enable_tpm
+        ec2,
+        image_name,
+        image_info,
+        snapshot_id,
+        is_public,
+        enable_tpm,
     )
 
     image_ids: dict[str, str] = {}
@@ -409,7 +424,7 @@ def upload_ami(
                 image_name,
                 ec2.meta.region_name,
                 regions,
-                public,
+                is_public,
                 best_effort_regions,
             )
         )