fix: restore pki_env from disk whenever necessary

crypto/src/e2e_identity/init_certificates.rs

@@ -1,5 +1,6 @@
 use crate::{e2e_identity::CrlRegistration, prelude::MlsCentral, CryptoError, CryptoResult};
-use core_crypto_keystore::entities::{E2eiAcmeCA, E2eiCrl, E2eiIntermediateCert, UniqueEntity};
+use core_crypto_keystore::entities::{E2eiAcmeCA, E2eiCrl, E2eiIntermediateCert, EntityBase, UniqueEntity};
+use mls_crypto_provider::MlsCryptoProvider;
 use openmls_traits::OpenMlsCryptoProvider;
 use std::ops::DerefMut;
 use wire_e2e_identity::prelude::x509::{
@@ -17,11 +18,6 @@ impl MlsCentral {
     /// # Parameters
     /// * `trust_anchor_pem` - PEM certificate to anchor as a Trust Root
     pub async fn e2ei_register_acme_ca(&mut self, trust_anchor_pem: String) -> CryptoResult<()> {
-        // Bail if we have an environment already setup
-        if self.e2ei_pki_env.is_some() {
-            return Err(CryptoError::ConsumerError);
-        };
-
         let pki_env = PkiEnvironment::init(PkiEnvironmentParams {
             intermediates: Default::default(),
             trust_roots: Default::default(),
@@ -150,7 +146,13 @@ impl MlsCentral {
     }
 
     pub(crate) async fn init_pki_env(&mut self) -> CryptoResult<()> {
-        let keystore = self.mls_backend.key_store();
+        self.e2ei_pki_env
+            .replace(Self::restore_pki_env(&self.mls_backend).await?);
+        Ok(())
+    }
+
+    pub(crate) async fn restore_pki_env(backend: &MlsCryptoProvider) -> CryptoResult<PkiEnvironment> {
+        let keystore = backend.key_store();
         let mut conn = keystore.borrow_conn().await?;
 
         let mut trust_roots = vec![];
@@ -161,19 +163,15 @@ impl MlsCentral {
             );
         }
 
-        drop(conn);
-
-        let intermediates = keystore
-            .find_all::<E2eiIntermediateCert>(Default::default())
+        let intermediates = E2eiIntermediateCert::find_all(&mut conn, Default::default())
             .await?
             .into_iter()
             .try_fold(vec![], |mut acc, inter| {
                 acc.push(x509_cert::Certificate::from_der(&inter.content)?);
                 CryptoResult::Ok(acc)
             })?;
 
-        let crls = keystore
-            .find_all::<E2eiCrl>(Default::default())
+        let crls = E2eiCrl::find_all(&mut conn, Default::default())
             .await?
             .into_iter()
             .try_fold(vec![], |mut acc, crl| {
@@ -188,9 +186,6 @@ impl MlsCentral {
             time_of_interest: None,
         };
 
-        let pki_env = PkiEnvironment::init(params).map_err(|e| CryptoError::E2eiError(e.into()))?;
-        self.e2ei_pki_env.replace(pki_env);
-
-        Ok(())
+        PkiEnvironment::init(params).map_err(|e| CryptoError::E2eiError(e.into()))
     }
 }

crypto/src/e2e_identity/mod.rs

@@ -19,8 +19,8 @@
 pub mod error;
 pub(crate) mod id;
 pub(crate) mod identity;
-#[cfg(not(target_family = "wasm"))]
 pub(crate) mod init_certificates;
+#[cfg(not(target_family = "wasm"))]
 pub(crate) mod refresh_token;
 pub(crate) mod rotate;
 pub(crate) mod stash;
@@ -538,7 +538,7 @@
 #[cfg(test)]
 pub mod tests {
     use itertools::Itertools;
     use openmls_traits::OpenMlsCryptoProvider;
     use serde_json::json;
     use wasm_bindgen_test::*;
 
@@ -624,7 +624,7 @@
         cc: MlsCentral,
         case: &TestCase,
         client_id: Option<&str>,
         is_renewal: bool,
         init: impl Fn(E2eiInitWrapper) -> InitFnReturn<'_>,
         // used to verify persisting the instance actually does restore it entirely
         restore: impl Fn(E2eiEnrollment, MlsCentral) -> RestoreFnReturn<'a>,
@@ -634,8 +634,10 @@
             if is_renewal {
                 let initial_refresh_token =
                     crate::e2e_identity::refresh_token::RefreshToken::from("initial-refresh-token".to_string());
-                let initial_refresh_token = E2eiRefreshToken::from(initial_refresh_token);
+                let initial_refresh_token =
+                    core_crypto_keystore::entities::E2eiRefreshToken::from(initial_refresh_token);
                 let mut conn = cc.mls_backend.key_store().borrow_conn().await?;
+                use core_crypto_keystore::entities::UniqueEntity as _;
                 initial_refresh_token.replace(&mut conn).await.unwrap();
             }
         }

crypto/src/mls/mod.rs

@@ -190,12 +190,13 @@ impl MlsCentral {
 
         // Restore persisted groups if there are any
         let mls_groups = Self::restore_groups(&mls_backend).await?;
+        let e2ei_pki_env = Some(Self::restore_pki_env(&mls_backend).await?);
 
         Ok(Self {
             mls_backend,
             mls_client,
             mls_groups,
-            e2ei_pki_env: None,
+            e2ei_pki_env,
             callbacks: None,
         })
     }
@@ -225,12 +226,13 @@ impl MlsCentral {
             None
         };
         let mls_groups = Self::restore_groups(&mls_backend).await?;
+        let e2ei_pki_env = Some(Self::restore_pki_env(&mls_backend).await?);
 
         Ok(Self {
             mls_backend,
             mls_client,
             mls_groups,
-            e2ei_pki_env: None,
+            e2ei_pki_env,
             callbacks: None,
         })
     }

crypto/src/mls/restore.rs

@@ -12,6 +12,7 @@ impl MlsCentral {
     #[cfg_attr(test, crate::idempotent)]
     pub async fn restore_from_disk(&mut self) -> CryptoResult<()> {
         self.mls_groups = Self::restore_groups(&self.mls_backend).await?;
+        self.e2ei_pki_env = Some(Self::restore_pki_env(&self.mls_backend).await?);
         Ok(())
     }