rialo-types 0.12.2

// Copyright (c) Subzero Labs, Inc.
// SPDX-License-Identifier: Apache-2.0

//! Core TEE types used in REX attestation

use borsh::{BorshDeserialize, BorshSerialize};
pub use rialo_tee_types::TeeType;
use serde::{Deserialize, Serialize};

// Re-export crypto types for backward compatibility
pub use crate::crypto::{PublicKey, TeeUserData, TEE_USER_DATA_SIZE, X25519_KEY_SIZE};

/// TEE platform attestation evidence
#[derive(
    Clone, Default, Eq, PartialEq, Serialize, Deserialize, BorshSerialize, BorshDeserialize,
)]
pub struct AttestationEvidence {
    pub tee_type: TeeType,
    pub evidence_data: Vec<u8>,
}

impl std::fmt::Debug for AttestationEvidence {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let truncated_data = if self.evidence_data.len() <= 16 {
            format!("{:?}", self.evidence_data)
        } else {
            format!(
                "{:?}... (truncated, total: {} bytes)",
                &self.evidence_data[..16],
                self.evidence_data.len()
            )
        };

        f.debug_struct("AttestationEvidence")
            .field("tee_type", &self.tee_type)
            .field("evidence_data", &truncated_data)
            .finish()
    }
}

/// SEV-SNP specific evidence data structure
///
/// This structure contains the raw SEV-SNP attestation data that gets
/// serialized/deserialized from the `evidence_data` field of `AttestationEvidence`.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct SevSnpEvidenceData {
    /// The raw SEV-SNP attestation report
    pub report: Vec<u8>,
    /// the VLEK or VCEK certificate
    pub vek_certificate: Vec<u8>,
}

/// Azure Confidential VM vTPM Evidence structure used for attestation verification
///
/// The vTPM-AK quote is the analogue of AWS Nitro's COSE-signed attestation document:
/// it carries PCR values from the Azure CVM's vTPM and a signature by the Attestation Key
/// (AK) over those PCRs together with a caller-supplied nonce.
///
/// Trust binding chain on Azure CVMs:
/// `user_data` (nonce) → (TPM AK signature) → AK pub → (`SHA256(RuntimeData)` in
/// `snp.report_data[..32]`) → SNP report → (AMD chain) → AMD root.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct AzureSnpVtpmEvidence {
    /// Serialized `az_snp_vtpm::vtpm::Quote` (bincode-encoded).
    ///
    /// Contains the TPM2B_ATTEST blob, the AK signature over it, and the bundled PCR
    /// digests whose hash matches the quote's `pcrDigest`.
    pub quote: Vec<u8>,
    /// Caller-supplied data bound to the quote as its nonce.
    ///
    /// On verification, the quote's recovered nonce must equal this value.
    pub nonce: Vec<u8>,
}

/// HCL-wrapped SEV-SNP attestation report + VCEK certificate, as produced by Azure
/// Confidential VMs.
///
/// Distinct from [`SevSnpEvidenceData`] (which carries a raw 1184-byte AMD
/// `AttestationReport` + VLEK) — Azure's `report` is the `az_cvm_vtpm::hcl::HclReport`
/// wire form and the certificate is a VCEK fetched from AMD KDS, not a VLEK.
///
/// HCL wire layout (frozen at `az-cvm-vtpm = 0.7.4`):
///
/// ```text
/// [   0 ..   32] AttestationHeader (8 × u32)
/// [  32 .. 1216] hardware SNP report (MAX_REPORT_SIZE)
/// [1216 .. 1236] IgvmRequestData fixed part (5 × u32)
/// [1236 ..     ] RuntimeData (JSON, embeds vTPM AK pub)
/// ```
///
/// The verifier parses the HCL wrapper to recover the raw SNP report and the AK pub.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct AzureSnpHclEvidence {
    /// HCL-wrapped SEV-SNP attestation report bytes (NOT a raw `AttestationReport`).
    pub report: Vec<u8>,
    /// VCEK certificate (DER), fetched from AMD KDS via `amd_kds::get_vcek`.
    pub vek_certificate: Vec<u8>,
}

/// Azure SEV-SNP specific evidence data structure
///
/// On Azure Confidential VMs the SNP report's `report_data` is set by the HCL layer to
/// `SHA256(RuntimeData)`, where `RuntimeData` embeds the vTPM AK public key. User data is
/// therefore bound through the AK and the vTPM quote (see `AzureSnpVtpmEvidence`).
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct AzureSevSnpEvidenceData {
    /// HCL-wrapped SEV-SNP attestation report + VCEK certificate fetched from AMD KDS.
    pub snp_evidence: AzureSnpHclEvidence,
    /// vTPM AK quote binding `user_data` (as nonce) and the platform PCR values.
    pub vtpm_evidence: AzureSnpVtpmEvidence,
}

/// AWS Nitro TPM COSE Evidence structure used for attestation verification
///
/// This structure contains the AWS Nitro TPM attestation data in COSE format along with
/// supporting metadata required for verification. The COSE report includes TPM-based
/// measurements and signatures that prove the identity and state of the AWS Nitro enclave.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct AwsNitroCoseEvidence {
    /// AWS COSE signed report containing TPM attestation data
    ///
    /// This report includes Platform Configuration Registers (PCRs)
    /// and digital signatures required to verify the authenticity and integrity of
    /// the enclave.
    pub report: Vec<u8>,
    /// Cryptographic nonce value used to prevent replay attacks
    ///
    /// This nonce is included in the attestation request and must match
    /// the value embedded in the COSE report to ensure freshness.
    pub nonce: Vec<u8>,
    /// Unix timestamp (in seconds) when the COSE report was generated
    ///
    /// Used to verify the attestation report's recency and detect
    /// potential replay attacks.
    pub time: u64,
}

/// AWS SEV-SNP specific evidence data structure optimized for AWS Nitro
///
/// This simplified structure leverages AWS Nitro's built-in attestation capabilities
/// and prepares for future AWS COSE signed reports that will include all TPM data.
/// The user_data contains the REX digest (32 bytes) that gets embedded in both
/// the SEV-SNP report and the AWS COSE report for binding.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct AwsSevSnpEvidenceData {
    /// SEV-SNP attestation report containing user_data
    pub snp_evidence: SevSnpEvidenceData,
    /// AWS Nitro COSE attestation report
    pub aws_nitro_cose_evidence: AwsNitroCoseEvidence,
}

/// Error type for TEE evidence parsing
#[derive(Debug, thiserror::Error)]
pub enum TeeEvidenceError {
    #[error("Invalid evidence: {0}")]
    InvalidEvidence(String),
}

impl TryFrom<&AttestationEvidence> for SevSnpEvidenceData {
    type Error = TeeEvidenceError;
    fn try_from(evidence: &AttestationEvidence) -> Result<Self, Self::Error> {
        parse_evidence(evidence, &[TeeType::SevSnp])
    }
}

impl TryFrom<&AttestationEvidence> for AwsSevSnpEvidenceData {
    type Error = TeeEvidenceError;
    fn try_from(evidence: &AttestationEvidence) -> Result<Self, Self::Error> {
        parse_evidence(evidence, &[TeeType::AwsSevSnp])
    }
}

impl TryFrom<&AttestationEvidence> for AzureSevSnpEvidenceData {
    type Error = TeeEvidenceError;
    fn try_from(evidence: &AttestationEvidence) -> Result<Self, Self::Error> {
        parse_evidence(evidence, &[TeeType::AzureSevSnp])
    }
}

/// Generic evidence parsing function that validates TEE type and deserializes evidence data
///
/// This function provides a unified way to parse evidence data across different TEE types,
/// validating that the evidence matches one of the expected TEE types before attempting
/// to deserialize the data into the target type.
pub fn parse_evidence<T>(
    evidence: &AttestationEvidence,
    expected_tee_types: &[TeeType],
) -> Result<T, TeeEvidenceError>
where
    T: for<'de> Deserialize<'de>,
{
    // Validate TEE type
    if !expected_tee_types.contains(&evidence.tee_type) {
        return Err(TeeEvidenceError::InvalidEvidence(format!(
            "Expected TEE types {:?}, got {:?}",
            expected_tee_types, evidence.tee_type
        )));
    }

    // Parse evidence data
    serde_json::from_slice(&evidence.evidence_data)
        .map_err(|e| TeeEvidenceError::InvalidEvidence(format!("Failed to parse evidence: {e}")))
}