fix: restore pki_env from disk whenever necessary

wireapp · Jan 24, 2024 · 8afd30a · 8afd30a
1 parent db0d451
commit 8afd30a
Show file tree

Hide file tree

Showing 4 changed files with 21 additions and 21 deletions.
diff --git a/crypto/src/e2e_identity/init_certificates.rs b/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()))
     }
 }
diff --git a/crypto/src/e2e_identity/mod.rs b/crypto/src/e2e_identity/mod.rs
@@ -19,8 +19,8 @@ pub mod enabled;
 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;
@@ -634,8 +634,10 @@ pub mod tests {
             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();
             }
         }

diff --git a/crypto/src/mls/mod.rs b/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,
         })
     }

diff --git a/crypto/src/mls/restore.rs b/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(())
     }