refactor(verifier): make image strategies exhaustive

dstack/verifier/src/verification.rs

@@ -23,7 +23,8 @@ use dstack_mr::{
 use dstack_types::VmConfig;
 use hex_literal::hex;
 use ra_tls::attestation::{
-    AppInfo, Attestation, AttestationQuote, AttestationVerifier, DstackVerifiedReport, NitroPcrs,
+    AppInfo, Attestation, AttestationQuote, AttestationVerifier, DstackAwsNitroTpmQuote,
+    DstackGcpTdxQuote, DstackNitroQuote, DstackVerifiedReport, NitroPcrs, SnpQuote, TdxQuote,
     TpmQuote, VerifiedAttestation, VersionedAttestation,
 };
 use serde::{Deserialize, Serialize};
@@ -190,6 +191,36 @@ struct ImagePaths {
     version: String,
 }
 
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+enum OsImageVerificationStrategy {
+    TdxFullDownload,
+    TdxLiteMeasurement,
+    SevSnpMeasurement,
+    GcpTdxMeasurement,
+    NitroEnclavePcrs,
+    AwsNitroTpmPcrs,
+}
+
+impl OsImageVerificationStrategy {
+    fn select(quote: &AttestationQuote, vm_config: &VmConfig) -> Self {
+        match quote {
+            AttestationQuote::DstackTdx(_) => {
+                if vm_config.tdx_attestation_variant.is_lite()
+                    || vm_config.tdx_measurement.is_some()
+                {
+                    Self::TdxLiteMeasurement
+                } else {
+                    Self::TdxFullDownload
+                }
+            }
+            AttestationQuote::DstackAmdSevSnp(_) => Self::SevSnpMeasurement,
+            AttestationQuote::DstackGcpTdx(_) => Self::GcpTdxMeasurement,
+            AttestationQuote::DstackNitroEnclave(_) => Self::NitroEnclavePcrs,
+            AttestationQuote::DstackAwsNitroTpm(_) => Self::AwsNitroTpmPcrs,
+        }
+    }
+}
+
 pub struct CvmVerifier {
     pub image_cache_dir: String,
     pub download_url: String,
@@ -676,48 +707,39 @@ impl CvmVerifier {
         let mut vm_config = attestation
             .decode_vm_config(&vm_config)
             .context("Failed to decode VM config")?;
-        match &attestation.quote {
-            AttestationQuote::DstackGcpTdx(quote) => {
+        match OsImageVerificationStrategy::select(&attestation.quote, &vm_config) {
+            OsImageVerificationStrategy::GcpTdxMeasurement => {
+                let AttestationQuote::DstackGcpTdx(quote) = &attestation.quote else {
+                    unreachable!("strategy selector returned the wrong GCP TDX branch")
+                };
                 self.verify_os_image_hash_for_gcp_tdx(&vm_config, &quote.tpm_quote)?;
             }
-            AttestationQuote::DstackTdx(_) => {
-                // New images carry a self-contained measurement document even
-                // when the boot kept the legacy attestation selector. Prefer
-                // that signed-MR-bound material; retain image download only
-                // for old legacy images which do not provide it.
-                if vm_config.tdx_attestation_variant.is_lite()
-                    || vm_config.tdx_measurement.is_some()
-                {
-                    self.verify_os_image_hash_for_dstack_tdx_lite(
-                        &vm_config,
-                        attestation,
-                        debug,
-                        details,
-                    )
-                    .await?;
-                } else {
-                    self.verify_os_image_hash_for_dstack_tdx(
-                        &vm_config,
-                        attestation,
-                        debug,
-                        details,
-                    )
+            OsImageVerificationStrategy::TdxLiteMeasurement => {
+                self.verify_os_image_hash_for_dstack_tdx_lite(
+                    &vm_config,
+                    attestation,
+                    debug,
+                    details,
+                )
+                .await?;
+            }
+            OsImageVerificationStrategy::TdxFullDownload => {
+                self.verify_os_image_hash_for_dstack_tdx(&vm_config, attestation, debug, details)
                     .await?;
-                }
             }
-            AttestationQuote::DstackNitroEnclave(_) => {
+            OsImageVerificationStrategy::NitroEnclavePcrs => {
                 let DstackVerifiedReport::DstackNitroEnclave(report) = &attestation.report else {
                     bail!("internal error: nitro quote without a verified nitro report");
                 };
                 self.verify_os_image_hash_for_nitro_enclave(&vm_config, &report.pcrs)?;
             }
-            AttestationQuote::DstackAwsNitroTpm(_) => {
+            OsImageVerificationStrategy::AwsNitroTpmPcrs => {
                 let DstackVerifiedReport::DstackAwsNitroTpm(report) = &attestation.report else {
                     bail!("internal error: NitroTPM quote without a verified NitroTPM report");
                 };
                 self.verify_os_image_hash_for_aws_nitro_tpm(&vm_config, &report.pcrs)?;
             }
-            AttestationQuote::DstackAmdSevSnp(_) => {
+            OsImageVerificationStrategy::SevSnpMeasurement => {
                 self.verify_os_image_hash_for_dstack_sev(
                     attestation,
                     &raw_config,
@@ -1317,6 +1339,179 @@ mod tests {
         Arc::new(AttestationVerifier::new_prod(None).unwrap())
     }
 
+    #[test]
+    fn platform_image_strategy_decision_table_is_exhaustive() {
+        let tdx = || {
+            AttestationQuote::DstackTdx(TdxQuote {
+                quote: Vec::new(),
+                event_log: Vec::new(),
+            })
+        };
+        let legacy: VmConfig = serde_json::from_value(serde_json::json!({})).unwrap();
+        let lite: VmConfig = serde_json::from_value(serde_json::json!({
+            "tdx_attestation_variant": "lite"
+        }))
+        .unwrap();
+        let gcp = AttestationQuote::DstackGcpTdx(DstackGcpTdxQuote {
+            tdx_quote: TdxQuote {
+                quote: Vec::new(),
+                event_log: Vec::new(),
+            },
+            tpm_quote: TpmQuote {
+                message: Vec::new(),
+                signature: Vec::new(),
+                pcr_values: Vec::new(),
+                ak_cert: Vec::new(),
+                platform: dstack_types::Platform::Gcp,
+                event_log: Vec::new(),
+            },
+        });
+        let rows = [
+            (
+                tdx(),
+                legacy.clone(),
+                OsImageVerificationStrategy::TdxFullDownload,
+            ),
+            (tdx(), lite, OsImageVerificationStrategy::TdxLiteMeasurement),
+            (
+                AttestationQuote::DstackAmdSevSnp(SnpQuote {
+                    report: Vec::new(),
+                    cert_chain: Vec::new(),
+                    mr_config: String::new(),
+                }),
+                legacy.clone(),
+                OsImageVerificationStrategy::SevSnpMeasurement,
+            ),
+            (
+                gcp,
+                legacy.clone(),
+                OsImageVerificationStrategy::GcpTdxMeasurement,
+            ),
+            (
+                AttestationQuote::DstackNitroEnclave(DstackNitroQuote {
+                    nsm_quote: Vec::new(),
+                }),
+                legacy.clone(),
+                OsImageVerificationStrategy::NitroEnclavePcrs,
+            ),
+            (
+                AttestationQuote::DstackAwsNitroTpm(DstackAwsNitroTpmQuote {
+                    attestation_doc: Vec::new(),
+                }),
+                legacy,
+                OsImageVerificationStrategy::AwsNitroTpmPcrs,
+            ),
+        ];
+        for (quote, vm_config, expected) in rows {
+            assert_eq!(
+                OsImageVerificationStrategy::select(&quote, &vm_config),
+                expected
+            );
+        }
+    }
+
+    #[test]
+    fn gcp_and_nitro_enclave_measurement_bindings_matrix() {
+        let verifier = test_verifier();
+
+        let nitro_pcrs = NitroPcrs {
+            pcr0: vec![0x10; 48],
+            pcr1: vec![0x11; 48],
+            pcr2: vec![0x12; 48],
+        };
+        let nitro_config: VmConfig = serde_json::from_value(serde_json::json!({
+            "os_image_hash": hex::encode(nitro_pcrs.image_hash()),
+        }))
+        .unwrap();
+        verifier
+            .verify_os_image_hash_for_nitro_enclave(&nitro_config, &nitro_pcrs)
+            .unwrap();
+        let mut changed_nitro = nitro_pcrs.clone();
+        changed_nitro.pcr2[0] ^= 1;
+        assert!(verifier
+            .verify_os_image_hash_for_nitro_enclave(&nitro_config, &changed_nitro)
+            .is_err());
+        let debug_nitro = NitroPcrs {
+            pcr0: vec![0; 48],
+            pcr1: vec![0; 48],
+            pcr2: vec![0; 48],
+        };
+        assert!(verifier
+            .verify_os_image_hash_for_nitro_enclave(&nitro_config, &debug_nitro)
+            .is_err());
+
+        let uki_hash = vec![0x24; 32];
+        let measurement = dstack_types::GcpOsImageMeasurement::new(uki_hash.clone()).unwrap();
+        let measurement_bytes = measurement.to_cbor_vec();
+        let checksum_file = format!(
+            "{}  measurement.gcp.cbor\n",
+            hex::encode(Sha256::digest(&measurement_bytes))
+        )
+        .into_bytes();
+        let os_image_hash = dstack_types::image_hash_from_sha256sum(&checksum_file);
+        let gcp_config: VmConfig = serde_json::from_value(serde_json::json!({
+            "os_image_hash": hex::encode(os_image_hash),
+            "gcp_measurement": dstack_types::GcpOsImageMeasurementDocument::new(
+                checksum_file,
+                measurement_bytes,
+            ),
+        }))
+        .unwrap();
+        let expected_pcr0 =
+            hex!("0cca9ec161b09288802e5a112255d21340ed5b797f5fe29cecccfd8f67b9f802");
+        let gcp_quote = |pcr0: Vec<u8>, event_28: Vec<u8>| TpmQuote {
+            message: Vec::new(),
+            signature: Vec::new(),
+            pcr_values: vec![tpm_types::PcrValue {
+                index: 0,
+                algorithm: "sha256".into(),
+                value: pcr0,
+            }],
+            ak_cert: Vec::new(),
+            platform: dstack_types::Platform::Gcp,
+            event_log: vec![
+                tpm_types::TpmEvent {
+                    pcr_index: 2,
+                    digest: vec![1; 32],
+                },
+                tpm_types::TpmEvent {
+                    pcr_index: 2,
+                    digest: vec![2; 32],
+                },
+                tpm_types::TpmEvent {
+                    pcr_index: 2,
+                    digest: event_28,
+                },
+            ],
+        };
+        verifier
+            .verify_os_image_hash_for_gcp_tdx(
+                &gcp_config,
+                &gcp_quote(expected_pcr0.to_vec(), uki_hash.clone()),
+            )
+            .unwrap();
+        assert!(verifier
+            .verify_os_image_hash_for_gcp_tdx(
+                &gcp_config,
+                &gcp_quote(vec![0; 32], uki_hash.clone()),
+            )
+            .is_err());
+        assert!(verifier
+            .verify_os_image_hash_for_gcp_tdx(
+                &gcp_config,
+                &gcp_quote(expected_pcr0.to_vec(), vec![0; 32]),
+            )
+            .is_err());
+        let mut missing_document = gcp_config;
+        missing_document.gcp_measurement = None;
+        assert!(verifier
+            .verify_os_image_hash_for_gcp_tdx(
+                &missing_document,
+                &gcp_quote(expected_pcr0.to_vec(), uki_hash),
+            )
+            .is_err());
+    }
+
     #[test]
     fn aws_os_image_check_requires_measurement() {
         let pcrs = aws_boot_pcrs(0x04);