Merge pull request #61 from flashbots/peg/test-dcap-verification

Improve logic for test DCAP verification

Author: ameba23

src/attestation/dcap.rs

@@ -24,44 +24,45 @@ pub async fn verify_dcap_attestation(
     expected_input_data: [u8; 64],
     pccs_url: Option<String>,
 ) -> Result<MultiMeasurements, DcapVerificationError> {
-    let measurements = if cfg!(not(test)) {
-        let now = std::time::SystemTime::now()
-            .duration_since(std::time::UNIX_EPOCH)?
-            .as_secs();
-        let quote = Quote::parse(&input)?;
-        tracing::info!("Verifying DCAP attestation: {quote:?}");
+    let now = std::time::SystemTime::now()
+        .duration_since(std::time::UNIX_EPOCH)?
+        .as_secs();
+    let quote = Quote::parse(&input)?;
+    tracing::info!("Verifying DCAP attestation: {quote:?}");
 
-        let ca = quote.ca()?;
-        let fmspc = hex::encode_upper(quote.fmspc()?);
-        let collateral = get_collateral_for_fmspc(
-            &pccs_url.clone().unwrap_or(PCS_URL.to_string()),
-            fmspc,
-            ca,
-            false, // Indicates not SGX
-        )
-        .await?;
-
-        let _verified_report = dcap_qvl::verify::verify(&input, &collateral, now)?;
+    let ca = quote.ca()?;
+    let fmspc = hex::encode_upper(quote.fmspc()?);
+    let collateral = get_collateral_for_fmspc(
+        &pccs_url.clone().unwrap_or(PCS_URL.to_string()),
+        fmspc,
+        ca,
+        false, // Indicates not SGX
+    )
+    .await?;
 
-        let measurements = MultiMeasurements::from_dcap_qvl_quote(&quote)?;
+    let _verified_report = dcap_qvl::verify::verify(&input, &collateral, now)?;
 
-        if get_quote_input_data(quote.report) != expected_input_data {
-            return Err(DcapVerificationError::InputMismatch);
-        }
-        measurements
-    } else {
-        // In tests we use mock quotes which will fail to verify
-        let quote = tdx_quote::Quote::from_bytes(&input)?;
-        if quote.report_input_data() != expected_input_data {
-            return Err(DcapVerificationError::InputMismatch);
-        }
+    let measurements = MultiMeasurements::from_dcap_qvl_quote(&quote)?;
 
-        MultiMeasurements::from_tdx_quote(&quote)
-    };
+    if get_quote_input_data(quote.report) != expected_input_data {
+        return Err(DcapVerificationError::InputMismatch);
+    }
 
     Ok(measurements)
 }
 
+pub fn mock_verify_dcap(
+    input: Vec<u8>,
+    expected_input_data: [u8; 64],
+) -> Result<MultiMeasurements, DcapVerificationError> {
+    // In tests we use mock quotes which will fail to verify
+    let quote = tdx_quote::Quote::from_bytes(&input)?;
+    if quote.report_input_data() != expected_input_data {
+        return Err(DcapVerificationError::InputMismatch);
+    }
+    Ok(MultiMeasurements::from_tdx_quote(&quote))
+}
+
 /// Create a mock quote for testing on non-confidential hardware
 #[cfg(test)]
 fn generate_quote(input: [u8; 64]) -> Result<Vec<u8>, QuoteGenerationError> {

src/attestation/mod.rs

@@ -285,12 +285,19 @@ impl AttestationVerifier {
                 .await?
             }
             _ => {
-                dcap::verify_dcap_attestation(
-                    attestation_exchange_message.attestation,
-                    expected_input_data,
-                    self.pccs_url.clone(),
-                )
-                .await?
+                if cfg!(test) {
+                    dcap::mock_verify_dcap(
+                        attestation_exchange_message.attestation,
+                        expected_input_data,
+                    )?
+                } else {
+                    dcap::verify_dcap_attestation(
+                        attestation_exchange_message.attestation,
+                        expected_input_data,
+                        self.pccs_url.clone(),
+                    )
+                    .await?
+                }
             }
         };