Persist Stronghold's changes only when its handle is dropped

iotaledger · Apr 26, 2024 · 096bb30 · 096bb30
1 parent 614719b
commit 096bb30
Show file tree

Hide file tree

Showing 3 changed files with 17 additions and 8 deletions.
diff --git a/identity_stronghold/src/stronghold_jwk_storage.rs b/identity_stronghold/src/stronghold_jwk_storage.rs
@@ -136,7 +136,6 @@ impl JwkStorage for StrongholdStorage {
           .with_custom_message("stronghold public key procedure failed")
           .with_source(err)
       })?;
-    persist_changes(self, stronghold).await?;
     let public_key: Vec<u8> = procedure_result.into();
 
     let mut params = JwkParamsOkp::new();
@@ -188,7 +187,6 @@ impl JwkStorage for StrongholdStorage {
           .with_custom_message("stronghold write secret failed")
           .with_source(err)
       })?;
-    persist_changes(self, stronghold).await?;
 
     Ok(key_id)
   }
@@ -263,7 +261,6 @@ impl JwkStorage for StrongholdStorage {
     if !deleted {
       return Err(KeyStorageError::new(KeyStorageErrorKind::KeyNotFound));
     }
-    persist_changes(self, stronghold).await?;
 
     Ok(())
   }
@@ -283,3 +280,17 @@ impl JwkStorage for StrongholdStorage {
     Ok(exists)
   }
 }
+
+/// Calls `persist_changes` when `StrongholdStorage` gets dropped.
+impl Drop for StrongholdStorage {
+  fn drop(&mut self) {
+    let secret_manager = std::mem::replace(&mut self.0, Arc::new(SecretManager::Placeholder));
+    tokio::spawn(async move {
+      let SecretManager::Stronghold(stronghold) = secret_manager.as_ref() else {
+        return;
+      };
+      let stronghold = stronghold.inner().await;
+      let _ = persist_changes(&secret_manager, stronghold).await;
+    });
+  }
+}
diff --git a/identity_stronghold/src/stronghold_jwk_storage_ext.rs b/identity_stronghold/src/stronghold_jwk_storage_ext.rs
@@ -73,8 +73,6 @@ impl JwkStorageBbsPlusExt for StrongholdStorage {
           .with_source(e)
       })?;
 
-    persist_changes(self, stronghold).await?;
-
     Ok(JwkGenOutput::new(kid, jwk))
   }
 

diff --git a/identity_stronghold/src/utils.rs b/identity_stronghold/src/utils.rs
@@ -6,14 +6,14 @@ use identity_storage::KeyStorageError;
 use identity_storage::KeyStorageErrorKind;
 use identity_storage::KeyStorageResult;
 use identity_verification::jws::JwsAlgorithm;
+use iota_sdk::client::secret::SecretManager;
 use iota_stronghold::Client;
 use iota_stronghold::ClientError;
 use iota_stronghold::Stronghold;
 use rand::distributions::DistString as _;
 use tokio::sync::MutexGuard;
 
 use crate::stronghold_key_type::StrongholdKeyType;
-use crate::StrongholdStorage;
 
 pub static IDENTITY_VAULT_PATH: &str = "iota_identity_vault";
 pub static IDENTITY_CLIENT_PATH: &[u8] = b"iota_identity_client";
@@ -54,7 +54,7 @@ fn load_or_create_client(stronghold: &Stronghold) -> KeyStorageResult<Client> {
 }
 
 pub async fn persist_changes(
-  secret_manager: &StrongholdStorage,
+  secret_manager: &SecretManager,
   stronghold: MutexGuard<'_, Stronghold>,
 ) -> KeyStorageResult<()> {
   stronghold.write_client(IDENTITY_CLIENT_PATH).map_err(|err| {
@@ -65,7 +65,7 @@ pub async fn persist_changes(
   // Must be dropped since `write_stronghold_snapshot` needs to acquire the stronghold lock.
   drop(stronghold);
 
-  match secret_manager.as_secret_manager() {
+  match secret_manager {
     iota_sdk::client::secret::SecretManager::Stronghold(stronghold_manager) => {
       stronghold_manager
         .write_stronghold_snapshot(None)