From 0c694862b7d0290f8a6da55d6f9f2019a52345de Mon Sep 17 00:00:00 2001
From: Kevin Wang <wy721@qq.com>
Date: Tue, 28 Jul 2026 08:52:29 +0000
Subject: [PATCH] fix(simulator): model measured Nitro enclave PCRs

---
 dstack/tee-simulator/src/nsm.rs              | 22 +++++++++++++++++++-
 dstack/tests/e2e/attestation/run-platform.sh | 10 +++++++++
 2 files changed, 31 insertions(+), 1 deletion(-)

diff --git a/dstack/tee-simulator/src/nsm.rs b/dstack/tee-simulator/src/nsm.rs
index 7a4ac8eab..b2cd1d5b1 100644
--- a/dstack/tee-simulator/src/nsm.rs
+++ b/dstack/tee-simulator/src/nsm.rs
@@ -46,6 +46,15 @@ impl Default for NsmState {
 }
 
 impl NsmState {
+    fn measured() -> Self {
+        let mut state = Self::default();
+        for index in 0..=2 {
+            let digest = Sha384::digest(format!("dstack-tee-simulator/nsm/pcr/{index}").as_bytes());
+            state.pcrs[index].copy_from_slice(&digest);
+        }
+        state
+    }
+
     fn handle(&mut self, generator: &NsmGenerator, request: Request) -> Response {
         match request {
             Request::DescribePCR { index } => {
@@ -376,7 +385,7 @@ pub fn run(config: &TeeSimulatorConfig) -> Result<()> {
         .set(Arc::new(NsmGenerator::from_seed(parse_seed(seed)?)?))
         .map_err(|_| anyhow::anyhow!("NSM simulator was already initialized"))?;
     STATE
-        .set(Mutex::new(NsmState::default()))
+        .set(Mutex::new(NsmState::measured()))
         .map_err(|_| anyhow::anyhow!("NSM simulator state was already initialized"))?;
 
     let device_name = CString::new("DEVNAME=nsm")?;
@@ -441,6 +450,17 @@ mod tests {
         }
     }
 
+    #[test]
+    fn measured_state_models_a_production_enclave() {
+        let state = NsmState::measured();
+        assert!(state.pcrs[..=2]
+            .iter()
+            .all(|value| value.iter().any(|byte| *byte != 0)));
+        assert!(state.pcrs[3..]
+            .iter()
+            .all(|value| value.iter().all(|byte| *byte == 0)));
+    }
+
     #[test]
     fn pcr_lifecycle_matches_nsm_semantics() {
         let generator = NsmGenerator::from_seed([0x51; 32]).unwrap();
diff --git a/dstack/tests/e2e/attestation/run-platform.sh b/dstack/tests/e2e/attestation/run-platform.sh
index 8b8fc2bc9..394670e48 100755
--- a/dstack/tests/e2e/attestation/run-platform.sh
+++ b/dstack/tests/e2e/attestation/run-platform.sh
@@ -38,6 +38,16 @@ elif [[ "$TEE_PLATFORM" == dstack-amd-sev-snp ]]; then
   dstack-util attest-json --input "$WORK/snp-fixture.bin" --output "$WORK/snp-fixture.json"
   VM_CONFIG=$(jq -c '.config | fromjson' "$WORK/snp-fixture.json")
   MR_CONFIG=$(jq -r .mr_config <<<"$VM_CONFIG")
+elif [[ "$TEE_PLATFORM" == dstack-nitro-enclave ]]; then
+  for index in 0 1 2; do
+    printf 'dstack-tee-simulator/nsm/pcr/%s' "$index" |
+      sha384sum | cut -d' ' -f1 > "$WORK/pcr$index"
+  done
+  OS_IMAGE_HASH=$(
+    cat "$WORK/pcr0" "$WORK/pcr1" "$WORK/pcr2" |
+      xxd -r -p | sha256sum | cut -d' ' -f1
+  )
+  VM_CONFIG=$(jq -cn --arg os "$OS_IMAGE_HASH" '{os_image_hash:$os}')
 elif [[ "$TEE_PLATFORM" == dstack-aws-nitro-tpm ]]; then
   ZERO_PCR=$(printf '00%.0s' $(seq 1 48))
   printf '%s' "$ZERO_PCR" > "$WORK/pcr4"