detee-sgx/src/config.rs

use crate::{RaTlsConfigBuilder, RaTlsError};

#[cfg(feature = "occlum")]
use crate::quote::{Quote, STATIC_QUOTE};

use rustls::{ClientConfig, ServerConfig};

pub type Measurement = [u8; 32];

#[derive(Default, Debug)]
pub struct RaTlsConfig {
    #[cfg(feature = "occlum")]
    pub(crate) allowed_instances: Vec<InstanceMeasurement>,
}

#[cfg(feature = "occlum")]
#[derive(Default, Debug, Clone)]
pub struct InstanceMeasurement {
    pub(crate) mrsigners: Option<Vec<Measurement>>,
    pub(crate) mrenclaves: Option<Vec<Measurement>>,
    pub(crate) product_ids: Option<Vec<u16>>,
    pub(crate) versions: Option<Vec<u16>>,
}

#[cfg(feature = "occlum")]
impl InstanceMeasurement {
    pub fn new() -> Self {
        Self::default()
    }

    pub fn with_mrsigners(self, mrsigners: Vec<Measurement>) -> Self {
        Self {
            mrsigners: Some(mrsigners),
            ..self
        }
    }

    pub fn with_mrenclaves(self, mrenclaves: Vec<Measurement>) -> Self {
        Self {
            mrenclaves: Some(mrenclaves),
            ..self
        }
    }

    pub fn with_product_ids(self, product_ids: Vec<u16>) -> Self {
        Self {
            product_ids: Some(product_ids),
            ..self
        }
    }

    pub fn with_versions(self, versions: Vec<u16>) -> Self {
        Self {
            versions: Some(versions),
            ..self
        }
    }

    pub fn with_current_mrsigner(self) -> Result<Self, RaTlsError> {
        let quote = Self::generate_static_empty_quote()?;
        let mrsigner_from_quote = quote.mrsigner().into();

        if let Some(mrsigners) = self.mrsigners {
            if mrsigners.contains(&mrsigner_from_quote) {
                // already contains the mrsigner
                Ok(self)
            } else {
                let mut mrsigners = self.mrsigners.clone().unwrap();
                mrsigners.push(mrsigner_from_quote);
                Ok(self.with_mrsigners(mrsigners))
            }
        } else {
            Ok(Self {
                mrsigners: Some(vec![mrsigner_from_quote]),
                ..self
            })
        }
    }

    pub fn with_current_mrenclave(self) -> Result<Self, RaTlsError> {
        let quote = Self::generate_static_empty_quote()?;
        let mrenclave_from_quote = quote.mrenclave().into();

        if let Some(mrenclaves) = self.mrenclaves {
            if mrenclaves.contains(&mrenclave_from_quote) {
                // already contains the mrenclave
                Ok(self)
            } else {
                let mut mrenclaves = self.mrenclaves.clone().unwrap();
                mrenclaves.push(mrenclave_from_quote);
                Ok(self.with_mrenclaves(mrenclaves))
            }
        } else {
            Ok(Self {
                mrenclaves: Some(vec![mrenclave_from_quote]),
                ..self
            })
        }
    }

    pub fn with_current_measurements() -> Result<Self, RaTlsError> {
        let quote = Self::generate_static_empty_quote()?;

        Ok(Self {
            mrsigners: Some(vec![quote.mrsigner().into()]),
            mrenclaves: Some(vec![quote.mrenclave().into()]),
            product_ids: Some(vec![quote.product_id()]),
            versions: Some(vec![quote.version()]),
        })
    }

    fn generate_static_empty_quote() -> Result<&'static Quote, RaTlsError> {
        Ok(STATIC_QUOTE.as_ref().map_err(|e| e.clone())?)
    }

    pub(crate) fn check_quote_measurements(&self, quote: &Quote) -> bool {
        let mut result = false;
        if let Some(mrsigners) = &self.mrsigners {
            result = true;
            let value = quote.mrsigner().into();
            if !mrsigners.contains(&value) {
                return false;
            }
        }
        if let Some(mrenclaves) = &self.mrenclaves {
            result = true;
            let value = quote.mrenclave().into();
            if !mrenclaves.contains(&value) {
                return false;
            }
        }

        if let Some(product_ids) = &self.product_ids {
            result = true;
            let value = quote.product_id().into();
            if !product_ids.contains(&value) {
                return false;
            }
        }

        if let Some(versions) = &self.versions {
            result = true;
            let value = quote.version().into();
            if !versions.contains(&value) {
                return false;
            }
        }

        result
    }
}

impl RaTlsConfig {
    pub fn new() -> Self {
        Self::default()
    }

    #[cfg(feature = "occlum")]
    pub fn allow_instance_measurement(mut self, instance_measurement: InstanceMeasurement) -> Self {
        self.allowed_instances.push(instance_measurement);
        self
    }

    #[cfg(feature = "occlum")]
    pub(crate) fn is_allowed_quote(&self, quote: &Quote) -> Result<(), RaTlsError> {
        match self
            .allowed_instances
            .iter()
            .any(|im| im.check_quote_measurements(quote))
        {
            true => Ok(()),
            false => Err(RaTlsError::QuoteError(format!(
                "{:?} is not allowed",
                quote
            ))),
        }
    }

    pub fn into_server_config(self) -> Result<ServerConfig, RaTlsError> {
        ServerConfig::from_ratls_config(self)
    }

    pub fn into_client_config(self) -> Result<ClientConfig, RaTlsError> {
        ClientConfig::from_ratls_config(self)
    }
}