Repository: google/OpenSK
File: src/api/attestation_store.rs (Rust)

Code before:
    fn get(&mut self, id: &Id) -> Result<Option<Attestation>, Error>;

    /// Sets the attestation for a given id.
    ///
    /// This function may not be supported.
    fn set(&mut self, id: &Id, attestation: Option<&Attestation>) -> Result<(), Error>;
}

/// Attestation store errors.
#[derive(Copy, Clone, Debug, PartialEq, Eq)]
#[non_exhaustive]
pub enum Error {
    Storage,
    Internal,
    NoSupport,
}

/// Keys of the environment store reserved for the attestation store.
pub const STORAGE_KEYS: &[usize] = &[1, 2];

/// Implements a default attestation store using the environment store.
///
/// Supports only one attestation at a time.
pub trait Helper: Env {
    /// Returns the current attestation id.
    fn attestation_id(&self) -> Id;
}

impl<T: Helper> AttestationStore for T {
    fn get(&mut self, id: &Id) -> Result<Option<Attestation>, Error> {
        if id != &self.attestation_id() {
            return Err(Error::NoSupport);
        }
        let private_key = self.store().find(PRIVATE_KEY_STORAGE_KEY)?;
        let certificate = self.store().find(CERTIFICATE_STORAGE_KEY)?;
        let (private_key, certificate) = match (private_key, certificate) {
            (Some(x), Some(y)) => (x, y),
            (None, None) => return Ok(None),
            _ => return Err(Error::Internal),
        };
        if private_key.len() != 32 {
            return Err(Error::Internal);
        }
        Ok(Some(Attestation {
            private_key: *array_ref![private_key, 0, 32],
            certificate,
        }))
    }

    fn set(&mut self, id: &Id, attestation: Option<&Attestation>) -> Result<(), Error> {
        if id != &self.attestation_id() {

Review comment (security):
What the idea behind this function? That's only useful if an `Env` will have exactly one attestation type? And if you enable enterprise attestation, you have to remember to also call this function here to keep the state consistent? Doesn't that make the Helper unusable when we have a setup where the incoming MakeCredential parameters decide whether we use Enterprise or Batch?

Code after:
    fn get(&mut self, id: &Id) -> Result<Option<Attestation>, Error>;

    /// Sets the attestation for a given id.
    ///
    /// This function may not be supported.
    fn set(&mut self, id: &Id, attestation: Option<&Attestation>) -> Result<(), Error>;
}

/// Attestation store errors.
#[derive(Copy, Clone, Debug, PartialEq, Eq)]
#[non_exhaustive]
pub enum Error {
    Storage,
    Internal,
    NoSupport,
}

/// Keys of the environment store reserved for the attestation store.
pub const STORAGE_KEYS: &[usize] = &[1, 2];

pub fn helper_get(env: &mut impl Env) -> Result<Option<Attestation>, Error> {
    let private_key = env.store().find(PRIVATE_KEY_STORAGE_KEY)?;
    let certificate = env.store().find(CERTIFICATE_STORAGE_KEY)?;
    let (private_key, certificate) = match (private_key, certificate) {
        (Some(x), Some(y)) => (x, y),
        (None, None) => return Ok(None),
        _ => return Err(Error::Internal),
    };
    if private_key.len() != 32 {
        return Err(Error::Internal);
    }
    Ok(Some(Attestation {
        private_key: *array_ref![private_key, 0, 32],
        certificate,
    }))
}

pub fn helper_set(env: &mut impl Env, attestation: Option<&Attestation>) -> Result<(), Error> {
    let updates = match attestation {
        None => [
            StoreUpdate::Remove {
                key: PRIVATE_KEY_STORAGE_KEY,
            },
            StoreUpdate::Remove {
                key: CERTIFICATE_STORAGE_KEY,
            },
        ],
        Some(attestation) => [
            StoreUpdate::Insert {
                key: PRIVATE_KEY_STORAGE_KEY,
                value: &attestation.private_key[..],