fix(tpm-qvl): 20 KB of duplicate PCR selections buys 4.7 seconds of CPU before any signature is checked

dstack/tpm-qvl/src/verify.rs

@@ -10,6 +10,7 @@ use dstack_types::Platform;
 use p256::ecdsa::{signature::hazmat::PrehashVerifier, Signature, VerifyingKey};
 use rsa::RsaPublicKey;
 use sha2::{Digest, Sha256};
+use std::collections::{HashMap, HashSet};
 use tracing::{debug, warn};
 use x509_parser::prelude::*;
 
@@ -96,6 +97,23 @@ pub fn verify_quote_with_ca(
         });
     }
 
+    // A TPM selects each PCR at most once per quote, so a repeated index is an
+    // impossible shape. Reject it here rather than let it through: both gates
+    // around this one are computable by whoever wrote the quote -- the
+    // selection list above is still unsigned, and pcr_digest below is a plain
+    // SHA-256 over the supplied values -- and the event-log replay that follows
+    // costs one pass over the log per entry in pcr_values. All of it runs
+    // before the AK signature is checked.
+    let mut seen_pcr_indices = HashSet::with_capacity(quote.pcr_values.len());
+    for pcr in &quote.pcr_values {
+        if !seen_pcr_indices.insert(pcr.index) {
+            return Err(VerificationError {
+                status,
+                error: anyhow!("duplicate PCR index {} in quote", pcr.index),
+            });
+        }
+    }
+
     // compute_pcr_digest() and the event-log replay below assume the SHA-256 PCR
     // bank. Reject other banks explicitly instead of silently failing with a
     // confusing "PCR digest mismatch".
@@ -315,31 +333,43 @@ fn compute_pcr_digest(pcr_values: &[PcrValue]) -> Result<Vec<u8>> {
 }
 
 fn verify_event_log(pcr_values: &[PcrValue], event_log: &[TpmEvent]) -> Result<()> {
-    for pcr in pcr_values {
-        let pcr_events: Vec<&TpmEvent> = event_log
-            .iter()
-            .filter(|e| e.pcr_index == pcr.index)
-            .collect();
+    // Bucket the log once instead of re-scanning it per quoted PCR. A real GCP
+    // CVM ships the whole firmware log (PCRs 0-11) while dstack quotes only a
+    // few of them, so entries for an unquoted PCR are normal and are dropped
+    // here rather than replayed.
+    let quoted: HashSet<u32> = pcr_values.iter().map(|p| p.index).collect();
+    let mut events_by_pcr: HashMap<u32, Vec<&TpmEvent>> = HashMap::new();
+    for event in event_log {
+        if quoted.contains(&event.pcr_index) {
+            events_by_pcr
+                .entry(event.pcr_index)
+                .or_default()
+                .push(event);
+        }
+    }
 
-        if pcr_events.is_empty() {
+    for pcr in pcr_values {
+        // Taking the bucket keeps the total replay bounded by the log length
+        // even if a duplicate index ever reaches this far.
+        let Some(pcr_events) = events_by_pcr.remove(&pcr.index) else {
             continue;
-        }
+        };
 
         // Replay PCR extension to verify Event Log matches quote
-        let mut replayed_pcr = vec![0u8; 32];
+        let mut replayed_pcr = [0u8; 32];
         for event in &pcr_events {
             let mut hasher = Sha256::new();
-            hasher.update(&replayed_pcr);
+            hasher.update(replayed_pcr);
             hasher.update(&event.digest);
-            replayed_pcr = hasher.finalize().to_vec();
+            replayed_pcr = hasher.finalize().into();
         }
 
-        if replayed_pcr != pcr.value {
+        if replayed_pcr[..] != pcr.value[..] {
             bail!(
                 "PCR {} replay mismatch: expected {}, got {}",
                 pcr.index,
                 hex::encode(&pcr.value),
-                hex::encode(&replayed_pcr)
+                hex::encode(replayed_pcr)
             );
         }
 
@@ -687,3 +717,172 @@ fn verify_ak_chain_with_collateral(
         }
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::GCP_ROOT_CA;
+    use std::sync::mpsc;
+    use std::time::{Duration, Instant};
+
+    /// The event-log replay runs before the AK certificate is parsed and before
+    /// the quote signature is checked, so everything below is work an
+    /// unauthenticated caller can ask for. A quote that gets this far only has
+    /// to satisfy two gates, and a test can synthesize both: the PCR selection
+    /// list lives in the still-unsigned TPMS_ATTEST, and `pcr_digest` is a
+    /// plain SHA-256 over the supplied PCR values.
+    fn attest_message(pcr_index: u8, selections: usize, pcr_digest: &[u8]) -> Vec<u8> {
+        let mut msg = Vec::new();
+        msg.extend_from_slice(&0xff544347u32.to_be_bytes()); // magic
+        msg.extend_from_slice(&0x8018u16.to_be_bytes()); // TPM_ST_ATTEST_QUOTE
+        msg.extend_from_slice(&0u16.to_be_bytes()); // qualified_signer
+        msg.extend_from_slice(&0u16.to_be_bytes()); // qualified_data
+        msg.extend_from_slice(&0u64.to_be_bytes()); // clock
+        msg.extend_from_slice(&0u32.to_be_bytes()); // reset_count
+        msg.extend_from_slice(&0u32.to_be_bytes()); // restart_count
+        msg.push(1); // safe
+        msg.extend_from_slice(&0u64.to_be_bytes()); // firmware_version
+        msg.extend_from_slice(&(selections as u32).to_be_bytes());
+        for _ in 0..selections {
+            msg.extend_from_slice(&0x000bu16.to_be_bytes()); // TPM_ALG_SHA256
+            msg.push(1); // sizeof_select
+            msg.push(1 << pcr_index); // bitmap
+        }
+        msg.extend_from_slice(&(pcr_digest.len() as u16).to_be_bytes());
+        msg.extend_from_slice(pcr_digest);
+        msg
+    }
+
+    fn replayed(digests: &[Vec<u8>]) -> Vec<u8> {
+        let mut value = vec![0u8; 32];
+        for digest in digests {
+            let mut hasher = Sha256::new();
+            hasher.update(&value);
+            hasher.update(digest);
+            value = hasher.finalize().to_vec();
+        }
+        value
+    }
+
+    fn pcr_digest_of(pcr_values: &[PcrValue]) -> Vec<u8> {
+        let mut hasher = Sha256::new();
+        for pcr in pcr_values {
+            hasher.update(&pcr.value);
+        }
+        hasher.finalize().to_vec()
+    }
+
+    /// A quote naming PCR 0 `n` times with a matching `n`-entry event log. Both
+    /// gates in front of the replay pass; only the replay itself is quadratic.
+    fn duplicate_pcr_quote(n: usize) -> TpmQuote {
+        let digests: Vec<Vec<u8>> = (0..n).map(|i| vec![i as u8; 32]).collect();
+        let value = replayed(&digests);
+        let pcr_values: Vec<PcrValue> = (0..n)
+            .map(|_| PcrValue {
+                index: 0,
+                algorithm: "sha256".to_string(),
+                value: value.clone(),
+            })
+            .collect();
+        let event_log: Vec<TpmEvent> = digests
+            .into_iter()
+            .map(|digest| TpmEvent {
+                pcr_index: 0,
+                digest,
+            })
+            .collect();
+        TpmQuote {
+            message: attest_message(0, n, &pcr_digest_of(&pcr_values)),
+            signature: Vec::new(),
+            pcr_values,
+            ak_cert: Vec::new(),
+            platform: Platform::Gcp,
+            event_log,
+        }
+    }
+
+    /// Nothing in the TPM data model lets a quote name the same PCR twice, but
+    /// the replay used to accept it and re-walk the whole event log once per
+    /// entry: `n` duplicate PCR values plus `n` events cost `n^2` SHA-256
+    /// extends, all of it before the AK signature is checked. ~20 KB of
+    /// TPMS_ATTEST and a few hundred KB of PCR values and events buy minutes of
+    /// a core.
+    #[test]
+    fn duplicate_pcr_indices_do_not_make_the_replay_quadratic() {
+        const N: usize = 5_000;
+        const DEADLINE: Duration = Duration::from_secs(5);
+
+        let quote = duplicate_pcr_quote(N);
+        let collateral = QuoteCollateral {
+            cert_chain_pem: String::new(),
+            crls: Vec::new(),
+            root_ca_crl: None,
+        };
+
+        // Run on a worker so a regression fails the test instead of hanging CI.
+        let (tx, rx) = mpsc::channel();
+        std::thread::spawn(move || {
+            let started = Instant::now();
+            let err = verify_quote_with_ca(&quote, &collateral, GCP_ROOT_CA)
+                .err()
+                .map(|e| format!("{:#}", e.error));
+            let _ = tx.send((err, started.elapsed()));
+        });
+
+        let (err, elapsed) = rx
+            .recv_timeout(DEADLINE)
+            .unwrap_or_else(|_| panic!("{N} duplicate PCR entries still busy after {DEADLINE:?}"));
+        let err = err.expect("a quote naming PCR 0 5000 times must be rejected");
+        assert!(
+            err.contains("duplicate PCR index"),
+            "expected a duplicate-index rejection, got: {err}"
+        );
+        assert!(
+            elapsed < DEADLINE,
+            "rejection took {elapsed:?}, expected well under {DEADLINE:?}"
+        );
+    }
+
+    /// The real PCR-2 digests a GCP CVM reported, captured in
+    /// `cc-eventlog/samples/tpm_eventlog.bin` (see its decode test): separator,
+    /// GPT event, UKI, kernel.
+    const GCP_PCR2_DIGESTS: [&str; 4] = [
+        "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119",
+        "00b8a357e652623798d1bbd16c375ec90fbed802b4269affa3e78e6eb19386cf",
+        "9ab14a46f858662a89adc102d2a57a13f52f75c1769d65a4c34edbbfc8855f0f",
+        "ade943a0a7a3189a3201ba17d7df778eb380cbd33ce5e361176e974ccf7cdedb",
+    ];
+
+    /// A device ships its whole firmware event log -- the captured GCP one
+    /// spans PCRs 0-11 -- while dstack quotes only a few PCRs, so entries for
+    /// an unquoted PCR are ordinary and must not make the replay fail.
+    #[test]
+    fn events_for_pcrs_outside_the_quote_are_ignored() {
+        let pcr2_digests: Vec<Vec<u8>> = GCP_PCR2_DIGESTS
+            .iter()
+            .map(|d| hex::decode(d).expect("fixture digest is hex"))
+            .collect();
+
+        let mut event_log = Vec::new();
+        for digest in &pcr2_digests {
+            // Interleave the unquoted PCRs the real log also carries.
+            for pcr_index in [0, 1, 3, 4, 5, 6, 7, 9, 11] {
+                event_log.push(TpmEvent {
+                    pcr_index,
+                    digest: vec![pcr_index as u8; 32],
+                });
+            }
+            event_log.push(TpmEvent {
+                pcr_index: 2,
+                digest: digest.clone(),
+            });
+        }
+
+        let pcr_values = vec![PcrValue {
+            index: 2,
+            algorithm: "sha256".to_string(),
+            value: replayed(&pcr2_digests),
+        }];
+        verify_event_log(&pcr_values, &event_log).expect("a full firmware log still replays");
+    }
+}