[STACKED on #851] fix(simulator): replay measured NitroTPM boot PCRs

dstack/dstack-types/src/lib.rs

@@ -940,6 +940,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/Cargo.toml

@@ -34,6 +34,7 @@ libloading = "0.8"
 pem.workspace = true
 aws-nitro-enclaves-nsm-api = "0.4"
 serde_cbor = "0.11"
+tpm2.workspace = true
 
 [dev-dependencies]
 dcap-qvl.workspace = true

dstack/tee-simulator/src/tpm.rs

@@ -19,8 +19,9 @@ 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::{TpmAlgId, TpmContext};
 
 const AK_ECC_CERT: &str = "0x01c10002";
 const AK_ECC_TEMPLATE: &str = "0x01c10003";
@@ -293,6 +294,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;
@@ -323,6 +330,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)
@@ -375,7 +421,7 @@ fn proxy_tpm_commands(
             let template = nv_write
                 .as_ref()
                 .context("NitroTPM vendor command without an NV request")?;
-            nsm_response = Some(handle_nsm_vendor_command(generator, template)?);
+            nsm_response = Some(handle_nsm_vendor_command(generator, template, backend)?);
             tpm_success_response()
         } else if code == TPM2_CC_NV_READ && nsm_response.is_some() {
             nv_read_response(
@@ -442,14 +488,21 @@ fn parse_nv_write(command: &[u8]) -> Result<Option<NvWriteTemplate>> {
 fn handle_nsm_vendor_command(
     generator: &NsmGenerator,
     template: &NvWriteTemplate,
+    backend: &mut UnixStream,
 ) -> Result<Vec<u8>> {
     let request: NsmRequest = serde_cbor::from_slice(&template.request)?;
     let response = match request {
         NsmRequest::Attestation { user_data, .. } => {
+            let mut tpm = TpmContext::from_stream(
+                backend
+                    .try_clone()
+                    .context("failed to clone NitroTPM stream")?,
+                "NitroTPM backend",
+            );
             let pcrs = [4u16, 7, 8, 12, 14]
                 .into_iter()
-                .map(|i| (i, vec![0; 48]))
-                .collect();
+                .map(|index| Ok((index, tpm.pcr_read_single(index.into(), TpmAlgId::Sha384)?)))
+                .collect::<Result<_>>()?;
             let document = generator.attest_with_pcrs(
                 user_data.as_ref().map(|v| v.as_slice()).unwrap_or_default(),
                 pcrs,

dstack/tpm2/src/commands.rs

@@ -8,6 +8,7 @@
 
 use anyhow::{Context, Result};
 use std::collections::HashSet;
+use std::io::{Read, Write};
 use tracing::debug;
 
 use super::constants::*;
@@ -32,6 +33,13 @@ impl TpmContext {
         Ok(Self { device })
     }
 
+    /// Create a TPM context over an already connected byte stream.
+    pub fn from_stream(stream: impl Read + Write + 'static, name: impl Into<String>) -> Self {
+        Self {
+            device: TpmDevice::from_stream(stream, name),
+        }
+    }
+
     /// Get the device path
     pub fn device_path(&self) -> &str {
         self.device.path()

dstack/tpm2/src/device.rs

@@ -7,7 +7,7 @@
 //! Provides low-level communication with TPM devices via /dev/tpmrm0 or /dev/tpm0.
 
 use anyhow::{bail, Context, Result};
-use std::fs::{File, OpenOptions};
+use std::fs::OpenOptions;
 use std::io::{Read, Write};
 use std::path::Path;
 
@@ -18,8 +18,12 @@ use super::marshal::*;
 const TPM_MAX_COMMAND_SIZE: usize = 4096;
 
 /// TPM device handle
+trait ReadWrite: Read + Write {}
+
+impl<T: Read + Write> ReadWrite for T {}
+
 pub struct TpmDevice {
-    file: File,
+    transport: Box<dyn ReadWrite>,
     path: String,
 }
 
@@ -36,11 +40,19 @@ impl TpmDevice {
             .with_context(|| format!("failed to open TPM device: {}", device_path))?;
 
         Ok(Self {
-            file,
+            transport: Box::new(file),
             path: device_path.to_string(),
         })
     }
 
+    /// Create a TPM device over an already connected byte stream.
+    pub fn from_stream(stream: impl Read + Write + 'static, name: impl Into<String>) -> Self {
+        Self {
+            transport: Box::new(stream),
+            path: name.into(),
+        }
+    }
+
     /// Detect and open the default TPM device
     pub fn detect() -> Result<Self> {
         if Path::new("/dev/tpmrm0").exists() {
@@ -59,19 +71,26 @@ impl TpmDevice {
 
     /// Send a command to the TPM and receive the response
     pub fn transmit(&mut self, command: &[u8]) -> Result<Vec<u8>> {
-        // Write command
-        self.file
+        self.transport
             .write_all(command)
             .context("failed to write TPM command")?;
 
-        // Read response
-        let mut response = vec![0u8; TPM_MAX_COMMAND_SIZE];
-        let n = self
-            .file
-            .read(&mut response)
-            .context("failed to read TPM response")?;
-
-        response.truncate(n);
+        // Read the fixed header first so stream transports cannot return a
+        // partial response and leave bytes for the next transaction.
+        let mut header = [0u8; 10];
+        self.transport
+            .read_exact(&mut header)
+            .context("failed to read TPM response header")?;
+        let response_size = u32::from_be_bytes(header[2..6].try_into().unwrap()) as usize;
+        if !(header.len()..=TPM_MAX_COMMAND_SIZE).contains(&response_size) {
+            bail!("invalid TPM response size: {response_size}");
+        }
+        let mut response = Vec::with_capacity(response_size);
+        response.extend_from_slice(&header);
+        response.resize(response_size, 0);
+        self.transport
+            .read_exact(&mut response[header.len()..])
+            .context("failed to read TPM response body")?;
         Ok(response)
     }
 

dstack/vmm/src/app.rs

@@ -1281,6 +1281,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")
@@ -1490,6 +1497,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")?;
@@ -1548,7 +1578,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": [],
@@ -1565,7 +1598,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());
@@ -1817,6 +1854,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 {
@@ -167,6 +170,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();
@@ -184,6 +198,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"),