fix(simulator): replay measured NitroTPM boot PCRs

dstack/dstack-types/src/lib.rs

@@ -1154,6 +1154,30 @@ pub struct TeeSimulatorConfig {
     /// JSON serialized VmConfig used to generate mock platform evidence.
     #[serde(default, skip_serializing_if = "Option::is_none")]
     pub vm_config: Option<String>,
+    /// Ordered SHA-384 PCR extensions used to reproduce the AWS boot state in
+    /// the development NitroTPM simulator.
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub aws_pcr_replay: Option<AwsPcrReplay>,
+}
+
+#[derive(Deserialize, Serialize, Debug, Clone, PartialEq, Eq)]
+pub struct AwsPcrReplay {
+    pub version: u32,
+    pub events: Vec<AwsPcrReplayEvent>,
+    #[serde(with = "hex_bytes")]
+    pub pcr4: Vec<u8>,
+    #[serde(with = "hex_bytes")]
+    pub pcr7: Vec<u8>,
+    #[serde(with = "hex_bytes")]
+    pub pcr12: Vec<u8>,
+}
+
+#[derive(Deserialize, Serialize, Debug, Clone, PartialEq, Eq)]
+pub struct AwsPcrReplayEvent {
+    pub pcr: u16,
+    pub event_type: String,
+    #[serde(with = "hex_bytes")]
+    pub digest: Vec<u8>,
 }
 
 #[derive(Deserialize, Serialize, Debug, Clone, Copy, Default, PartialEq, Eq, Encode, Decode)]

dstack/tee-simulator/src/tpm.rs

@@ -19,7 +19,7 @@ use std::{
 
 use anyhow::{bail, Context, Result};
 use aws_nitro_enclaves_nsm_api::api::{Request as NsmRequest, Response as NsmResponse};
-use dstack_types::TeeSimulatorConfig;
+use dstack_types::{AwsPcrReplay, TeeSimulatorConfig};
 use mock_attestation::{nsm::NsmGenerator, parse_seed, server::MockCollateralState};
 use tpm2::{add_command_capability, TpmAlgId, TpmCc, TpmContext};
 
@@ -340,6 +340,12 @@ pub fn run_nitro_vtpm(runtime_dir: &Path, config: &TeeSimulatorConfig) -> Result
     drop(swtpm_stream);
     fs_err::write(runtime_dir.join("swtpm.pid"), child.id().to_string())?;
 
+    let replay = config
+        .aws_pcr_replay
+        .as_ref()
+        .context("tee_simulator.aws_pcr_replay is required for NitroTPM")?;
+    replay_aws_boot_pcrs(&mut simulator, replay)?;
+
     let (control, mut proxy, tpm_num) = create_vtpm_proxy()?;
     let proxy_thread = thread::spawn(move || {
         let _control = control;
@@ -370,6 +376,45 @@ pub fn run_nitro_vtpm(runtime_dir: &Path, config: &TeeSimulatorConfig) -> Result
     result
 }
 
+fn replay_aws_boot_pcrs(backend: &mut UnixStream, replay: &AwsPcrReplay) -> Result<()> {
+    anyhow::ensure!(replay.version == 1, "unsupported AWS PCR replay version");
+    let mut tpm = TpmContext::from_stream(
+        backend
+            .try_clone()
+            .context("failed to clone NitroTPM stream for PCR replay")?,
+        "NitroTPM replay backend",
+    );
+    for event in &replay.events {
+        anyhow::ensure!(
+            matches!(event.pcr, 4 | 7 | 12),
+            "AWS PCR replay contains unsupported PCR {}",
+            event.pcr
+        );
+        anyhow::ensure!(
+            event.digest.len() == 48,
+            "AWS PCR replay event digest must be SHA-384"
+        );
+        tpm.pcr_extend(event.pcr.into(), &event.digest, TpmAlgId::Sha384)
+            .with_context(|| {
+                format!(
+                    "failed to replay {} into PCR{}",
+                    event.event_type, event.pcr
+                )
+            })?;
+    }
+    for (index, expected) in [(4u16, &replay.pcr4), (7, &replay.pcr7), (12, &replay.pcr12)] {
+        anyhow::ensure!(expected.len() == 48, "expected PCR{index} must be SHA-384");
+        let actual = tpm.pcr_read_single(index.into(), TpmAlgId::Sha384)?;
+        anyhow::ensure!(
+            actual == *expected,
+            "replayed PCR{index} mismatch: expected={}, actual={}",
+            hex::encode(expected),
+            hex::encode(actual)
+        );
+    }
+    Ok(())
+}
+
 fn create_vtpm_proxy() -> Result<(std::fs::File, std::fs::File, u32)> {
     let control = std::fs::OpenOptions::new()
         .read(true)

dstack/vmm/src/app.rs

@@ -1278,6 +1278,13 @@ pub(crate) fn sync_tee_simulator_config(
     let sys_config: dstack_types::SysConfig = serde_json::from_str(sys_config)?;
     let mut simulator_config = simulator_config.clone();
     simulator_config.mr_config = sys_config.mr_config;
+    let vm_config_value: serde_json::Value = serde_json::from_str(&sys_config.vm_config)?;
+    simulator_config.aws_pcr_replay = vm_config_value
+        .get("aws_pcr_replay")
+        .cloned()
+        .map(serde_json::from_value)
+        .transpose()
+        .context("invalid aws_pcr_replay in vm_config")?;
     simulator_config.vm_config = Some(sys_config.vm_config);
     fs::write(path, serde_json::to_vec(&simulator_config)?)
         .context("failed to write TEE simulator config")
@@ -1492,6 +1499,29 @@ fn make_vm_config(
     })?;
     // For backward compatibility
     config["spec_version"] = serde_json::Value::from(1);
+    if is_aws_nitro_tpm {
+        let replay = image
+            .aws_pcr_replay
+            .as_ref()
+            .context("AWS NitroTPM simulation requires measurement.aws.replay.json")?;
+        let replay_measurement = dstack_types::AwsOsImageMeasurement::from_boot_pcrs(
+            &replay.pcr4,
+            &replay.pcr7,
+            &replay.pcr12,
+        )
+        .map_err(anyhow::Error::msg)?;
+        let image_measurement = image
+            .aws_measurement
+            .as_ref()
+            .context("AWS NitroTPM image is missing measurement.aws.cbor")?
+            .decode_measurement()
+            .map_err(anyhow::Error::msg)?;
+        anyhow::ensure!(
+            replay_measurement == image_measurement,
+            "measurement.aws.replay.json does not match measurement.aws.cbor"
+        );
+        config["aws_pcr_replay"] = serde_json::to_value(replay)?;
+    }
     if is_amd_sev_snp {
         if let Some(mr_config) = mr_config {
             MrConfigV3::from_document(&mr_config).context("Invalid mr_config document")?;
@@ -1627,7 +1657,10 @@ mod tests {
             ..Default::default()
         };
         let mr_config = r#"{"version":3}"#;
-        let vm_config = r#"{"image":"dev"}"#;
+        let vm_config = format!(
+            r#"{{"image":"dev","aws_pcr_replay":{{"version":1,"events":[],"pcr4":"{zero}","pcr7":"{zero}","pcr12":"{zero}"}}}}"#,
+            zero = "00".repeat(48)
+        );
         let sys_config = serde_json::json!({
             "kms_urls": [],
             "gateway_urls": [],
@@ -1644,7 +1677,11 @@ mod tests {
         assert_eq!(written.mock_attestation_seed, config.mock_attestation_seed);
         assert_eq!(written.collateral_base_url, config.collateral_base_url);
         assert_eq!(written.mr_config.as_deref(), Some(mr_config));
-        assert_eq!(written.vm_config.as_deref(), Some(vm_config));
+        assert_eq!(written.vm_config.as_deref(), Some(vm_config.as_str()));
+        assert_eq!(
+            written.aws_pcr_replay.as_ref().map(|replay| replay.version),
+            Some(1)
+        );
 
         sync_tee_simulator_config(dir.path(), None, &sys_config)?;
         assert!(!dir.path().join(TEE_SIMULATOR_CONFIG).exists());
@@ -1937,6 +1974,7 @@ mod tests {
             sev_measurement: None,
             gcp_measurement: None,
             aws_measurement: None,
+            aws_pcr_replay: None,
         }
     }
 

dstack/vmm/src/app/image.rs

@@ -8,13 +8,14 @@ use std::path::{Path, PathBuf};
 
 use anyhow::{bail, Context, Result};
 use dstack_types::{
-    AwsOsImageMeasurementDocument, GcpOsImageMeasurementDocument, SevOsImageMeasurementDocument,
-    TdxOsImageMeasurementDocument, GCP_MEASUREMENT_FILENAME, SNP_MEASUREMENT_FILENAME,
-    TDX_MEASUREMENT_FILENAME,
+    AwsOsImageMeasurementDocument, AwsPcrReplay, GcpOsImageMeasurementDocument,
+    SevOsImageMeasurementDocument, TdxOsImageMeasurementDocument, GCP_MEASUREMENT_FILENAME,
+    SNP_MEASUREMENT_FILENAME, TDX_MEASUREMENT_FILENAME,
 };
 use serde::{Deserialize, Serialize};
 
 const AWS_MEASUREMENT_FILENAME: &str = "measurement.aws.cbor";
+const AWS_PCR_REPLAY_FILENAME: &str = "measurement.aws.replay.json";
 
 #[derive(Debug, Serialize, Deserialize)]
 pub struct ImageInfo {
@@ -86,6 +87,8 @@ pub struct Image {
     pub gcp_measurement: Option<GcpOsImageMeasurementDocument>,
     /// AWS NitroTPM no-image-download measurement material.
     pub aws_measurement: Option<AwsOsImageMeasurementDocument>,
+    /// AWS boot events consumed only by the development NitroTPM simulator.
+    pub aws_pcr_replay: Option<AwsPcrReplay>,
 }
 
 impl Image {
@@ -171,6 +174,17 @@ impl Image {
             AWS_MEASUREMENT_FILENAME,
             AwsOsImageMeasurementDocument::new,
         )?;
+        let aws_pcr_replay_path = base_path.join(AWS_PCR_REPLAY_FILENAME);
+        let aws_pcr_replay = if aws_pcr_replay_path.exists() {
+            Some(
+                serde_json::from_slice(&fs::read(&aws_pcr_replay_path).with_context(|| {
+                    format!("failed to read {}", aws_pcr_replay_path.display())
+                })?)
+                .with_context(|| format!("failed to parse {}", aws_pcr_replay_path.display()))?,
+            )
+        } else {
+            None
+        };
         if info.version.is_empty() {
             // Older images does not have version field. Fallback to the version of the image folder name
             info.version = guess_version(&base_path).unwrap_or_default();
@@ -188,6 +202,7 @@ impl Image {
             sev_measurement,
             gcp_measurement,
             aws_measurement,
+            aws_pcr_replay,
         }
         .ensure_exists()
     }

dstack/vmm/src/app/qemu.rs

@@ -1090,6 +1090,7 @@ mod tests {
                 sev_measurement: None,
                 gcp_measurement: None,
                 aws_measurement: None,
+                aws_pcr_replay: None,
             },
             cid: 100,
             workdir: PathBuf::from("/does-not-exist/vm-1"),

os/image/README.md

@@ -14,8 +14,10 @@ When a UKI image is built (`ENABLE_UKI_IMAGE=1`), assemble **always** produces:
 - `measurement.gcp.cbor` — UKI Authenticode binding for GCP TPM
 - `measurement.aws.cbor` — NitroTPM boot binding for AWS EC2
   (`boot_pcr_digest = sha256(PCR4||PCR7||PCR12)`)
+- `measurement.aws.replay.json` — ordered SHA-384 boot-event digests used by
+  the development NitroTPM simulator to reproduce PCR4/7/12
 
-Both are listed in `sha256sum.txt`, so
+All measurement artifacts are listed in `sha256sum.txt`, so
 `digest.txt = sha256(sha256sum.txt) = os_image_hash` is fixed at build time.
 Deploy tooling (`dstack-cloud prepare`) only **embeds** these files into
 `VmConfig`; it must not recompute PCRs (that would change the image identity).
@@ -25,7 +27,10 @@ AWS PCR precompute requires a pinned host `nitro-tpm-pcr-compute` binary (Rust,
 `NITRO_TPM_PCR_COMPUTE_BIN` or install it on `PATH`, for example with
 `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.
+The assembler captures that pinned tool's per-event debug trace, converts it
+to the replay document, and verifies that replaying the events produces the
+tool's reported PCR values. If the tool is missing or the replay does not
+match, UKI assembly fails.
 
 `mk-image-mr.sh <release.tar.gz>` creates the flattened, rootfs-free
 `mr_<digest>.tar.gz` bundle consumed by verifier/KMS image-download endpoints.

os/image/assemble.sh

@@ -545,14 +545,18 @@ if [[ "$UKI_CREATED" = "1" ]]; then
         echo "measurement.aws.cbor must be fixed at assemble time for a stable os_image_hash." >&2
         exit 1
     fi
-    echo "Generating AWS PCRs via host ${pcr_compute_bin}"
+    echo "Generating AWS PCRs and replay events via host ${pcr_compute_bin}"
     pcr_args=(--image "$uki_abs")
     # Secure Boot variable stores (optional; affects PCR7)
     [[ -n "${NITRO_TPM_PCR_PK:-}" ]] && pcr_args+=(--PK "$NITRO_TPM_PCR_PK")
     [[ -n "${NITRO_TPM_PCR_KEK:-}" ]] && pcr_args+=(--KEK "$NITRO_TPM_PCR_KEK")
     [[ -n "${NITRO_TPM_PCR_DB:-}" ]] && pcr_args+=(--db "$NITRO_TPM_PCR_DB")
-    pcr_json=$("$pcr_compute_bin" "${pcr_args[@]}") \
+    pcr_trace="${OUTPUT_DIR}/aws-pcr-compute.trace"
+    pcr_json_path="${OUTPUT_DIR}/aws-pcrs.json"
+    pcr_json=$(RUST_LOG=nitro_tpm_pcr_compute=debug \
+        "$pcr_compute_bin" "${pcr_args[@]}" 2>"$pcr_trace") \
         || { echo "Error: nitro-tpm-pcr-compute failed" >&2; exit 1; }
+    printf '%s\n' "$pcr_json" > "$pcr_json_path"
 
     pcr4=$(jq -r '.Measurements.PCR4 // empty' <<<"$pcr_json")
     pcr7=$(jq -r '.Measurements.PCR7 // empty' <<<"$pcr_json")
@@ -565,8 +569,11 @@ if [[ "$UKI_CREATED" = "1" ]]; then
     echo "Generating measurement.aws.cbor via ${DSTACK_MR_BIN}"
     "${DSTACK_MR_BIN}" aws-measurement-cbor "$pcr4" "$pcr7" "$pcr12" \
         > "${OUTPUT_DIR}/measurement.aws.cbor"
-    # Keep a machine-readable side-car for verifier-side PCR comparison.
-    printf '%s\n' "$pcr_json" > "${OUTPUT_DIR}/aws-pcrs.json"
+    python3 "$(dirname "$0")/aws-pcr-replay.py" \
+        --trace "$pcr_trace" \
+        --measurements "$pcr_json_path" \
+        --output "${OUTPUT_DIR}/measurement.aws.replay.json"
+    rm "$pcr_trace"
     HAVE_MEASUREMENT_AWS=1
 fi
 
@@ -579,7 +586,7 @@ if [ "$HAVE_MEASUREMENT_GCP" = "1" ]; then
     CHECKSUM_FILES+=(measurement.gcp.cbor)
 fi
 if [ "$HAVE_MEASUREMENT_AWS" = "1" ]; then
-    CHECKSUM_FILES+=(measurement.aws.cbor)
+    CHECKSUM_FILES+=(measurement.aws.cbor measurement.aws.replay.json)
 fi
 (
     cd "${OUTPUT_DIR}/"
@@ -605,7 +612,7 @@ if [ "$DSTACK_TAR_RELEASE" = "1" ]; then
         BARE_METAL_FILES+=(measurement.gcp.cbor)
     fi
     if [ "$HAVE_MEASUREMENT_AWS" = "1" ]; then
-        BARE_METAL_FILES+=(measurement.aws.cbor)
+        BARE_METAL_FILES+=(measurement.aws.cbor measurement.aws.replay.json)
     fi
     BARE_METAL_TAR_FILES=()
     for file in "${BARE_METAL_FILES[@]}"; do
@@ -618,7 +625,7 @@ if [ "$DSTACK_TAR_RELEASE" = "1" ]; then
     if [[ "$UKI_CREATED" = "1" ]]; then
         rm -rf "${IMAGE_TAR_UKI}"
         echo "Archiving UKI image to ${IMAGE_TAR_UKI}"
-        UKI_FILES=(disk.raw digest.txt sha256sum.txt measurement.gcp.cbor measurement.aws.cbor)
+        UKI_FILES=(disk.raw digest.txt sha256sum.txt measurement.gcp.cbor measurement.aws.cbor measurement.aws.replay.json)
         UKI_TAR_FILES=()
         for file in "${UKI_FILES[@]}"; do
             UKI_TAR_FILES+=("$TAR_DIR_NAME/$file")

os/image/aws-pcr-replay.py

@@ -0,0 +1,62 @@
+#!/usr/bin/env python3
+# SPDX-FileCopyrightText: Copyright (c) 2026 Phala Network
+# SPDX-License-Identifier: Apache-2.0
+
+"""Build a NitroTPM PCR replay document from the pinned AWS tool trace."""
+
+import argparse
+import hashlib
+import json
+import re
+from pathlib import Path
+
+EVENT_RE = re.compile(r"\[PCR(4|7|12)\]\s+([A-Z0-9_]+):\s+SHA384:([0-9a-fA-F]{96})")
+
+
+def extend(current: bytes, digest: bytes) -> bytes:
+    """Extend a SHA-384 PCR value with one measured digest."""
+    return hashlib.sha384(current + digest).digest()
+
+
+def main() -> None:
+    """Generate and validate the NitroTPM PCR replay document."""
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--trace", type=Path, required=True)
+    parser.add_argument("--measurements", type=Path, required=True)
+    parser.add_argument("--output", type=Path, required=True)
+    args = parser.parse_args()
+
+    events = []
+    replayed = {4: bytes(48), 7: bytes(48), 12: bytes(48)}
+    for match in EVENT_RE.finditer(args.trace.read_text()):
+        pcr = int(match.group(1))
+        digest = bytes.fromhex(match.group(3))
+        replayed[pcr] = extend(replayed[pcr], digest)
+        events.append(
+            {"pcr": pcr, "event_type": match.group(2), "digest": digest.hex()}
+        )
+
+    if not events:
+        raise SystemExit("no NitroTPM PCR events found in tool trace")
+
+    computed = json.loads(args.measurements.read_text())["Measurements"]
+    expected = {pcr: computed[f"PCR{pcr}"].lower() for pcr in replayed}
+    for pcr, value in replayed.items():
+        if value.hex() != expected[pcr]:
+            raise SystemExit(
+                f"PCR{pcr} replay mismatch: expected={expected[pcr]}, "
+                f"replayed={value.hex()}"
+            )
+
+    document = {
+        "version": 1,
+        "events": events,
+        "pcr4": expected[4],
+        "pcr7": expected[7],
+        "pcr12": expected[12],
+    }
+    args.output.write_text(json.dumps(document, indent=2) + "\n")
+
+
+if __name__ == "__main__":
+    main()