Extending Atlas

Atlas is designed to be easily extended with new TEE verifiers. This guide walks through adding support for a new TEE (e.g., AMD SEV-SNP, ARM CCA).

Overview

Adding a new TEE verifier involves:

Create module structure for your TEE
Define report type and add to Report enum
Implement AtlsVerifier trait
Create config and builder
Create policy type implementing IntoVerifier
Add variants to Policy and Verifier enums
Re-export in lib.rs

Step 1: Create module structure

Create a new module for your TEE implementation:

core/src/
└── my_tee/
    ├── mod.rs
    ├── verifier.rs
    ├── config.rs
    └── policy.rs

In my_tee/mod.rs, re-export your types:

mod config;
mod policy;
mod verifier;

pub use config::{MyTeeVerifierConfig, MyTeeVerifierBuilder};
pub use policy::MyTeePolicy;
pub use verifier::MyTeeVerifier;

Step 2: Define report type

Add a variant to the Report enum in core/src/verifier.rs:

pub enum Report {
    Tdx(VerifiedReport),
    MyTee(MyTeeReport),  // Add your variant
}

impl Report {
    pub fn as_my_tee(&self) -> Option<&MyTeeReport> {
        match self {
            Report::MyTee(r) => Some(r),
            _ => None,
        }
    }

    pub fn into_my_tee(self) -> Option<MyTeeReport> {
        match self {
            Report::MyTee(r) => Some(r),
            _ => None,
        }
    }
}

Step 3: Implement AtlsVerifier

Create your verifier in my_tee/verifier.rs:

use crate::error::AtlsVerificationError;
use crate::verifier::{AsyncByteStream, AtlsVerifier, Report};

pub struct MyTeeVerifier {
    config: MyTeeVerifierConfig,
}

impl MyTeeVerifier {
    pub fn new(config: MyTeeVerifierConfig) -> Result<Self, AtlsVerificationError> {
        // Validate configuration
        Ok(Self { config })
    }

    pub fn builder() -> MyTeeVerifierBuilder {
        MyTeeVerifierBuilder::new()
    }

    async fn fetch_evidence<S>(
        &self,
        stream: &mut S,
        hostname: &str,
    ) -> Result<Vec<u8>, AtlsVerificationError>
    where
        S: AsyncByteStream,
    {
        // Fetch attestation evidence from server
        // e.g., HTTP POST to /attestation endpoint
        todo!()
    }

    fn verify_evidence(&self, evidence: &[u8]) -> Result<MyTeeReport, AtlsVerificationError> {
        // Verify evidence cryptographically
        // e.g., validate signatures, check measurements
        todo!()
    }

    fn verify_cert_binding(
        &self,
        peer_cert: &[u8],
        evidence: &[u8],
    ) -> Result<(), AtlsVerificationError> {
        // Verify certificate is bound to the evidence
        // e.g., cert hash in attestation report
        todo!()
    }

    fn verify_measurements(&self, report: &MyTeeReport) -> Result<(), AtlsVerificationError> {
        // Verify measurements against policy
        todo!()
    }
}

Native implementation

Implement AtlsVerifier for native targets:

#[cfg(not(target_arch = "wasm32"))]
impl AtlsVerifier for MyTeeVerifier {
    async fn verify<S>(
        &self,
        stream: &mut S,
        peer_cert: &[u8],
        session_ekm: &[u8],
        hostname: &str,
    ) -> Result<Report, AtlsVerificationError>
    where
        S: AsyncByteStream,
    {
        // 1. Fetch attestation evidence from server
        let evidence = self.fetch_evidence(stream, hostname).await?;

        // 2. Verify evidence cryptographically
        let report = self.verify_evidence(&evidence)?;

        // 3. Verify certificate binding
        self.verify_cert_binding(peer_cert, &evidence)?;

        // 4. Verify session EKM binding (optional, TEE-specific)
        // Bind session_ekm to evidence to prevent relay attacks

        // 5. Verify measurements against policy
        self.verify_measurements(&report)?;

        Ok(Report::MyTee(report))
    }
}

WASM implementation

Provide a WASM-compatible implementation:

#[cfg(target_arch = "wasm32")]
impl AtlsVerifier for MyTeeVerifier {
    async fn verify<S>(
        &self,
        stream: &mut S,
        peer_cert: &[u8],
        session_ekm: &[u8],
        hostname: &str,
    ) -> Result<Report, AtlsVerificationError>
    where
        S: AsyncByteStream,
    {
        // Same implementation as native version
        // Only the trait bounds differ
    }
}

The WASM implementation is typically identical to native. The main difference is trait bounds (no Send requirement).

Step 4: Create config and builder

In my_tee/config.rs, define configuration and builder:

use crate::error::AtlsVerificationError;
use super::MyTeeVerifier;

#[derive(Debug, Clone)]
pub struct MyTeeVerifierConfig {
    pub expected_measurement: Option<String>,
    pub allowed_status: Vec<String>,
    // ... other config fields
}

impl Default for MyTeeVerifierConfig {
    fn default() -> Self {
        Self {
            expected_measurement: None,
            allowed_status: vec!["Valid".into()],
        }
    }
}

pub struct MyTeeVerifierBuilder {
    config: MyTeeVerifierConfig,
}

impl MyTeeVerifierBuilder {
    pub fn new() -> Self {
        Self {
            config: MyTeeVerifierConfig::default(),
        }
    }

    pub fn expected_measurement(mut self, m: impl Into<String>) -> Self {
        self.config.expected_measurement = Some(m.into());
        self
    }

    pub fn allowed_status(mut self, statuses: Vec<String>) -> Self {
        self.config.allowed_status = statuses;
        self
    }

    pub fn build(self) -> Result<MyTeeVerifier, AtlsVerificationError> {
        MyTeeVerifier::new(self.config)
    }
}

Step 5: Create policy type

In my_tee/policy.rs, create a serde-compatible policy:

use serde::{Deserialize, Serialize};
use crate::error::AtlsVerificationError;
use crate::verifier::IntoVerifier;
use super::{MyTeeVerifier, MyTeeVerifierConfig};

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MyTeePolicy {
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub expected_measurement: Option<String>,
    
    #[serde(default = "default_allowed_status")]
    pub allowed_status: Vec<String>,
}

fn default_allowed_status() -> Vec<String> {
    vec!["Valid".into()]
}

impl Default for MyTeePolicy {
    fn default() -> Self {
        Self {
            expected_measurement: None,
            allowed_status: default_allowed_status(),
        }
    }
}

impl IntoVerifier for MyTeePolicy {
    type Verifier = MyTeeVerifier;

    fn into_verifier(self) -> Result<Self::Verifier, AtlsVerificationError> {
        let config = MyTeeVerifierConfig {
            expected_measurement: self.expected_measurement,
            allowed_status: self.allowed_status,
        };
        MyTeeVerifier::new(config)
    }
}

Step 6: Add to enums

Add to Policy enum

In core/src/policy.rs:

#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "type")]
pub enum Policy {
    #[serde(rename = "dstack_tdx")]
    DstackTdx(DstackTdxPolicy),
    
    #[serde(rename = "my_tee")]
    MyTee(MyTeePolicy),  // Add variant
}

impl Policy {
    pub fn into_verifier(self) -> Result<Verifier, AtlsVerificationError> {
        match self {
            Policy::DstackTdx(p) => Ok(Verifier::DstackTdx(p.into_verifier()?)),
            Policy::MyTee(p) => Ok(Verifier::MyTee(p.into_verifier()?)),  // Add arm
        }
    }
}

Add to Verifier enum

In core/src/verifier.rs:

pub enum Verifier {
    DstackTdx(DstackTDXVerifier),
    MyTee(MyTeeVerifier),  // Add variant
}

#[cfg(not(target_arch = "wasm32"))]
impl AtlsVerifier for Verifier {
    async fn verify<S>(
        &self,
        stream: &mut S,
        peer_cert: &[u8],
        session_ekm: &[u8],
        hostname: &str,
    ) -> Result<Report, AtlsVerificationError>
    where
        S: AsyncByteStream,
    {
        match self {
            Verifier::DstackTdx(v) => v.verify(stream, peer_cert, session_ekm, hostname).await,
            Verifier::MyTee(v) => v.verify(stream, peer_cert, session_ekm, hostname).await,  // Add arm
        }
    }
}

// Similar for wasm32 implementation

Step 7: Re-export in lib.rs

In core/src/lib.rs:

pub mod my_tee;

pub use my_tee::{
    MyTeePolicy,
    MyTeeVerifier,
    MyTeeVerifierBuilder,
    MyTeeVerifierConfig,
};

Testing your verifier

Create tests in my_tee/verifier.rs:

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_builder() {
        let verifier = MyTeeVerifier::builder()
            .expected_measurement("abc123")
            .build()
            .unwrap();
        
        assert!(verifier.config.expected_measurement.is_some());
    }

    #[test]
    fn test_policy_conversion() {
        let policy = MyTeePolicy {
            expected_measurement: Some("abc123".into()),
            allowed_status: vec!["Valid".into()],
        };
        
        let verifier = policy.into_verifier().unwrap();
        assert!(verifier.config.expected_measurement.is_some());
    }
}

Common patterns

Session binding

Bind the TLS session to the attestation evidence to prevent relay attacks:

// Include session_ekm in the evidence request
let nonce = generate_nonce();
let evidence_request = format!("{{\"nonce\": \"{}\", \"ekm\": \"{}\"}}", 
    hex::encode(nonce),
    hex::encode(session_ekm)
);

// Server includes SHA256(nonce || ekm) in attestation report_data
// Verify it matches
let expected = sha256(&[nonce, session_ekm].concat());
if evidence.report_data != expected {
    return Err(AtlsVerificationError::SessionBindingFailed);
}

Collateral caching

Cache verification collateral (CA certs, CRLs) to reduce network requests:

use std::sync::{Arc, RwLock};
use std::collections::HashMap;

struct CachedCollateral {
    data: Collateral,
    cached_at: u64,
}

pub struct MyTeeVerifier {
    config: MyTeeVerifierConfig,
    collateral_cache: Arc<RwLock<HashMap<String, CachedCollateral>>>,
}

impl MyTeeVerifier {
    async fn get_collateral(&self, key: &str) -> Result<Collateral, AtlsVerificationError> {
        // Check cache first
        if let Ok(guard) = self.collateral_cache.read() {
            if let Some(cached) = guard.get(key) {
                if !is_expired(cached.cached_at) {
                    return Ok(cached.data.clone());
                }
            }
        }

        // Fetch and cache
        let collateral = fetch_collateral(key).await?;
        if let Ok(mut guard) = self.collateral_cache.write() {
            guard.insert(key.to_string(), CachedCollateral {
                data: collateral.clone(),
                cached_at: now(),
            });
        }
        Ok(collateral)
    }
}

Design

Protocol Specification

Extending Atlas

Overview

Step 1: Create module structure

Step 2: Define report type

Step 3: Implement AtlsVerifier

Native implementation

WASM implementation

Step 4: Create config and builder

Step 5: Create policy type

Step 6: Add to enums

Add to Policy enum

Add to Verifier enum

Step 7: Re-export in lib.rs

Testing your verifier

Common patterns

Session binding

Collateral caching

See also

Build docs developers (and LLMs) love

Design

Protocol Specification

​Overview

​Step 1: Create module structure

​Step 2: Define report type

​Step 3: Implement AtlsVerifier

​Native implementation

​WASM implementation

​Step 4: Create config and builder

​Step 5: Create policy type

​Step 6: Add to enums

​Add to Policy enum

​Add to Verifier enum

​Step 7: Re-export in lib.rs

​Testing your verifier

​Common patterns

​Session binding

​Collateral caching

​See also

Build docs developers (and LLMs) love

Overview

Step 1: Create module structure

Step 2: Define report type

Step 3: Implement AtlsVerifier

Native implementation

WASM implementation

Step 4: Create config and builder

Step 5: Create policy type

Step 6: Add to enums

Add to Policy enum

Add to Verifier enum

Step 7: Re-export in lib.rs

Testing your verifier

Common patterns

Session binding

Collateral caching

See also