From 0c0c902d26c01c5c6f3b6598c92fdd7655b9d21e Mon Sep 17 00:00:00 2001
From: Aleksandr Karbyshev <karbyshev@users.noreply.github.com>
Date: Thu, 7 Mar 2024 22:47:17 +0100
Subject: [PATCH] Atomic `derive_store_spending_key_from_mnemonic_code`

---
 crates/apps/src/lib/cli/wallet.rs |  1 +
 crates/sdk/src/wallet/mod.rs      | 70 ++++++++++++++++++++-----------
 2 files changed, 46 insertions(+), 25 deletions(-)

diff --git a/crates/apps/src/lib/cli/wallet.rs b/crates/apps/src/lib/cli/wallet.rs
index eac47ec5a86..78c1834adb2 100644
--- a/crates/apps/src/lib/cli/wallet.rs
+++ b/crates/apps/src/lib/cli/wallet.rs
@@ -228,6 +228,7 @@ fn shielded_key_derive(
                 prompt_bip39_passphrase,
                 encryption_password,
             )
+            .expect("Failed to update the wallet storage.")
             .unwrap_or_else(|| {
                 edisplay_line!(io, "Failed to derive a key.");
                 display_line!(io, "No changes are persisted. Exiting.");
diff --git a/crates/sdk/src/wallet/mod.rs b/crates/sdk/src/wallet/mod.rs
index 1cbde642cf5..f1c21e162f4 100644
--- a/crates/sdk/src/wallet/mod.rs
+++ b/crates/sdk/src/wallet/mod.rs
@@ -843,26 +843,19 @@ impl<U: WalletIo> Wallet<U> {
         Some(derive_hd_secret_key(
             scheme,
             seed.as_bytes(),
-            derivation_path.clone(),
+            derivation_path,
         ))
     }
 
-    /// Restore a spending key from the user mnemonic code (read from stdin)
-    /// using a given ZIP32 derivation path and insert it into the store with
-    /// the provided alias, converted to lower case.
-    /// The key is encrypted with the provided password. If no password
-    /// provided, will prompt for password from stdin.
-    /// Stores the key in decrypted key cache and returns the alias of the key
-    /// and a reference-counting pointer to the key.
-    pub fn derive_store_spending_key_from_mnemonic_code(
-        &mut self,
-        alias: String,
-        alias_force: bool,
+    /// XXX OK
+    /// Derive a spending key from the user mnemonic code (read from stdin)
+    /// using a given ZIP32 derivation path.
+    pub fn derive_spending_key_from_mnemonic_code(
+        &self,
         derivation_path: DerivationPath,
         mnemonic_passphrase: Option<(Mnemonic, Zeroizing<String>)>,
         prompt_bip39_passphrase: bool,
-        password: Option<Zeroizing<String>>,
-    ) -> Option<(String, ExtendedSpendingKey)> {
+    ) -> Option<ExtendedSpendingKey> {
         let (mnemonic, passphrase) =
             if let Some(mnemonic_passphrase) = mnemonic_passphrase {
                 mnemonic_passphrase
@@ -876,17 +869,7 @@ impl<U: WalletIo> Wallet<U> {
                 (mnemonic, passphrase)
             };
         let seed = Seed::new(&mnemonic, &passphrase);
-        let spend_key =
-            derive_hd_spending_key(seed.as_bytes(), derivation_path.clone());
-
-        self.insert_spending_key(
-            alias,
-            alias_force,
-            spend_key,
-            password,
-            Some(derivation_path),
-        )
-        .map(|alias| (alias, spend_key))
+        Some(derive_hd_spending_key(seed.as_bytes(), derivation_path))
     }
 
     /// Generate a spending key similarly to how it's done for keypairs
@@ -1442,4 +1425,41 @@ impl<U: WalletIo + WalletStorage> Wallet<U> {
         .flatten()
         .transpose()
     }
+
+    // XXX OK
+    /// Derive a spending key from the user mnemonic code (read from stdin)
+    /// using a given ZIP32 derivation path and insert it into the store with
+    /// the provided alias, converted to lower case.
+    /// The key is encrypted with the provided password. If no password
+    /// provided, will prompt for password from stdin.
+    /// Stores the key in decrypted key cache and returns the alias of the key
+    /// and the derived spending key.
+    pub fn derive_store_spending_key_from_mnemonic_code(
+        &mut self,
+        alias: String,
+        alias_force: bool,
+        derivation_path: DerivationPath,
+        mnemonic_passphrase: Option<(Mnemonic, Zeroizing<String>)>,
+        prompt_bip39_passphrase: bool,
+        password: Option<Zeroizing<String>>,
+    ) -> Result<Option<(String, ExtendedSpendingKey)>, LoadStoreError> {
+        self.derive_spending_key_from_mnemonic_code(
+            derivation_path.clone(),
+            mnemonic_passphrase,
+            prompt_bip39_passphrase,
+        )
+        .map(|spend_key| {
+            self.insert_spending_key_atomic(
+                alias,
+                alias_force,
+                spend_key,
+                password,
+                Some(derivation_path),
+            )
+            .map(|o| o.map(|alias| (alias, spend_key)))
+            .transpose()
+        })
+        .flatten()
+        .transpose()
+    }
 }