refactor(simulator): move simulator helpers out of guest-agent crate

Author: kvinwang

Cargo.lock

@@ -17,6 +17,7 @@ path = "src/main.rs"
 anyhow.workspace = true
 clap.workspace = true
 serde.workspace = true
+serde_json.workspace = true
 tracing.workspace = true
 tracing-subscriber.workspace = true
 rocket.workspace = true

guest-agent-simulator/Cargo.toml

@@ -2,16 +2,14 @@
 //
 // SPDX-License-Identifier: Apache-2.0
 
+mod simulator;
+
 use std::sync::Arc;
 
 use anyhow::{bail, Context, Result};
 use clap::Parser;
 use dstack_guest_agent::{
-    backend::{
-        load_versioned_attestation, simulated_attest_response,
-        simulated_certificate_attestation, simulated_info_attestation,
-        simulated_quote_response, PlatformBackend,
-    },
+    backend::PlatformBackend,
     config::{self, Config},
     AppState, run_server,
 };
@@ -58,20 +56,20 @@ impl SimulatorPlatform {
 }
 
 impl PlatformBackend for SimulatorPlatform {
-    fn attestation_for_info(&self) -> Result<Option<Attestation>> {
-        Ok(Some(simulated_info_attestation(&self.attestation)))
+    fn attestation_for_info(&self) -> Result<Attestation> {
+        Ok(simulator::simulated_info_attestation(&self.attestation))
     }
 
     fn certificate_attestation(&self, pubkey: &[u8]) -> Result<VersionedAttestation> {
-        Ok(simulated_certificate_attestation(&self.attestation, pubkey))
+        Ok(simulator::simulated_certificate_attestation(&self.attestation, pubkey))
     }
 
     fn quote_response(&self, report_data: [u8; 64], vm_config: &str) -> Result<GetQuoteResponse> {
-        simulated_quote_response(&self.attestation, report_data, vm_config)
+        simulator::simulated_quote_response(&self.attestation, report_data, vm_config)
     }
 
     fn attest_response(&self, _report_data: [u8; 64]) -> Result<AttestResponse> {
-        Ok(simulated_attest_response(&self.attestation))
+        Ok(simulator::simulated_attest_response(&self.attestation))
     }
 
     fn emit_event(&self, event: &str, _payload: &[u8]) -> Result<()> {
@@ -93,7 +91,7 @@ async fn main() -> Result<()> {
         .extract()
         .context("Failed to extract simulator core config")?;
     warn!(attestation_file = %sim_config.simulator.attestation_file, "starting dstack guest-agent simulator");
-    let attestation = load_versioned_attestation(&sim_config.simulator.attestation_file)?;
+    let attestation = simulator::load_versioned_attestation(&sim_config.simulator.attestation_file)?;
     let state = AppState::new(sim_config.core, Arc::new(SimulatorPlatform::new(attestation)))
         .await
         .context("Failed to create simulator app state")?;
@@ -106,7 +104,7 @@ mod tests {
     use super::*;
 
     fn load_fixture_platform() -> SimulatorPlatform {
-        let fixture = load_versioned_attestation(
+        let fixture = simulator::load_versioned_attestation(
             std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join("../guest-agent/fixtures/attestation.bin"),
         )
         .expect("fixture attestation should load");
@@ -125,6 +123,6 @@ mod tests {
         let platform = load_fixture_platform();
         let cert_attestation = platform.certificate_attestation(b"test-public-key").unwrap();
         assert!(cert_attestation.decode_app_info(false).is_ok());
-        assert!(platform.attestation_for_info().unwrap().is_some());
+        let _ = platform.attestation_for_info().unwrap();
     }
 }

guest-agent-simulator/src/main.rs

@@ -0,0 +1,66 @@
+// SPDX-FileCopyrightText: © 2025 Phala Network <dstack@phala.network>
+//
+// SPDX-License-Identifier: Apache-2.0
+
+use std::path::Path;
+
+use anyhow::{Context, Result};
+use dstack_guest_agent_rpc::{AttestResponse, GetQuoteResponse};
+use std::fs;
+use ra_rpc::Attestation;
+use ra_tls::attestation::{
+    QuoteContentType, VersionedAttestation, TDX_QUOTE_REPORT_DATA_RANGE,
+};
+
+pub fn load_versioned_attestation(path: impl AsRef<Path>) -> Result<VersionedAttestation> {
+    let path = path.as_ref();
+    let attestation_bytes = fs::read(path).with_context(|| {
+        format!(
+            "Failed to read simulator attestation file: {}",
+            path.display()
+        )
+    })?;
+    VersionedAttestation::from_scale(&attestation_bytes)
+        .context("Failed to decode simulator attestation")
+}
+
+pub fn simulated_quote_response(
+    attestation: &VersionedAttestation,
+    report_data: [u8; 64],
+    vm_config: &str,
+) -> Result<GetQuoteResponse> {
+    let VersionedAttestation::V0 { attestation } = attestation.clone();
+    let mut attestation = attestation;
+    let Some(quote) = attestation.tdx_quote_mut() else {
+        return Err(anyhow::anyhow!("Quote not found"));
+    };
+
+    quote.quote[TDX_QUOTE_REPORT_DATA_RANGE].copy_from_slice(&report_data);
+    Ok(GetQuoteResponse {
+        quote: quote.quote.to_vec(),
+        event_log: serde_json::to_string(&quote.event_log)
+            .context("Failed to serialize event log")?,
+        report_data: report_data.to_vec(),
+        vm_config: vm_config.to_string(),
+    })
+}
+
+pub fn simulated_attest_response(attestation: &VersionedAttestation) -> AttestResponse {
+    AttestResponse {
+        attestation: attestation.to_scale(),
+    }
+}
+
+pub fn simulated_info_attestation(attestation: &VersionedAttestation) -> Attestation {
+    attestation.clone().into_inner()
+}
+
+pub fn simulated_certificate_attestation(
+    attestation: &VersionedAttestation,
+    pubkey: &[u8],
+) -> VersionedAttestation {
+    let mut attestation = attestation.clone();
+    let report_data = QuoteContentType::RaTlsCert.to_report_data(pubkey);
+    attestation.set_report_data(report_data);
+    attestation
+}

guest-agent-simulator/src/simulator.rs

@@ -2,19 +2,14 @@
 //
 // SPDX-License-Identifier: Apache-2.0
 
-use std::path::Path;
-
 use anyhow::{Context, Result};
 use dstack_attest::emit_runtime_event;
 use dstack_guest_agent_rpc::{AttestResponse, GetQuoteResponse};
-use fs_err as fs;
 use ra_rpc::Attestation;
-use ra_tls::attestation::{
-    QuoteContentType, VersionedAttestation, TDX_QUOTE_REPORT_DATA_RANGE,
-};
+use ra_tls::attestation::{QuoteContentType, VersionedAttestation};
 
 pub trait PlatformBackend: Send + Sync {
-    fn attestation_for_info(&self) -> Result<Option<Attestation>>;
+    fn attestation_for_info(&self) -> Result<Attestation>;
     fn certificate_attestation(&self, pubkey: &[u8]) -> Result<VersionedAttestation>;
     fn quote_response(&self, report_data: [u8; 64], vm_config: &str) -> Result<GetQuoteResponse>;
     fn attest_response(&self, report_data: [u8; 64]) -> Result<AttestResponse>;
@@ -25,8 +20,8 @@ pub trait PlatformBackend: Send + Sync {
 pub struct RealPlatform;
 
 impl PlatformBackend for RealPlatform {
-    fn attestation_for_info(&self) -> Result<Option<Attestation>> {
-        Ok(Attestation::local().ok())
+    fn attestation_for_info(&self) -> Result<Attestation> {
+        Attestation::local().context("Failed to get local attestation")
     }
 
     fn certificate_attestation(&self, pubkey: &[u8]) -> Result<VersionedAttestation> {
@@ -59,56 +54,3 @@ impl PlatformBackend for RealPlatform {
         emit_runtime_event(event, payload)
     }
 }
-
-pub fn load_versioned_attestation(path: impl AsRef<Path>) -> Result<VersionedAttestation> {
-    let path = path.as_ref();
-    let attestation_bytes = fs::read(path).with_context(|| {
-        format!(
-            "Failed to read simulator attestation file: {}",
-            path.display()
-        )
-    })?;
-    VersionedAttestation::from_scale(&attestation_bytes)
-        .context("Failed to decode simulator attestation")
-}
-
-pub fn simulated_quote_response(
-    attestation: &VersionedAttestation,
-    report_data: [u8; 64],
-    vm_config: &str,
-) -> Result<GetQuoteResponse> {
-    let VersionedAttestation::V0 { attestation } = attestation.clone();
-    let mut attestation = attestation;
-    let Some(quote) = attestation.tdx_quote_mut() else {
-        return Err(anyhow::anyhow!("Quote not found"));
-    };
-
-    quote.quote[TDX_QUOTE_REPORT_DATA_RANGE].copy_from_slice(&report_data);
-    Ok(GetQuoteResponse {
-        quote: quote.quote.to_vec(),
-        event_log: serde_json::to_string(&quote.event_log)
-            .context("Failed to serialize event log")?,
-        report_data: report_data.to_vec(),
-        vm_config: vm_config.to_string(),
-    })
-}
-
-pub fn simulated_attest_response(attestation: &VersionedAttestation) -> AttestResponse {
-    AttestResponse {
-        attestation: attestation.to_scale(),
-    }
-}
-
-pub fn simulated_info_attestation(attestation: &VersionedAttestation) -> Attestation {
-    attestation.clone().into_inner()
-}
-
-pub fn simulated_certificate_attestation(
-    attestation: &VersionedAttestation,
-    pubkey: &[u8],
-) -> VersionedAttestation {
-    let mut attestation = attestation.clone();
-    let report_data = QuoteContentType::RaTlsCert.to_report_data(pubkey);
-    attestation.set_report_data(report_data);
-    attestation
-}

guest-agent/src/backend.rs

@@ -62,7 +62,7 @@ struct AppStateInner {
 }
 
 impl AppStateInner {
-    fn info_attestation(&self) -> Result<Option<ra_rpc::Attestation>> {
+    fn info_attestation(&self) -> Result<ra_rpc::Attestation> {
         self.platform.attestation_for_info()
     }
 
@@ -178,9 +178,7 @@ pub struct InternalRpcHandler {
 
 pub async fn get_info(state: &AppState, external: bool) -> Result<AppInfo> {
     let hide_tcb_info = external && !state.config().app_compose.public_tcbinfo;
-    let Some(attestation) = state.inner.info_attestation()? else {
-        return Ok(AppInfo::default());
-    };
+    let attestation = state.inner.info_attestation()?;
     let app_info = attestation
         .decode_app_info(false)
         .context("Failed to decode app info")?;
@@ -651,14 +649,7 @@ impl RpcCall<AppState> for ExternalRpcHandler {
 #[cfg(test)]
 mod tests {
     use super::*;
-    use crate::{
-        backend::{
-            load_versioned_attestation, simulated_attest_response,
-            simulated_certificate_attestation, simulated_info_attestation,
-            simulated_quote_response, PlatformBackend,
-        },
-        config::{AppComposeWrapper, Config},
-    };
+    use crate::{backend::PlatformBackend, config::{AppComposeWrapper, Config}};
     use dstack_guest_agent_rpc::{GetAttestationForAppKeyRequest, SignRequest};
     use ra_tls::attestation::VersionedAttestation;
     use dstack_types::{AppCompose, AppKeys, KeyProvider};
@@ -807,24 +798,40 @@ pNs85uhOZE8z2jr8Pg==
         }
 
         impl PlatformBackend for TestSimulatorPlatform {
-            fn attestation_for_info(&self) -> Result<Option<ra_rpc::Attestation>> {
-                Ok(Some(simulated_info_attestation(&self.attestation)))
+            fn attestation_for_info(&self) -> Result<ra_rpc::Attestation> {
+                Ok(self.attestation.clone().into_inner())
             }
 
             fn certificate_attestation(&self, pubkey: &[u8]) -> Result<VersionedAttestation> {
-                Ok(simulated_certificate_attestation(&self.attestation, pubkey))
+                let mut attestation = self.attestation.clone();
+                let report_data = ra_tls::attestation::QuoteContentType::RaTlsCert.to_report_data(pubkey);
+                attestation.set_report_data(report_data);
+                Ok(attestation)
             }
 
             fn quote_response(
                 &self,
                 report_data: [u8; 64],
                 vm_config: &str,
             ) -> Result<GetQuoteResponse> {
-                simulated_quote_response(&self.attestation, report_data, vm_config)
+                let ra_tls::attestation::VersionedAttestation::V0 { attestation } = self.attestation.clone();
+                let mut attestation = attestation;
+                let Some(quote) = attestation.tdx_quote_mut() else {
+                    return Err(anyhow::anyhow!("Quote not found"));
+                };
+                quote.quote[ra_tls::attestation::TDX_QUOTE_REPORT_DATA_RANGE].copy_from_slice(&report_data);
+                Ok(GetQuoteResponse {
+                    quote: quote.quote.to_vec(),
+                    event_log: serde_json::to_string(&quote.event_log).context("Failed to serialize event log")?,
+                    report_data: report_data.to_vec(),
+                    vm_config: vm_config.to_string(),
+                })
             }
 
             fn attest_response(&self, _report_data: [u8; 64]) -> Result<AttestResponse> {
-                Ok(simulated_attest_response(&self.attestation))
+                Ok(AttestResponse {
+                    attestation: self.attestation.to_scale(),
+                })
             }
 
             fn emit_event(&self, _event: &str, _payload: &[u8]) -> Result<()> {
@@ -839,7 +846,7 @@ pNs85uhOZE8z2jr8Pg==
             cert_client: dummy_cert_client,
             demo_cert: RwLock::new(String::new()),
             platform: Arc::new(TestSimulatorPlatform {
-                attestation: load_versioned_attestation(temp_attestation_file.path()).unwrap(),
+                attestation: VersionedAttestation::from_scale(&std::fs::read(temp_attestation_file.path()).unwrap()).unwrap(),
             }),
         };