core_crypto/transaction_context/

key_package.rs

//! This module contains all transactional behavior related to key packages

use std::time::Duration;

use core_crypto_keystore::{
    entities::{StoredEncryptionKeyPair, StoredHpkePrivateKey, StoredKeyPackage},
    traits::FetchFromDatabase as _,
};
use openmls::prelude::{CryptoConfig, Lifetime};

use super::{Error, Result, TransactionContext};
use crate::{
    ConversationConfiguration, CredentialRef, Keypackage, KeypackageRef, KeystoreError, RecursiveError,
    mls::key_package::KeypackageExt as _,
};

/// Default lifetime of all generated Keypackages. Matches the limit defined in openmls
pub const KEYPACKAGE_DEFAULT_LIFETIME: Duration = Duration::from_secs(60 * 60 * 24 * 28 * 3); // ~3 months

impl TransactionContext {
    /// Generate a [Keypackage] from the referenced credential.
    ///
    /// Makes no attempt to look up or prune existing keypackges.
    ///
    /// If `lifetime` is set, the keypackages will expire that span into the future.
    /// If it is unset, [`KEYPACKAGE_DEFAULT_LIFETIME`]
    /// is used.
    ///
    /// As a side effect, stores the keypackages and some related data in the keystore.
    pub async fn generate_key_package(
        &self,
        credential_ref: &CredentialRef,
        lifetime: Option<Duration>,
    ) -> Result<Keypackage> {
        let inner = self.inner().await?;
        let lifetime = Lifetime::new(lifetime.unwrap_or(KEYPACKAGE_DEFAULT_LIFETIME).as_secs());
        let credential = credential_ref
            .load(&inner.transaction)
            .await
            .map_err(RecursiveError::context("loading credential"))?;
        let config = CryptoConfig {
            ciphersuite: credential.cipher_suite.into(),
            version: openmls::versions::ProtocolVersion::default(),
        };

        Keypackage::builder()
            .leaf_node_capabilities(ConversationConfiguration::default_leaf_capabilities())
            .key_package_lifetime(lifetime)
            .build(
                config,
                &self.crypto_provider().await?,
                &credential.signature_key_pair,
                credential.to_mls_credential_with_key(),
            )
            .await
            .map_err(Error::key_package_new())
    }

    /// Get all [`KeypackageRef`]s known to the keystore.
    pub async fn get_key_package_refs(&self) -> Result<Vec<KeypackageRef>> {
        let session = self.session().await?;
        session
            .get_keypackage_refs()
            .await
            .map_err(RecursiveError::context("getting all key package refs for transaction"))
            .map_err(Into::into)
    }

    /// Remove one [`Keypackage`] from the database.
    ///
    /// Succeeds silently if the keypackage does not exist in the database.
    ///
    /// Implementation note: this must first load and deserialize the keypackage,
    /// then remove items from three distinct tables.
    pub async fn remove_key_package(&self, kp_ref: &KeypackageRef) -> Result<()> {
        let Some(kp) = self
            .session()
            .await?
            .load_key_package(kp_ref)
            .await
            .map_err(RecursiveError::context("loading key packages on session"))?
        else {
            return Ok(());
        };

        let inner = self.inner().await?;
        inner
            .transaction
            .remove_borrowed::<StoredKeyPackage>(kp_ref.hash_ref())
            .await
            .map_err(KeystoreError::wrap("removing key package from keystore"))?;
        inner
            .transaction
            .remove_borrowed::<StoredHpkePrivateKey>(kp.hpke_init_key().as_slice())
            .await
            .map_err(KeystoreError::wrap("removing private key from keystore"))?;
        inner
            .transaction
            .remove_borrowed::<StoredEncryptionKeyPair>(kp.leaf_node().encryption_key().as_slice())
            .await
            .map_err(KeystoreError::wrap("removing encryption keypair from keystore"))?;

        Ok(())
    }

    /// Remove all keypackages associated with this credential.
    ///
    /// This is fairly expensive as it must first load all keypackages, then delete those matching the credential.
    ///
    /// Implementation note: once it makes it as far as having a list of keypackages, does _not_ short-circuit
    /// if removing one returns an error. In that case, only the first produced error is returned.
    /// This helps ensure that as many keypackages for the given credential ref are removed as possible.
    pub async fn remove_key_packages_for(&self, credential_ref: &CredentialRef) -> Result<()> {
        let inner = self.inner().await?;
        let credential = credential_ref
            .load(&inner.transaction)
            .await
            .map_err(RecursiveError::context("loading credential"))?;
        let signature_public_key = credential.signature_key_pair.public();

        let mut first_err = None;
        macro_rules! try_retain_err {
            ($e:expr) => {
                match $e {
                    Err(err) => {
                        if first_err.is_none() {
                            first_err = Some(Error::from(err));
                        }
                        continue;
                    }
                    Ok(val) => val,
                }
            };
        }

        let session = self.session().await?;
        for keypackage in session
            .get_key_packages()
            .await
            .map_err(RecursiveError::context("loading key packages"))?
            .into_iter()
            .filter(|keypackage| keypackage.leaf_node().signature_key().as_slice() == signature_public_key)
        {
            let kp_ref = try_retain_err!(keypackage.make_ref());
            try_retain_err!(self.remove_key_package(&kp_ref).await);
        }

        match first_err {
            None => Ok(()),
            Some(err) => Err(err),
        }
    }

    /// Restore a key package that was deleted in this transaction by removing it from the deleted list. This is
    /// idempotent: if the key package doesn't exist in the deleted list, do nothing.
    ///
    /// NOTE: This will only work if the key package has been added in an earlier transaction, because otherwise,
    /// removing its id from the deleted list wouldn't suffice: we'd need to replay its insertion.
    pub(crate) async fn restore_key_package(&self, key_package_ref: &[u8]) -> Result<()> {
        let inner = self.inner().await?;

        inner
            .transaction
            .restore::<StoredKeyPackage>(key_package_ref)
            .await
            .map_err(KeystoreError::wrap(
                "restoring key package deleted in current transaction",
            ))?;

        let Some(key_package) = inner
            .transaction
            .get_borrowed::<StoredKeyPackage>(key_package_ref)
            .await
            .map_err(KeystoreError::wrap("loading keypackage from database"))?
            .map(crate::mls::session::key_package::from_stored)
            .transpose()
            .map_err(RecursiveError::context("loading key package"))?
        else {
            return Ok(());
        };

        inner
            .transaction
            .restore::<StoredHpkePrivateKey>(key_package.hpke_init_key().as_slice())
            .await
            .map_err(KeystoreError::wrap("restoring private key from keystore"))?;
        inner
            .transaction
            .restore::<StoredEncryptionKeyPair>(key_package.leaf_node().encryption_key().as_slice())
            .await
            .map_err(KeystoreError::wrap("restoring encryption keypair from keystore"))?;

        Ok(())
    }
}