rialo-types 0.2.0-alpha.0

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

//! Core TEE types used in oracle attestation

use borsh::{BorshDeserialize, BorshSerialize};
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()
    }
}

// TODO: Decouple Cloud provider and TEE types
// <https://linear.app/subzero-labs/issue/SUB-931/decouple-cloud-provider-from-tee-type-in-core-types>
#[derive(
    Clone,
    Copy,
    Debug,
    Default,
    Eq,
    Hash,
    PartialEq,
    Serialize,
    Deserialize,
    BorshSerialize,
    BorshDeserialize,
)]
#[borsh(use_discriminant = true)]
#[non_exhaustive]
pub enum TeeType {
    SevSnp = 1,
    #[default]
    AwsSevSnp = 2,
    AzureSevSnp = 3,
    IntelTdx = 4,
    #[cfg(feature = "testing")]
    Mock = 254,
}

impl std::fmt::Display for TeeType {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            TeeType::AzureSevSnp => write!(f, "Azure SEV-SNP"),
            TeeType::AwsSevSnp => write!(f, "AWS SEV-SNP EC2 with NitroTPM"),
            TeeType::SevSnp => write!(f, "Generic SEV-SNP"),
            TeeType::IntelTdx => write!(f, "Intel TDX"),
            // This is a special TEE type used for testing purposes only.
            // It should never be included in production.
            #[cfg(feature = "testing")]
            TeeType::Mock => write!(f, "Mock TEE"),
        }
    }
}

/// 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 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 AzureSevSnpEvidenceData {
    /// SEV-SNP attestation report containing user_data
    pub snp_evidence: SevSnpEvidenceData,
}

/// 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 oracle 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}")))
}