shade-attestation 0.1.1

use crate::{
    // app_compose::AppCompose,
    collateral::Collateral,
    measurements::{FullMeasurements, MeasurementsError},
    quote::QuoteBytes,
    report_data::ReportData,
    tcb_info::{EventLog, HexBytes, TcbInfo},
};

use alloc::{
    format,
    string::{String, ToString},
    vec::Vec,
};
use borsh::{BorshDeserialize, BorshSerialize};
use core::fmt;
use dcap_qvl::verify::VerifiedReport;
use derive_more::Constructor;
use serde::{Deserialize, Serialize};
use serde_json::json;
use sha2::{Digest as _, Sha256, Sha384};

/// Expected TCB status for a successfully verified TEE quote.
const EXPECTED_QUOTE_STATUS: &str = "UpToDate";

// DSTACK_EVENT_TYPE is defined in https://github.com/Dstack-TEE/dstack/blob/cfa4cc4e8a4f525d537883b1a0ba5d9fbfd87f1e/tdx-attest/src/lib.rs#L28
// It is the same for all events
const DSTACK_EVENT_TYPE: u32 = 134217729;

const COMPOSE_HASH_EVENT: &str = "compose-hash";
pub(crate) const KEY_PROVIDER_EVENT: &str = "key-provider";

const RTMR3_INDEX: u32 = 3;

#[derive(Clone, Constructor, Serialize, Deserialize, BorshDeserialize, BorshSerialize)]
pub struct DstackAttestation {
    pub quote: QuoteBytes,
    pub collateral: Collateral,
    pub tcb_info: TcbInfo,
}

#[derive(Clone, Debug, PartialEq, Eq, thiserror::Error)]
pub enum VerificationError {
    #[error("could not parse embedded measurements: {0}")]
    EmbeddedMeasurementsParsing(MeasurementsError),
    #[error("dcap verification failed: {0}")]
    DcapVerification(String),
    #[error("verification report is not TD10")]
    ReportNotTd10,
    #[error("TCB status `{0}` is not up to date")]
    TcbStatusNotUpToDate(String),
    #[error("ouststanding advisories reported: {0}")]
    NonEmptyAdvisoryIds(String),
    #[error("wrong {name} hash (found {found} expected {expected})")]
    WrongHash {
        name: &'static str,
        found: String,
        expected: String,
    },
    #[error("invalid event type {0}")]
    InvalidEventType(u32),
    #[error("failed to decode event digest `{0}`")]
    EventDecoding(String),
    #[error("failed to parse app compose JSON: {0}")]
    AppComposeParsing(String),
    #[error("no {0} event in event log")]
    MissingEvent(&'static str),
    #[error("duplicate {0} events in event log")]
    DuplicateEvent(&'static str),
    #[error("invalid app compose config: `{0}`")]
    InvalidAppComposeConfig(String),
    #[error("app-compose event payload had an unexpected size of {0}")]
    AppComposeEventPayloadWrongSize(usize),
    #[error("app-compose event payload `{0}` is not a hex string")]
    AppComposeEventPayloadNotHex(String),
    #[error(
        "the attestation certificate with timestap {attestation_time} has expired since {expiry_time}"
    )]
    ExpiredCertificate {
        attestation_time: u64,
        expiry_time: u64,
    },
    #[error("PPID must be 16 bytes, got {0}")]
    PpidWrongSize(usize),
    #[error("the mock attestation is invalid per definition")]
    InvalidMockAttestation,
    #[error("custom error: `{0}`")]
    Custom(String),
}

impl fmt::Debug for DstackAttestation {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        const MAX_BYTES: usize = 2048;

        fn truncate_debug<T: fmt::Debug>(value: &T, max_bytes: usize) -> String {
            let debug_str = format!("{:?}", value);
            if debug_str.len() <= max_bytes {
                debug_str
            } else {
                format!(
                    "{}... (truncated {} bytes)",
                    &debug_str[..max_bytes],
                    debug_str.len() - max_bytes
                )
            }
        }

        f.debug_struct("DstackAttestation")
            .field("quote", &truncate_debug(&self.quote, MAX_BYTES))
            .field("collateral", &truncate_debug(&self.collateral, MAX_BYTES))
            .field("tcb_info", &truncate_debug(&self.tcb_info, MAX_BYTES))
            .finish()
    }
}

/// Produces a valid dummy attestation for tests
pub fn create_mock_dstack_attestation() -> DstackAttestation {
    let collateral_json = json!({
        "tcb_info_issuer_chain": "",
        "tcb_info": "",
        "tcb_info_signature": "",
        "qe_identity_issuer_chain": "",
        "qe_identity": "",
        "qe_identity_signature": "",
        "pck_crl_issuer_chain": "",
        "root_ca_crl": "",
        "pck_crl": ""
    });
    let collateral = Collateral::try_from_json(collateral_json).expect("mock collateral is valid");
    let tcb_info = TcbInfo {
        mrtd: HexBytes::default(),
        rtmr0: HexBytes::default(),
        rtmr1: HexBytes::default(),
        rtmr2: HexBytes::default(),
        rtmr3: HexBytes::default(),
        os_image_hash: None,
        compose_hash: HexBytes::default(),
        device_id: HexBytes::default(),
        app_compose: String::new(),
        event_log: Vec::new(),
    };
    DstackAttestation::new(QuoteBytes::from(Vec::new()), collateral, tcb_info)
}

impl DstackAttestation {
    /// Checks whether this attestation is valid
    /// with respect to expected values of:
    /// - report_data: must be measured correctly in RTMR3
    /// - timestamp_seconds: current UNIX time in seconds
    /// - accepted_measurements: set of accepted RTMRs and key-provider event digest.
    ///   If any element in the set is valid, the function accepts the attestation as
    ///   valid.
    /// - accepted_ppids: set of accepted PPIDs. PPID in the attestation must match one of the allowed PPIDs.
    ///
    /// Returns the `FullMeasurements` that matched and the verified PPID if verification succeeds.
    pub fn verify(
        &self,
        expected_report_data: ReportData,
        timestamp_seconds: u64,
        accepted_measurements: &[FullMeasurements],
        accepted_ppids: &[HexBytes<16>],
    ) -> Result<(FullMeasurements, HexBytes<16>), VerificationError> {
        let verification_result =
            dcap_qvl::verify::verify(&self.quote, &self.collateral, timestamp_seconds)
                .map_err(|e| VerificationError::DcapVerification(e.to_string()))?;

        let report_data = verification_result
            .report
            .as_td10()
            .ok_or(VerificationError::ReportNotTd10)?;

        // Verify all attestation components
        self.verify_tcb_status(&verification_result)?;
        self.verify_report_data(&expected_report_data, report_data)?;
        let ppid = self.verify_ppid(verification_result.ppid, accepted_ppids)?;

        self.verify_rtmr3(report_data, &self.tcb_info)?;
        self.verify_app_compose(&self.tcb_info)?;

        let measurements =
            self.verify_any_measurements(report_data, &self.tcb_info, accepted_measurements)?;
        Ok((measurements, ppid))
    }

    /// Replays RTMR3 from the event log by hashing all relevant events together and verifies all
    /// digests are correct
    fn verify_event_log_rtmr3(
        event_log: &[EventLog],
        expected_digest: [u8; 48],
    ) -> Result<(), VerificationError> {
        let mut digest = [0u8; 48];

        let filtered_events = event_log.iter().filter(|e| e.imr == RTMR3_INDEX);

        for event in filtered_events {
            // In Dstack, all events measured in RTMR3 are of type DSTACK_EVENT_TYPE
            if event.event_type != DSTACK_EVENT_TYPE {
                return Err(VerificationError::InvalidEventType(event.event_type));
            }
            let mut hasher = Sha384::new();
            hasher.update(digest);
            let payload_bytes = match hex::decode(&event.event_payload) {
                Ok(bytes) => bytes,
                Err(e) => {
                    tracing::error!("Failed to decode hex string for: {:?}", e);
                    return Err(VerificationError::EventDecoding(hex::encode(*event.digest)));
                }
            };
            let expected_digest =
                Self::event_digest(event.event_type, &event.event, &payload_bytes);
            compare_hashes("event_digest", event.digest.as_slice(), &expected_digest)?;

            hasher.update(event.digest.as_slice());

            digest = hasher.finalize().into();
        }

        compare_hashes("event_log", &digest, &expected_digest)
    }

    fn validate_app_compose_payload(
        expected_event_payload_hex: &str,
        app_compose: &str,
    ) -> Result<(), VerificationError> {
        let expected_payload = match hex::decode(expected_event_payload_hex) {
            Ok(bytes) => match <[u8; 32]>::try_from(bytes.as_slice()) {
                Ok(expected_bytes) => expected_bytes,
                Err(_) => {
                    return Err(VerificationError::AppComposeEventPayloadWrongSize(
                        bytes.len(),
                    ));
                }
            },
            Err(e) => {
                tracing::error!(
                    "Failed to decode hex string for compose-hash event: {:?}",
                    e
                );
                return Err(VerificationError::AppComposeEventPayloadNotHex(
                    expected_event_payload_hex.to_string(),
                ));
            }
        };

        let app_compose_hash: [u8; 32] = Sha256::digest(app_compose.as_bytes()).into();

        compare_hashes("app_compose_payload", &app_compose_hash, &expected_payload)
    }

    /// Verifies TCB status and security advisories.
    fn verify_tcb_status(
        &self,
        verification_result: &VerifiedReport,
    ) -> Result<(), VerificationError> {
        // The "UpToDate" TCB status indicates that the measured platform components (CPU
        // microcode, firmware, etc.) match the latest known good values published by Intel
        // and do not require any updates or mitigations.
        let status_is_up_to_date = verification_result.status == EXPECTED_QUOTE_STATUS;

        // Advisory IDs indicate known security vulnerabilities or issues with the TEE.
        // For a quote to be considered secure, there should be no outstanding advisories.
        let no_security_advisories = verification_result.advisory_ids.is_empty();

        status_is_up_to_date.or_err(|| {
            VerificationError::TcbStatusNotUpToDate(verification_result.status.clone())
        })?;

        no_security_advisories.or_err(|| {
            VerificationError::NonEmptyAdvisoryIds(verification_result.advisory_ids.join(", "))
        })?;

        Ok(())
    }

    /// Verifies report data matches expected values.
    fn verify_report_data(
        &self,
        expected: &ReportData,
        actual: &dcap_qvl::quote::TDReport10,
    ) -> Result<(), VerificationError> {
        // Check the report data from the report matches the expected report data.
        compare_hashes("report_data", &actual.report_data, &expected.to_bytes())
    }

    /// Verifies PPID is in the allowed PPIDs list. Returns the matched PPID on success.
    fn verify_ppid(
        &self,
        ppid: Vec<u8>,
        accepted_ppids: &[HexBytes<16>],
    ) -> Result<HexBytes<16>, VerificationError> {
        // In the future we'll change this to checking device_id inside of PPID

        let ppid_array = match <[u8; 16]>::try_from(ppid.as_slice()) {
            Ok(array) => array,
            Err(_) => {
                return Err(VerificationError::PpidWrongSize(ppid.len()));
            }
        };
        let ppid_hex_bytes = HexBytes::from(ppid_array);
        if !accepted_ppids.contains(&ppid_hex_bytes) {
            return Err(VerificationError::Custom(format!(
                "PPID {} is not in the allowed PPIDs list",
                hex::encode(ppid_hex_bytes.as_ref())
            )));
        }
        Ok(ppid_hex_bytes)
    }

    /// Try to verify static RTMRs, key_provider_digest and app_compose_hash against multiple expected measurement sets.
    /// Returns the matching `ExpectedMeasurements` if any set matches; otherwise, returns a WrongHash error.
    fn verify_any_measurements(
        &self,
        report_data: &dcap_qvl::quote::TDReport10,
        tcb_info: &TcbInfo,
        accepted_measurements: &[FullMeasurements],
    ) -> Result<FullMeasurements, VerificationError> {
        for expected in accepted_measurements {
            if self
                .verify_static_rtmrs(report_data, tcb_info, expected)
                .is_ok()
                && self
                    .verify_key_provider_digest(tcb_info, &expected.key_provider_event_digest)
                    .is_ok()
                && self
                    .verify_app_compose_hash(tcb_info, &expected.app_compose_hash_payload)
                    .is_ok()
            {
                return Ok(*expected); // found a valid match
            }
        }

        Err(VerificationError::WrongHash {
            name: "expected_measurements",
            expected: "one of the embedded TCB info sets (prod or dev)".into(),
            found: "none matched".into(),
        })
    }
    /// Verifies static RTMRs match expected values.
    fn verify_static_rtmrs(
        &self,
        report_data: &dcap_qvl::quote::TDReport10,
        tcb_info: &TcbInfo,
        expected_measurements: &FullMeasurements,
    ) -> Result<(), VerificationError> {
        // Check if the RTMRs match the expected values. To learn more about RTMRs and
        // their significance, refer to the TDX documentation:
        // - https://phala.network/posts/understanding-tdx-attestation-reports-a-developers-guide
        // - https://www.kernel.org/doc/Documentation/x86/tdx.rst
        compare_hashes(
            "rtmr0_report_data",
            &report_data.rt_mr0,
            &expected_measurements.rtmrs.rtmr0,
        )?;
        compare_hashes(
            "rtmr1_report_data",
            &report_data.rt_mr1,
            &expected_measurements.rtmrs.rtmr1,
        )?;
        compare_hashes(
            "rtmr2_report_data",
            &report_data.rt_mr2,
            &expected_measurements.rtmrs.rtmr2,
        )?;
        compare_hashes(
            "mrtd_report_data",
            &report_data.mr_td,
            &expected_measurements.rtmrs.mrtd,
        )?;

        compare_hashes(
            "rtmr0_tcb_info",
            tcb_info.rtmr0.as_slice(),
            &expected_measurements.rtmrs.rtmr0,
        )?;
        compare_hashes(
            "rtmr1_tcb_info",
            tcb_info.rtmr1.as_slice(),
            &expected_measurements.rtmrs.rtmr1,
        )?;
        compare_hashes(
            "rtmr2_tcb_info",
            tcb_info.rtmr2.as_slice(),
            &expected_measurements.rtmrs.rtmr2,
        )?;
        compare_hashes(
            "mtrd_tcb_info",
            tcb_info.mrtd.as_slice(),
            &expected_measurements.rtmrs.mrtd,
        )
    }

    /// Verifies RTMR3 by replaying event log.
    fn verify_rtmr3(
        &self,
        report_data: &dcap_qvl::quote::TDReport10,
        tcb_info: &TcbInfo,
    ) -> Result<(), VerificationError> {
        compare_hashes("rtmr3", tcb_info.rtmr3.as_slice(), &report_data.rt_mr3)?;

        Self::verify_event_log_rtmr3(&tcb_info.event_log, report_data.rt_mr3)
    }

    /// Verifies the app compose hash from RTMR3 event matches the one in TCB info.
    /// and that the app compose hashed in the tcb info matches the hashes provided
    fn verify_app_compose(&self, tcb_info: &TcbInfo) -> Result<(), VerificationError> {
        // Allow any app compose configuration as long as the hash matches (in verify_app_compose_hash)
        // let app_compose: AppCompose = serde_json::from_str(&tcb_info.app_compose)
        //     .map_err(|e| VerificationError::AppComposeParsing(e.to_string()))?;

        // Self::validate_app_compose_config(&app_compose).or_err(|| {
        //     VerificationError::InvalidAppComposeConfig(tcb_info.app_compose.to_string())
        // })?;

        let app_compose_event = tcb_info.get_single_event(COMPOSE_HASH_EVENT)?;

        compare_hex_hashes(
            "app_compose_event_hash",
            &app_compose_event.event_payload,
            &hex::encode(*tcb_info.compose_hash),
        )?;

        Self::validate_app_compose_payload(&app_compose_event.event_payload, &tcb_info.app_compose)
    }

    // /// Validates app compose configuration against expected security requirements.
    // fn validate_app_compose_config(app_compose: &AppCompose) -> bool {
    //     app_compose.manifest_version == 2
    //         && app_compose.runner == "docker-compose"
    //         && !app_compose.kms_enabled
    //         && app_compose.gateway_enabled == Some(false)
    //         && app_compose.public_logs
    //         && app_compose.public_sysinfo
    //         && app_compose.local_key_provider_enabled
    //         && app_compose.allowed_envs.is_empty()
    //         && app_compose.no_instance_id
    //         && app_compose.pre_launch_script.is_none()
    // }

    /// Verifies local key-provider event digest matches the expected digest.
    fn verify_key_provider_digest(
        &self,
        tcb_info: &TcbInfo,
        expected_digest: &[u8; 48],
    ) -> Result<(), VerificationError> {
        let key_provider_event = tcb_info.get_single_event(KEY_PROVIDER_EVENT)?;

        compare_hashes(
            "key_provider",
            key_provider_event.digest.as_slice(),
            expected_digest,
        )
    }

    /// Verifies app compose hash in tcb info matches the expected hash in the FullMeasurements.
    fn verify_app_compose_hash(
        &self,
        tcb_info: &TcbInfo,
        expected_hash_payload: &[u8; 32],
    ) -> Result<(), VerificationError> {
        compare_hashes(
            "app_compose_hash",
            tcb_info.compose_hash.as_slice(),
            expected_hash_payload,
        )
    }

    // Implementation taken to match Dstack's https://github.com/Dstack-TEE/dstack/blob/cfa4cc4e8a4f525d537883b1a0ba5d9fbfd87f1e/cc-eventlog/src/lib.rs#L54
    fn event_digest(event_type: u32, event: &str, payload: &[u8]) -> [u8; 48] {
        let mut hasher = Sha384::new();
        hasher.update(event_type.to_ne_bytes());
        hasher.update(b":");
        hasher.update(event.as_bytes());
        hasher.update(b":");
        hasher.update(payload);
        hasher.finalize().into()
    }
}

fn compare_hashes(
    name: &'static str,
    found: &[u8],
    expected: &[u8],
) -> Result<(), VerificationError> {
    (found == expected).or_err(|| VerificationError::WrongHash {
        name,
        found: hex::encode(found),
        expected: hex::encode(expected),
    })
}

fn compare_hex_hashes<S: ToString + Eq>(
    name: &'static str,
    found: S,
    expected: S,
) -> Result<(), VerificationError> {
    (found == expected).or_err(|| VerificationError::WrongHash {
        name,
        found: found.to_string(),
        expected: expected.to_string(),
    })
}

pub trait OrErr {
    fn or_err<Error>(self, err: impl FnOnce() -> Error) -> Result<(), Error>;
}

impl OrErr for bool {
    fn or_err<Error>(self, err: impl FnOnce() -> Error) -> Result<(), Error> {
        self.then_some(()).ok_or_else(err)
    }
}

pub trait GetSingleEvent {
    fn get_single_event(&self, event_name: &'static str) -> Result<&EventLog, VerificationError>;
}

impl GetSingleEvent for TcbInfo {
    fn get_single_event(&self, event_name: &'static str) -> Result<&EventLog, VerificationError> {
        let mut events = self
            .event_log
            .iter()
            .filter(|event| event.event == event_name && event.imr == RTMR3_INDEX);

        let Some(event) = events.next() else {
            return Err(VerificationError::MissingEvent(event_name));
        };

        if events.next().is_some() {
            Err(VerificationError::DuplicateEvent(event_name))
        } else {
            Ok(event)
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::measurements::{FullMeasurements, create_mock_full_measurements_hex};
    use alloc::vec;
    use dcap_qvl::quote::{Report, TDReport10};
    use dcap_qvl::tcb_info::{TcbStatus, TcbStatusWithAdvisory};
    use dcap_qvl::verify::VerifiedReport;

    fn td_report() -> TDReport10 {
        TDReport10 {
            tee_tcb_svn: [0; 16],
            mr_seam: [0; 48],
            mr_signer_seam: [0; 48],
            seam_attributes: [0; 8],
            td_attributes: [0; 8],
            xfam: [0; 8],
            mr_td: [0; 48],
            mr_config_id: [0; 48],
            mr_owner: [0; 48],
            mr_owner_config: [0; 48],
            rt_mr0: [0; 48],
            rt_mr1: [0; 48],
            rt_mr2: [0; 48],
            rt_mr3: [0; 48],
            report_data: [0; 64],
        }
    }

    fn verified_report(status: &str, advisory_ids: Vec<String>) -> VerifiedReport {
        VerifiedReport {
            status: status.to_string(),
            advisory_ids,
            report: Report::TD10(td_report()),
            ppid: Vec::new(),
            qe_status: TcbStatusWithAdvisory::new(TcbStatus::UpToDate, Vec::new()),
            platform_status: TcbStatusWithAdvisory::new(TcbStatus::UpToDate, Vec::new()),
        }
    }

    fn empty_tcb_info() -> TcbInfo {
        TcbInfo {
            mrtd: HexBytes::from([0u8; 48]),
            rtmr0: HexBytes::from([0u8; 48]),
            rtmr1: HexBytes::from([0u8; 48]),
            rtmr2: HexBytes::from([0u8; 48]),
            rtmr3: HexBytes::from([0u8; 48]),
            os_image_hash: None,
            compose_hash: HexBytes::from([0u8; 32]),
            device_id: HexBytes::from([0u8; 32]),
            app_compose: String::new(),
            event_log: Vec::new(),
        }
    }

    // -------- verify_tcb_status --------

    // "UpToDate" + no advisories passes.
    #[test]
    fn verify_tcb_status_accepts_up_to_date_no_advisories() {
        let attestation = create_mock_dstack_attestation();
        let report = verified_report("UpToDate", Vec::new());
        assert_eq!(attestation.verify_tcb_status(&report), Ok(()));
    }

    // Any non-"UpToDate" status fails.
    #[test]
    fn verify_tcb_status_rejects_out_of_date_status() {
        let attestation = create_mock_dstack_attestation();
        let report = verified_report("OutOfDate", Vec::new());
        assert_eq!(
            attestation.verify_tcb_status(&report),
            Err(VerificationError::TcbStatusNotUpToDate(
                "OutOfDate".to_string()
            ))
        );
    }

    // Outstanding advisories fail even when status is "UpToDate".
    #[test]
    fn verify_tcb_status_rejects_non_empty_advisories() {
        let attestation = create_mock_dstack_attestation();
        let report = verified_report("UpToDate", vec!["INTEL-SA-00001".to_string()]);
        assert_eq!(
            attestation.verify_tcb_status(&report),
            Err(VerificationError::NonEmptyAdvisoryIds(
                "INTEL-SA-00001".to_string()
            ))
        );
    }

    // -------- verify_report_data --------

    // 64-byte report_data equal to expected passes.
    #[test]
    fn verify_report_data_accepts_matching_64_bytes() {
        let attestation = create_mock_dstack_attestation();
        let expected = ReportData::from([7u8; 64]);
        let mut report = td_report();
        report.report_data = expected.to_bytes();
        assert_eq!(attestation.verify_report_data(&expected, &report), Ok(()));
    }

    // One byte different fails (byte-exact compare).
    #[test]
    fn verify_report_data_rejects_one_byte_difference() {
        let attestation = create_mock_dstack_attestation();
        let expected = ReportData::from([7u8; 64]);
        let mut report = td_report();
        report.report_data = expected.to_bytes();
        report.report_data[0] = 0xFF;
        assert!(matches!(
            attestation.verify_report_data(&expected, &report),
            Err(VerificationError::WrongHash {
                name: "report_data",
                ..
            })
        ));
    }

    // -------- verify_ppid --------

    // 16-byte PPID in the accepted list returns the matched PPID.
    #[test]
    fn verify_ppid_accepts_known_ppid() {
        let attestation = create_mock_dstack_attestation();
        let ppid_bytes = [1u8; 16];
        let accepted = vec![HexBytes::from(ppid_bytes)];
        assert_eq!(
            attestation.verify_ppid(ppid_bytes.to_vec(), &accepted),
            Ok(HexBytes::from(ppid_bytes))
        );
    }

    // Empty accepted list fails closed.
    #[test]
    fn verify_ppid_rejects_empty_accepted_list() {
        let attestation = create_mock_dstack_attestation();
        let accepted: Vec<HexBytes<16>> = Vec::new();
        assert!(matches!(
            attestation.verify_ppid([1u8; 16].to_vec(), &accepted),
            Err(VerificationError::Custom(_))
        ));
    }

    // Valid-shape but unknown PPID is rejected with the documented message.
    #[test]
    fn verify_ppid_rejects_unknown_ppid() {
        let attestation = create_mock_dstack_attestation();
        let accepted = vec![HexBytes::from([1u8; 16])];
        let err = attestation
            .verify_ppid([9u8; 16].to_vec(), &accepted)
            .unwrap_err();
        match err {
            VerificationError::Custom(msg) => {
                assert!(msg.contains("not in the allowed PPIDs list"));
            }
            other => panic!("expected Custom, got {other:?}"),
        }
    }

    // 32-byte PPID is rejected and the error message says "16 bytes".
    #[test]
    fn verify_ppid_wrong_size_32_bytes() {
        let attestation = create_mock_dstack_attestation();
        let accepted = vec![HexBytes::from([0u8; 16])];
        let err = attestation
            .verify_ppid(vec![0u8; 32], &accepted)
            .unwrap_err();
        assert_eq!(err, VerificationError::PpidWrongSize(32));
        assert!(err.to_string().contains("16 bytes"));
    }

    // Empty PPID is also size-rejected (size=0).
    #[test]
    fn verify_ppid_wrong_size_zero_bytes() {
        let attestation = create_mock_dstack_attestation();
        let accepted = vec![HexBytes::from([0u8; 16])];
        assert_eq!(
            attestation.verify_ppid(Vec::new(), &accepted),
            Err(VerificationError::PpidWrongSize(0))
        );
    }

    // -------- verify_static_rtmrs --------

    // All 8 RTMR/MRTD comparisons match → Ok.
    #[test]
    fn verify_static_rtmrs_accepts_when_all_match() {
        let attestation = create_mock_dstack_attestation();
        let report = td_report();
        let tcb = empty_tcb_info();
        let measurements: FullMeasurements = create_mock_full_measurements_hex().into();
        assert_eq!(
            attestation.verify_static_rtmrs(&report, &tcb, &measurements),
            Ok(())
        );
    }

    // Report-side RTMR0 mismatch fails.
    #[test]
    fn verify_static_rtmrs_rejects_rtmr0_mismatch() {
        let attestation = create_mock_dstack_attestation();
        let mut report = td_report();
        report.rt_mr0 = [0xFFu8; 48];
        let tcb = empty_tcb_info();
        let measurements: FullMeasurements = create_mock_full_measurements_hex().into();
        assert!(matches!(
            attestation.verify_static_rtmrs(&report, &tcb, &measurements),
            Err(VerificationError::WrongHash {
                name: "rtmr0_report_data",
                ..
            })
        ));
    }

    // Report-side MRTD mismatch fails.
    #[test]
    fn verify_static_rtmrs_rejects_mrtd_mismatch() {
        let attestation = create_mock_dstack_attestation();
        let mut report = td_report();
        report.mr_td = [0xABu8; 48];
        let tcb = empty_tcb_info();
        let measurements: FullMeasurements = create_mock_full_measurements_hex().into();
        assert!(matches!(
            attestation.verify_static_rtmrs(&report, &tcb, &measurements),
            Err(VerificationError::WrongHash {
                name: "mrtd_report_data",
                ..
            })
        ));
    }

    // TcbInfo-side RTMR1 mismatch fails (report and TcbInfo must agree).
    #[test]
    fn verify_static_rtmrs_rejects_tcb_info_rtmr1_mismatch() {
        let attestation = create_mock_dstack_attestation();
        let report = td_report();
        let mut tcb = empty_tcb_info();
        tcb.rtmr1 = HexBytes::from([0xAAu8; 48]);
        let measurements: FullMeasurements = create_mock_full_measurements_hex().into();
        assert!(matches!(
            attestation.verify_static_rtmrs(&report, &tcb, &measurements),
            Err(VerificationError::WrongHash {
                name: "rtmr1_tcb_info",
                ..
            })
        ));
    }

    // -------- verify_rtmr3 --------

    // RTMR3 hash mismatch fails before event-log replay (happy path is the
    // fixture test below).
    #[test]
    fn verify_rtmr3_rejects_hash_mismatch_before_event_log_replay() {
        let attestation = create_mock_dstack_attestation();
        let mut report = td_report();
        report.rt_mr3 = [0xFFu8; 48];
        let tcb = empty_tcb_info();
        assert!(matches!(
            attestation.verify_rtmr3(&report, &tcb),
            Err(VerificationError::WrongHash { name: "rtmr3", .. })
        ));
    }

    // -------- verify_app_compose_hash --------

    // compose_hash equal to the 32-byte expected payload passes.
    #[test]
    fn verify_app_compose_hash_accepts_matching_payload() {
        let attestation = create_mock_dstack_attestation();
        let mut tcb = empty_tcb_info();
        let payload = [0xABu8; 32];
        tcb.compose_hash = HexBytes::from(payload);
        assert_eq!(attestation.verify_app_compose_hash(&tcb, &payload), Ok(()));
    }

    // One byte changed in the expected payload fails.
    #[test]
    fn verify_app_compose_hash_rejects_single_byte_change() {
        let attestation = create_mock_dstack_attestation();
        let mut tcb = empty_tcb_info();
        let payload = [0xABu8; 32];
        tcb.compose_hash = HexBytes::from(payload);
        let mut tampered = payload;
        tampered[5] = 0x00;
        assert!(matches!(
            attestation.verify_app_compose_hash(&tcb, &tampered),
            Err(VerificationError::WrongHash {
                name: "app_compose_hash",
                ..
            })
        ));
    }

    // -------- validate_app_compose_payload --------

    // SHA-256 of the JSON string matches the hex event payload → Ok.
    #[test]
    fn validate_app_compose_payload_accepts_matching_hash() {
        let app_compose = r#"{"manifest_version":2}"#;
        let hash: [u8; 32] = Sha256::digest(app_compose.as_bytes()).into();
        let event_payload = hex::encode(hash);
        assert_eq!(
            DstackAttestation::validate_app_compose_payload(&event_payload, app_compose),
            Ok(())
        );
    }

    // One byte different in the event payload fails.
    #[test]
    fn validate_app_compose_payload_rejects_byte_flipped_payload() {
        let app_compose = r#"{"manifest_version":2}"#;
        let hash: [u8; 32] = Sha256::digest(app_compose.as_bytes()).into();
        let mut tampered = hash;
        tampered[0] = 0xFF;
        let event_payload = hex::encode(tampered);
        assert!(matches!(
            DstackAttestation::validate_app_compose_payload(&event_payload, app_compose),
            Err(VerificationError::WrongHash {
                name: "app_compose_payload",
                ..
            })
        ));
    }

    // -------- verify_any_measurements --------

    // No matching entry in the accepted list → WrongHash.
    #[test]
    fn verify_any_measurements_returns_wrong_hash_when_no_match() {
        let attestation = create_mock_dstack_attestation();
        let report = td_report();
        let tcb = empty_tcb_info();
        let mut tampered: FullMeasurements = create_mock_full_measurements_hex().into();
        tampered.rtmrs.rtmr0 = [0xFFu8; 48];
        assert!(matches!(
            attestation.verify_any_measurements(&report, &tcb, &[tampered]),
            Err(VerificationError::WrongHash {
                name: "expected_measurements",
                ..
            })
        ));
    }

    // -------- compare_hashes --------

    // Equal byte slices pass.
    #[test]
    fn compare_hashes_accepts_equal_byte_slices() {
        assert_eq!(compare_hashes("test", &[1, 2, 3], &[1, 2, 3]), Ok(()));
    }

    // One byte different fails.
    #[test]
    fn compare_hashes_rejects_one_byte_difference() {
        assert!(matches!(
            compare_hashes("test", &[1, 2, 3], &[1, 9, 3]),
            Err(VerificationError::WrongHash { name: "test", .. })
        ));
    }

    // Different lengths fail.
    #[test]
    fn compare_hashes_rejects_different_lengths() {
        assert!(matches!(
            compare_hashes("test", &[1, 2, 3], &[1, 2]),
            Err(VerificationError::WrongHash { name: "test", .. })
        ));
    }

    // -------- event_digest --------

    // Same input → same SHA-384 digest.
    #[test]
    fn event_digest_is_deterministic() {
        let a = DstackAttestation::event_digest(1, "compose-hash", b"payload");
        let b = DstackAttestation::event_digest(1, "compose-hash", b"payload");
        assert_eq!(a, b);
    }

    // Changing the event name changes the digest.
    #[test]
    fn event_digest_differs_on_event_change() {
        let a = DstackAttestation::event_digest(1, "compose-hash", b"x");
        let b = DstackAttestation::event_digest(1, "key-provider", b"x");
        assert_ne!(a, b);
    }

    // Changing the payload changes the digest.
    #[test]
    fn event_digest_differs_on_payload_change() {
        let a = DstackAttestation::event_digest(1, "compose-hash", b"a");
        let b = DstackAttestation::event_digest(1, "compose-hash", b"b");
        assert_ne!(a, b);
    }

    // -------- fixture-based happy paths --------
    //
    // The tests above use synthesised inputs. The two below use a TcbInfo
    // captured from a live TDX VM so the orchestrating verify_rtmr3 and
    // verify_app_compose run against a real event log.

    const TEST_TCB_INFO_STRING: &str = include_str!("../assets/tcb_info.json");

    fn fixture_tcb_info() -> TcbInfo {
        serde_json::from_str(TEST_TCB_INFO_STRING).expect("tcb_info.json is valid")
    }

    // Real event log replays cleanly to tcb_info.rtmr3.
    #[test]
    fn verify_rtmr3_accepts_fixture_event_log_replay() {
        let attestation = create_mock_dstack_attestation();
        let tcb = fixture_tcb_info();
        let mut report = td_report();
        report.rt_mr3 = *tcb.rtmr3;
        assert_eq!(attestation.verify_rtmr3(&report, &tcb), Ok(()));
    }

    // Real "compose-hash" event payload matches SHA-256 of app_compose.
    #[test]
    fn verify_app_compose_accepts_fixture() {
        let attestation = create_mock_dstack_attestation();
        let tcb = fixture_tcb_info();
        assert_eq!(attestation.verify_app_compose(&tcb), Ok(()));
    }
}