From a6ee4639129fa2631622f5249e9ec8816949e383 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Wed, 5 Aug 2026 18:31:58 -0700 Subject: [PATCH 1/2] fix(simulator): replay measured NitroTPM boot PCRs --- dstack/dstack-types/src/lib.rs | 24 +++++++++++++ dstack/tee-simulator/src/tpm.rs | 47 +++++++++++++++++++++++++- dstack/vmm/src/app.rs | 42 +++++++++++++++++++++-- dstack/vmm/src/app/image.rs | 21 ++++++++++-- dstack/vmm/src/app/qemu.rs | 1 + os/image/README.md | 9 +++-- os/image/assemble.sh | 21 ++++++++---- os/image/aws-pcr-replay.py | 60 +++++++++++++++++++++++++++++++++ 8 files changed, 210 insertions(+), 15 deletions(-) create mode 100644 os/image/aws-pcr-replay.py diff --git a/dstack/dstack-types/src/lib.rs b/dstack/dstack-types/src/lib.rs index 1eb78a29b..bc1e00cee 100644 --- a/dstack/dstack-types/src/lib.rs +++ b/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, + /// 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, +} + +#[derive(Deserialize, Serialize, Debug, Clone, PartialEq, Eq)] +pub struct AwsPcrReplay { + pub version: u32, + pub events: Vec, + #[serde(with = "hex_bytes")] + pub pcr4: Vec, + #[serde(with = "hex_bytes")] + pub pcr7: Vec, + #[serde(with = "hex_bytes")] + pub pcr12: Vec, +} + +#[derive(Deserialize, Serialize, Debug, Clone, PartialEq, Eq)] +pub struct AwsPcrReplayEvent { + pub pcr: u16, + pub event_type: String, + #[serde(with = "hex_bytes")] + pub digest: Vec, } #[derive(Deserialize, Serialize, Debug, Clone, Copy, Default, PartialEq, Eq, Encode, Decode)] diff --git a/dstack/tee-simulator/src/tpm.rs b/dstack/tee-simulator/src/tpm.rs index 6272a0692..3e319841e 100644 --- a/dstack/tee-simulator/src/tpm.rs +++ b/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) diff --git a/dstack/vmm/src/app.rs b/dstack/vmm/src/app.rs index 7a9a32b81..7508592ce 100644 --- a/dstack/vmm/src/app.rs +++ b/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, } } diff --git a/dstack/vmm/src/app/image.rs b/dstack/vmm/src/app/image.rs index 2b74f2b2b..274b4e9c7 100644 --- a/dstack/vmm/src/app/image.rs +++ b/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, /// AWS NitroTPM no-image-download measurement material. pub aws_measurement: Option, + /// AWS boot events consumed only by the development NitroTPM simulator. + pub aws_pcr_replay: Option, } 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() } diff --git a/dstack/vmm/src/app/qemu.rs b/dstack/vmm/src/app/qemu.rs index f84b44fc8..327af5e75 100644 --- a/dstack/vmm/src/app/qemu.rs +++ b/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"), diff --git a/os/image/README.md b/os/image/README.md index dd6244d05..1b51c4bb4 100644 --- a/os/image/README.md +++ b/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 ` creates the flattened, rootfs-free `mr_.tar.gz` bundle consumed by verifier/KMS image-download endpoints. diff --git a/os/image/assemble.sh b/os/image/assemble.sh index 9764f955f..eca70594d 100755 --- a/os/image/assemble.sh +++ b/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") diff --git a/os/image/aws-pcr-replay.py b/os/image/aws-pcr-replay.py new file mode 100644 index 000000000..2a15c854a --- /dev/null +++ b/os/image/aws-pcr-replay.py @@ -0,0 +1,60 @@ +#!/usr/bin/env python3 +"""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: + return hashlib.sha384(current + digest).digest() + + +def main() -> None: + 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() From f59fcc00b7a1f15d7a18ea8bc6976715dfbc73af Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Wed, 5 Aug 2026 19:41:32 -0700 Subject: [PATCH 2/2] chore(image): satisfy replay helper linters --- os/image/aws-pcr-replay.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/os/image/aws-pcr-replay.py b/os/image/aws-pcr-replay.py index 2a15c854a..e27431064 100644 --- a/os/image/aws-pcr-replay.py +++ b/os/image/aws-pcr-replay.py @@ -1,4 +1,7 @@ #!/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 @@ -7,17 +10,16 @@ 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})" -) +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)