fix(os): harden AWS guest image release validation

Author: kvinwang

docs/aws-ec2-production-verifier-runbook.md

@@ -74,7 +74,8 @@ After the release artifact hashes and AWS PCR references match the published
 release evidence, deploy with **`dstack-cloud`** (`platform: aws`). The CLI
 imports the local UKI `disk.raw` as an Attestable AMI (UEFI + NitroTPM v2.0) when
 `aws_config.ami_id` is empty, builds the shared config disk with
-`aws_measurement` from `measurement.aws.cbor`, and launches the instance.
+`aws_measurement` from `measurement.aws.cbor`, imports a minimal GPT data-disk
+template whose partition is labeled `dstack-data`, and launches the instance.
 
 ```bash
 # one-time: install CLI
@@ -88,21 +89,23 @@ cd my-aws-app
 
 dstack-cloud pull dstack-0.6.0   # UKI package must include assemble-time measurement.aws.cbor
 dstack-cloud prepare             # embeds fixed os_image_hash + aws_measurement (no PCR recompute)
-dstack-cloud deploy              # import AMI if needed, import shared snapshot, run-instances
+dstack-cloud deploy              # import AMI/shared/labeled-data snapshots, run-instances
 
 dstack-cloud status
 dstack-cloud logs --follow
 ```
 
-Record the resulting AMI id, shared snapshot, and instance id from
+Record the resulting AMI id, shared/data snapshots, and instance id from
 `state.json` / `dstack-cloud status` in the release review package.
 
 Recompute AWS reference PCRs and the UKI AuthentiCode hash from the release
 UKI with version-pinned host tools (do not rely on an unpinned container
 toolchain for reference measurements):
 
 ```bash
-cargo install --git https://github.com/aws/NitroTPM-Tools --locked nitro-tpm-pcr-compute
+cargo install --git https://github.com/aws/NitroTPM-Tools \
+  --rev d76d6eeebd4169b00a3c3af9858852d48f40e748 \
+  --locked nitro-tpm-pcr-compute   # aws/NitroTPM-Tools v1.1.2
 # The UKI is EFI/BOOT/BOOTX64.EFI in the 256 MiB EFI partition at 1 MiB.
 dd if=published/disk.raw of=efi.img bs=1M skip=1 count=256 status=none
 mcopy -i efi.img ::EFI/BOOT/BOOTX64.EFI dstack-uki.efi

dstack/scripts/bin/dstack-cloud

@@ -310,7 +310,7 @@ class AwsConfig:
     boot_image_tar: str = ""  # explicit path override
 
     # Data disk
-    data_size: int = 20  # GiB for new empty data volume
+    data_size: int = 20  # GiB for the volume created from the labeled template
     data_snapshot: str = ""  # optional existing snapshot for data disk
 
     # Shared-disk device names (dstack host-shared + data)
@@ -324,7 +324,7 @@ class AwsConfig:
     security_group_ids: List[str] = field(default_factory=list)
     iam_instance_profile: str = ""  # name or ARN
 
-    # S3 used when importing AMI / shared snapshot from local disk.raw
+    # S3 used when importing AMI / shared / labeled-data snapshots from raw disks
     s3_bucket: str = ""
     s3_prefix: str = "dstack-import"
     import_role_name: str = ""  # optional vmimport role override
@@ -384,6 +384,7 @@ class DeploymentState:
     instance_id: str = ""
     ami_id: str = ""
     shared_snapshot: str = ""
+    data_snapshot: str = ""  # Data template snapshot created by dstack-cloud
     external_ip: str = ""
     internal_ip: str = ""
     status: str = ""  # RUNNING, STOPPED, TERMINATED, etc.
@@ -731,25 +732,7 @@ class CloudDeploymentManager:
 
         # Create a minimal raw disk image with GPT partition table
         with tempfile.TemporaryDirectory() as tmpdir:
-            raw_file = os.path.join(tmpdir, "disk.raw")
-
-            # Create a 10MB sparse file (enough for GPT)
-            disk_size_bytes = 10 * 1024 * 1024
-            with open(raw_file, 'wb') as f:
-                f.truncate(disk_size_bytes)
-
-            # Create GPT partition table with dstack-data partition using sgdisk
-            # -o: clear and create new GPT
-            # -n 1:0:0: create partition 1, start at first available, end at last available
-            # -c 1:dstack-data: set partition 1 name (PARTLABEL) to dstack-data
-            result = subprocess.run(
-                ["sgdisk", "-o", "-n", "1:0:0", "-c", "1:dstack-data", raw_file],
-                capture_output=True, text=True
-            )
-            if result.returncode != 0:
-                raise RuntimeError(f"Failed to create GPT partition table: {result.stderr}")
-
-            logger.debug("Created GPT partition table with dstack-data label")
+            self._build_data_raw_disk(Path(tmpdir))
 
             # Compress to tar.gz for upload
             tar_file = os.path.join(tmpdir, "disk.tar.gz")
@@ -1697,6 +1680,29 @@ class CloudDeploymentManager:
             logger.warning(".user-config not found, skipping")
         return raw_file
 
+    def _build_data_raw_disk(self, work_path: Optional[Path] = None) -> Path:
+        """Build a minimal GPT disk with a PARTLABEL=dstack-data partition."""
+        import shutil
+
+        if not shutil.which("sgdisk"):
+            raise FileNotFoundError(
+                "'sgdisk' command not found. Install the gdisk package."
+            )
+        if work_path is None:
+            work_path = Path(tempfile.mkdtemp(prefix="dstack-data-"))
+        work_path.mkdir(parents=True, exist_ok=True)
+        raw_file = work_path / "disk.raw"
+        subprocess.run(["truncate", "-s", "10M", str(raw_file)], check=True)
+        subprocess.run(
+            [
+                "sgdisk", "-o", "-n", "1:0:0", "-c", "1:dstack-data",
+                str(raw_file),
+            ],
+            check=True,
+            capture_output=True,
+        )
+        return raw_file
+
     def _aws_import_snapshot_from_raw(
         self,
         config: AwsConfig,
@@ -1931,14 +1937,31 @@ class CloudDeploymentManager:
                 pass
 
         # Block device mappings: shared + data
+        managed_data_snapshot = ""
+        if config.data_snapshot:
+            data_snapshot = config.data_snapshot
+        else:
+            logger.info("building labeled data disk template...")
+            data_raw = self._build_data_raw_disk()
+            try:
+                data_token = f"data-{config.instance_name}-{int(time.time())}"[:64]
+                data_snapshot = self._aws_import_snapshot_from_raw(
+                    config, data_raw, f"dstack-data-{config.instance_name}", data_token
+                )
+                managed_data_snapshot = data_snapshot
+            finally:
+                try:
+                    import shutil
+                    shutil.rmtree(data_raw.parent, ignore_errors=True)
+                except Exception:
+                    pass
+
         data_ebs: Dict[str, Any] = {
             "DeleteOnTermination": True,
             "VolumeType": config.volume_type,
+            "SnapshotId": data_snapshot,
+            "VolumeSize": int(config.data_size),
         }
-        if config.data_snapshot:
-            data_ebs["SnapshotId"] = config.data_snapshot
-        else:
-            data_ebs["VolumeSize"] = int(config.data_size)
 
         mappings = [
             {
@@ -2019,6 +2042,7 @@ class CloudDeploymentManager:
             instance_id=instance_id,
             ami_id=ami_id,
             shared_snapshot=shared_snapshot,
+            data_snapshot=managed_data_snapshot,
             external_ip=external_ip,
             internal_ip=internal_ip,
             status=inst.get("State", {}).get("Name", "running").upper(),
@@ -2032,6 +2056,8 @@ class CloudDeploymentManager:
         logger.info(f"Instance ID: {instance_id}")
         logger.info(f"AMI: {ami_id}")
         logger.info(f"Shared snapshot: {shared_snapshot}")
+        if managed_data_snapshot:
+            logger.info(f"Data template snapshot: {managed_data_snapshot}")
         logger.info(f"Public IP: {external_ip or '(none)'}")
         logger.info(f"Private IP: {internal_ip or '(none)'}")
         logger.info("")
@@ -2758,6 +2784,13 @@ class CloudDeploymentManager:
                     "--region", state.region,
                     "--snapshot-id", state.shared_snapshot,
                 ], capture=True, check=False)
+            if not keep_images and state.data_snapshot:
+                logger.info(f"Deleting data template snapshot {state.data_snapshot}...")
+                self._run_aws([
+                    "ec2", "delete-snapshot",
+                    "--region", state.region,
+                    "--snapshot-id", state.data_snapshot,
+                ], capture=True, check=False)
             state_path = self.work_dir / STATE_FILE
             if state_path.exists():
                 state_path.unlink()

dstack/tpm2/src/device.rs

@@ -320,7 +320,10 @@ mod tests {
         let bytes = cmd.finalize();
 
         assert_eq!(&bytes[0..2], &[0x80, 0x02]);
-        assert_eq!(u32::from_be_bytes(bytes[2..6].try_into().unwrap()), 25);
+        assert_eq!(
+            u32::from_be_bytes([bytes[2], bytes[3], bytes[4], bytes[5]]),
+            25
+        );
         assert_eq!(&bytes[6..10], &[0x20, 0x00, 0x00, 0x01]);
         assert_eq!(&bytes[10..14], &[0x01, 0x80, 0x10, 0x00]);
         assert_eq!(&bytes[14..18], &[0x01, 0x80, 0x10, 0x00]);
@@ -358,7 +361,10 @@ mod tests {
             0x00, 0x00, 0x00, 0x00, // TPM_RC_SUCCESS
         ];
 
-        let parsed = TpmResponse::parse(&response).unwrap();
+        let parsed = match TpmResponse::parse(&response) {
+            Ok(parsed) => parsed,
+            Err(error) => panic!("failed to parse minimal TPM response: {error}"),
+        };
         assert!(parsed.is_success());
         assert!(parsed.data.is_empty());
     }

os/common/rootfs/dstack-prepare.sh

@@ -149,31 +149,6 @@ has_luks_header() {
 	return 1
 }
 
-create_data_partition() {
-	local disk="$1"
-	log "Creating GPT partition table on ${disk}..."
-	if ! command -v sgdisk >/dev/null 2>&1; then
-		log "Error: sgdisk not available, cannot create partition table"
-		return 1
-	fi
-	# Create GPT with single partition filling entire disk
-	sgdisk -Z "$disk" >/dev/null || true # Zap any existing data
-	sgdisk -n 1:1MiB:0 -c 1:dstack-data -t 1:8300 "$disk" >/dev/null || return 1
-	# Trigger kernel to re-read partition table
-	blockdev --rereadpt "$disk" >/dev/null || true
-	udevadm settle >/dev/null || sleep 1
-	part_device=$(
-		lsblk -nr -o PATH "$disk" 2>/dev/null | sed -n '2p'
-	)
-	if [ -n "$part_device" ] && [ -b "$part_device" ]; then
-		log "Created partition: $part_device"
-		echo "$part_device"
-		return 0
-	fi
-	log "Failed to create partition"
-	return 1
-}
-
 choose_data_device() {
 	local override="$1"
 	local dev=""
@@ -216,16 +191,11 @@ choose_data_device() {
 		return 1
 	fi
 
-	# 3.3. /dev/vdb is empty, create partition table
-	log "Empty disk detected at /dev/vdb, creating dstack-data partition..."
-	local new_partition
-	new_partition=$(create_data_partition /dev/vdb)
-	if [ -z "$new_partition" ]; then
-		log "Error: Failed to create partition on /dev/vdb"
-		return 1
-	fi
-	echo "$new_partition"
-	return 0
+	# Do not guess which blank disk is persistent storage or destructively
+	# initialize it. Cloud deployment tooling must provision PARTLABEL=dstack-data.
+	log "Error: Empty /dev/vdb has no 'dstack-data' partition"
+	log "Provision a labeled data-disk image or specify dstack.data_device"
+	return 1
 }
 
 DATA_DEVICE_OVERRIDE=$(get_cmdline_value "dstack.data_device" || true)

os/image/README.md

@@ -23,7 +23,8 @@ Deploy tooling (`dstack-cloud prepare`) only **embeds** these files into
 AWS PCR precompute requires a pinned host `nitro-tpm-pcr-compute` binary (Rust,
 [aws/NitroTPM-Tools](https://github.com/aws/NitroTPM-Tools)). Set
 `NITRO_TPM_PCR_COMPUTE_BIN` or install it on `PATH`, for example with
-`cargo install --git https://github.com/aws/NitroTPM-Tools --locked nitro-tpm-pcr-compute`.
+`cargo install --git https://github.com/aws/NitroTPM-Tools --rev d76d6eeebd4169b00a3c3af9858852d48f40e748 --locked nitro-tpm-pcr-compute`
+(aws/NitroTPM-Tools v1.1.2).
 If it is missing, UKI assembly fails.
 
 `mk-image-mr.sh <release.tar.gz>` creates the flattened, rootfs-free

os/image/assemble.sh

@@ -459,7 +459,7 @@ if [[ "$UKI_CREATED" = "1" ]]; then
     if [[ -z "$pcr_compute_bin" ]]; then
         echo "Error: cannot produce measurement.aws.cbor for UKI image." >&2
         echo "Install the pinned host tool:" >&2
-        echo "  cargo install --git https://github.com/aws/NitroTPM-Tools --locked nitro-tpm-pcr-compute" >&2
+        echo "  cargo install --git https://github.com/aws/NitroTPM-Tools --rev d76d6eeebd4169b00a3c3af9858852d48f40e748 --locked nitro-tpm-pcr-compute" >&2
         echo "  # or set NITRO_TPM_PCR_COMPUTE_BIN=/path/to/nitro-tpm-pcr-compute" >&2
         echo "measurement.aws.cbor must be fixed at assemble time for a stable os_image_hash." >&2
         exit 1

os/yocto/layers/meta-dstack/recipes-kernel/linux/files/dstack.cfg

@@ -47,3 +47,15 @@ CONFIG_WLAN=n
 # "verity: Cannot initialize hash function (-ENOENT)".
 CONFIG_CRYPTO_SHA256=y
 CONFIG_CRYPTO_SHA512=y
+
+# Production kernel self-protection baseline.
+CONFIG_SECURITY_DMESG_RESTRICT=y
+CONFIG_IO_STRICT_DEVMEM=y
+CONFIG_INIT_ON_ALLOC_DEFAULT_ON=y
+CONFIG_INIT_ON_FREE_DEFAULT_ON=y
+CONFIG_SLAB_FREELIST_RANDOM=y
+CONFIG_SLAB_FREELIST_HARDENED=y
+CONFIG_HARDENED_USERCOPY=y
+CONFIG_FORTIFY_SOURCE=y
+CONFIG_MAGIC_SYSRQ=n
+CONFIG_LEGACY_TIOCSTI=n

os/yocto/repro-build/Dockerfile.repro

@@ -11,6 +11,10 @@ ENV RUSTUP_HOME=/usr/local/rustup \
     CARGO_HOME=/usr/local/cargo \
     PATH=/usr/local/cargo/bin:$PATH
 
+# aws/NitroTPM-Tools v1.1.2. This binary contributes release measurements, so
+# pin the source revision rather than following the repository default branch.
+ARG NITRO_TPM_TOOLS_REV=d76d6eeebd4169b00a3c3af9858852d48f40e748
+
 # Set timezone
 RUN ln -snf /usr/share/zoneinfo/$TZ /etc/localtime && echo $TZ > /etc/timezone
 
@@ -26,6 +30,7 @@ RUN apt update && apt install -y \
     ca-certificates \
     curl \
     git \
+    jq \
     file \
     gawk \
     wget \
@@ -42,6 +47,8 @@ RUN curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | \
     sh -s -- -y --profile minimal --default-toolchain 1.92 && \
     rustup component add rustfmt clippy rust-analyzer && \
     rustup target add wasm32-unknown-unknown x86_64-unknown-linux-musl thumbv6m-none-eabi && \
+    cargo install --git https://github.com/aws/NitroTPM-Tools \
+      --rev "$NITRO_TPM_TOOLS_REV" --locked nitro-tpm-pcr-compute && \
     chmod -R a+rX "$RUSTUP_HOME" "$CARGO_HOME" && \
     cargo --version && \
     rustc --version

os/yocto/repro-build/check.sh

@@ -31,12 +31,14 @@ COMPARE_IMAGE_WHITELIST=(
     "initramfs.cpio.gz"
     "metadata.json"
     "measurement.gcp.cbor"
+    "measurement.aws.cbor"
     "measurement.snp.cbor"
     "measurement.tdx.cbor"
     "ovmf.fd"
     "rootfs.img.parted.verity"
     "sha256sum.txt"
     "auth_hash.txt"
+    "aws-pcrs.json"
     "gcp/efi-root/EFI/BOOT/BOOTX64.EFI"
 )