refactor(cert): remove attestation override from shared flow

kvinwang · kvinwang · commit 1b70779fcf30 · 2026-03-19T13:51:42.000Z
diff --git a/cert-client/src/lib.rs b/cert-client/src/lib.rs
@@ -93,23 +93,24 @@ impl CertRequestClient {
         }
     }
 
-    pub async fn request_cert(
+    pub async fn request_cert(&self, key: &KeyPair, config: CertConfigV2) -> Result<Vec<String>> {
+        let pubkey = key.public_key_der();
+        let report_data = QuoteContentType::RaTlsCert.to_report_data(&pubkey);
+        let attestation = ra_rpc::Attestation::quote(&report_data)
+            .context("Failed to get quote for cert pubkey")?
+            .into_versioned();
+        self.sign_cert_with_attestation(key, config, attestation).await
+    }
+
+    pub async fn sign_cert_with_attestation(
         &self,
         key: &KeyPair,
         config: CertConfigV2,
-        attestation_override: Option<VersionedAttestation>,
+        mut attestation: VersionedAttestation,
     ) -> Result<Vec<String>> {
         let pubkey = key.public_key_der();
         let report_data = QuoteContentType::RaTlsCert.to_report_data(&pubkey);
-        let attestation = match attestation_override {
-            Some(mut attestation) => {
-                attestation.set_report_data(report_data);
-                attestation
-            }
-            None => ra_rpc::Attestation::quote(&report_data)
-                .context("Failed to get quote for cert pubkey")?
-                .into_versioned(),
-        };
+        attestation.set_report_data(report_data);
 
         let csr = CertSigningRequestV2 {
             confirm: "please sign cert:".to_string(),
diff --git a/guest-agent-simulator/src/main.rs b/guest-agent-simulator/src/main.rs
@@ -8,7 +8,8 @@ use anyhow::{bail, Context, Result};
 use clap::Parser;
 use dstack_guest_agent::{
     backend::{
-        load_versioned_attestation, simulated_attest_response, simulated_info_attestation,
+        load_versioned_attestation, simulated_attest_response,
+        simulated_certificate_attestation, simulated_info_attestation,
         simulated_quote_response, PlatformBackend,
     },
     config::{self, Config},
@@ -61,8 +62,8 @@ impl PlatformBackend for SimulatorPlatform {
         Ok(Some(simulated_info_attestation(&self.attestation)))
     }
 
-    fn attestation_override(&self) -> Result<Option<VersionedAttestation>> {
-        Ok(Some(self.attestation.clone()))
+    fn certificate_attestation(&self, pubkey: &[u8]) -> Result<VersionedAttestation> {
+        Ok(simulated_certificate_attestation(&self.attestation, pubkey))
     }
 
     fn quote_response(&self, report_data: [u8; 64], vm_config: &str) -> Result<GetQuoteResponse> {
@@ -120,9 +121,10 @@ mod tests {
     }
 
     #[test]
-    fn simulator_provides_attestation_override() {
+    fn simulator_provides_certificate_attestation() {
         let platform = load_fixture_platform();
-        assert!(platform.attestation_override().unwrap().is_some());
+        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());
     }
 }
diff --git a/guest-agent/src/backend.rs b/guest-agent/src/backend.rs
@@ -9,11 +9,13 @@ 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::{VersionedAttestation, TDX_QUOTE_REPORT_DATA_RANGE};
+use ra_tls::attestation::{
+    QuoteContentType, VersionedAttestation, TDX_QUOTE_REPORT_DATA_RANGE,
+};
 
 pub trait PlatformBackend: Send + Sync {
     fn attestation_for_info(&self) -> Result<Option<Attestation>>;
-    fn attestation_override(&self) -> Result<Option<VersionedAttestation>>;
+    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>;
     fn emit_event(&self, event: &str, payload: &[u8]) -> Result<()>;
@@ -27,8 +29,11 @@ impl PlatformBackend for RealPlatform {
         Ok(Attestation::local().ok())
     }
 
-    fn attestation_override(&self) -> Result<Option<VersionedAttestation>> {
-        Ok(None)
+    fn certificate_attestation(&self, pubkey: &[u8]) -> Result<VersionedAttestation> {
+        let report_data = QuoteContentType::RaTlsCert.to_report_data(pubkey);
+        Ok(Attestation::quote(&report_data)
+            .context("Failed to get quote for cert pubkey")?
+            .into_versioned())
     }
 
     fn quote_response(&self, report_data: [u8; 64], vm_config: &str) -> Result<GetQuoteResponse> {
@@ -97,3 +102,13 @@ pub fn simulated_attest_response(attestation: &VersionedAttestation) -> AttestRe
 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
+}
diff --git a/guest-agent/src/rpc_service.rs b/guest-agent/src/rpc_service.rs
@@ -26,7 +26,7 @@ use k256::ecdsa::SigningKey;
 use or_panic::ResultOrPanic;
 use ra_rpc::{CallContext, RpcCall};
 use ra_tls::{
-    attestation::{QuoteContentType, VersionedAttestation, DEFAULT_HASH_ALGORITHM},
+    attestation::{QuoteContentType, DEFAULT_HASH_ALGORITHM},
     cert::CertConfigV2,
     kdf::{derive_key, derive_p256_key_pair_from_bytes},
 };
@@ -62,20 +62,25 @@ struct AppStateInner {
 }
 
 impl AppStateInner {
-    fn attestation_override(&self) -> Result<Option<VersionedAttestation>> {
-        self.platform.attestation_override()
-    }
-
     fn info_attestation(&self) -> Result<Option<ra_rpc::Attestation>> {
         self.platform.attestation_for_info()
     }
 
+    async fn issue_cert(&self, key: &KeyPair, config: CertConfigV2) -> Result<Vec<String>> {
+        let attestation = self
+            .platform
+            .certificate_attestation(&key.public_key_der())
+            .context("Failed to get certificate attestation")?;
+        self.cert_client
+            .sign_cert_with_attestation(key, config, attestation)
+            .await
+            .context("Failed to sign the CSR")
+    }
+
     async fn request_demo_cert(&self) -> Result<String> {
         let key = KeyPair::generate().context("Failed to generate demo key")?;
-        let attestation_override = self.attestation_override()?;
         let demo_cert = self
-            .cert_client
-            .request_cert(
+            .issue_cert(
                 &key,
                 CertConfigV2 {
                     org_name: None,
@@ -88,7 +93,6 @@ impl AppStateInner {
                     not_after: None,
                     not_before: None,
                 },
-                attestation_override,
             )
             .await
             .context("Failed to get app cert")?
@@ -260,18 +264,7 @@ impl DstackGuestRpc for InternalRpcHandler {
             not_after: request.not_after,
             not_before: request.not_before,
         };
-        let attestation_override = self
-            .state
-            .inner
-            .attestation_override()
-            .context("Failed to load platform attestation override")?;
-        let certificate_chain = self
-            .state
-            .inner
-            .cert_client
-            .request_cert(&derived_key, config, attestation_override)
-            .await
-            .context("Failed to sign the CSR")?;
+        let certificate_chain = self.state.inner.issue_cert(&derived_key, config).await?;
         Ok(GetTlsKeyResponse {
             key: derived_key.serialize_pem(),
             certificate_chain,
@@ -497,18 +490,7 @@ impl TappdRpc for InternalRpcHandlerV0 {
             not_before: None,
             not_after: None,
         };
-        let attestation_override = self
-            .state
-            .inner
-            .attestation_override()
-            .context("Failed to load platform attestation override")?;
-        let certificate_chain = self
-            .state
-            .inner
-            .cert_client
-            .request_cert(&derived_key, config, attestation_override)
-            .await
-            .context("Failed to sign the CSR")?;
+        let certificate_chain = self.state.inner.issue_cert(&derived_key, config).await?;
         Ok(GetTlsKeyResponse {
             key: derived_key.serialize_pem(),
             certificate_chain,
@@ -671,12 +653,14 @@ mod tests {
     use super::*;
     use crate::{
         backend::{
-            load_versioned_attestation, simulated_attest_response, simulated_info_attestation,
+            load_versioned_attestation, simulated_attest_response,
+            simulated_certificate_attestation, simulated_info_attestation,
             simulated_quote_response, PlatformBackend,
         },
         config::{AppComposeWrapper, Config},
     };
     use dstack_guest_agent_rpc::{GetAttestationForAppKeyRequest, SignRequest};
+    use ra_tls::attestation::VersionedAttestation;
     use dstack_types::{AppCompose, AppKeys, KeyProvider};
     use ed25519_dalek::ed25519::signature::hazmat::PrehashVerifier;
     use ed25519_dalek::{
@@ -827,8 +811,8 @@ pNs85uhOZE8z2jr8Pg==
                 Ok(Some(simulated_info_attestation(&self.attestation)))
             }
 
-            fn attestation_override(&self) -> Result<Option<VersionedAttestation>> {
-                Ok(Some(self.attestation.clone()))
+            fn certificate_attestation(&self, pubkey: &[u8]) -> Result<VersionedAttestation> {
+                Ok(simulated_certificate_attestation(&self.attestation, pubkey))
             }
 
             fn quote_response(

Original file line number	Diff line number	Diff line change
`@@ -8,7 +8,8 @@ use anyhow::{bail, Context, Result};`
`8`	`8`	`use clap::Parser;`
`9`	`9`	`use dstack_guest_agent::{`
`10`	`10`	`backend::{`
`11`		`- load_versioned_attestation, simulated_attest_response, simulated_info_attestation,`
	`11`	`+ load_versioned_attestation, simulated_attest_response,`
	`12`	`+ simulated_certificate_attestation, simulated_info_attestation,`
`12`	`13`	`simulated_quote_response, PlatformBackend,`
`13`	`14`	`},`
`14`	`15`	`config::{self, Config},`
`@@ -61,8 +62,8 @@ impl PlatformBackend for SimulatorPlatform {`
`61`	`62`	`Ok(Some(simulated_info_attestation(&self.attestation)))`
`62`	`63`	`}`
`63`	`64`
`64`		`- fn attestation_override(&self) -> Result<Option<VersionedAttestation>> {`
`65`		`- Ok(Some(self.attestation.clone()))`
	`65`	`+ fn certificate_attestation(&self, pubkey: &[u8]) -> Result<VersionedAttestation> {`
	`66`	`+ Ok(simulated_certificate_attestation(&self.attestation, pubkey))`
`66`	`67`	`}`
`67`	`68`
`68`	`69`	`fn quote_response(&self, report_data: [u8; 64], vm_config: &str) -> Result<GetQuoteResponse> {`
`@@ -120,9 +121,10 @@ mod tests {`
`120`	`121`	`}`
`121`	`122`
`122`	`123`	`#[test]`
`123`		`- fn simulator_provides_attestation_override() {`
	`124`	`+ fn simulator_provides_certificate_attestation() {`
`124`	`125`	`let platform = load_fixture_platform();`
`125`		`- assert!(platform.attestation_override().unwrap().is_some());`
	`126`	`+ let cert_attestation = platform.certificate_attestation(b"test-public-key").unwrap();`
	`127`	`+ assert!(cert_attestation.decode_app_info(false).is_ok());`
`126`	`128`	`assert!(platform.attestation_for_info().unwrap().is_some());`
`127`	`129`	`}`
`128`	`130`	`}`