Improve WalletBuilder (#1941)

* rename wallet data * remove awaits and update signatures * make it compile * rename more wallet data mentions * first impl * make it compile * unmut wallet * remove mut * remove more mut * update save load test * clean up * one more mut * refactor * error handling * fix tests * temporarily disable test * re-include test * review * some nits * revert backup rename; remove a function * revert restore rename * visibility; rename * set alias * ) Co-authored-by: Thoralf-M <[email protected]> * doc Co-authored-by: Thibault Martinez <[email protected]> * doc fixes * fix * get rid of the enum * little refactor and CI fix * fix todo * fix test * nits * clean up * nice up * only verify ed25519 address at first wallet creation * review * fix comment Co-authored-by: Thibault Martinez <[email protected]> * suggestions Co-authored-by: Thibault Martinez <[email protected]> Co-authored-by: Thoralf-M <[email protected]> --------- Co-authored-by: Thoralf-M <[email protected]> Co-authored-by: Thibault Martinez <[email protected]> Co-authored-by: Thibault Martinez <[email protected]>
iotaledger · Feb 21, 2024 · 6b3e64f · 6b3e64f
1 parent 43cc073
commit 6b3e64f
Show file tree

Hide file tree

Showing 9 changed files with 129 additions and 102 deletions.
diff --git a/bindings/core/src/lib.rs b/bindings/core/src/lib.rs
@@ -45,9 +45,9 @@ pub fn init_logger(config: String) -> std::result::Result<(), fern_logger::Error
 #[serde(rename_all = "camelCase")]
 pub struct WalletOptions {
     pub address: Option<Bech32Address>,
-    pub alias: Option<String>,
     #[serde(with = "option_bip44", default)]
     pub bip_path: Option<Bip44>,
+    pub alias: Option<String>,
     pub client_options: Option<ClientOptions>,
     #[derivative(Debug(format_with = "OmittedDebug::omitted_fmt"))]
     pub secret_manager: Option<SecretManagerDto>,
@@ -60,13 +60,13 @@ impl WalletOptions {
         self
     }
 
-    pub fn with_alias(mut self, alias: impl Into<Option<String>>) -> Self {
-        self.alias = alias.into();
+    pub fn with_bip_path(mut self, bip_path: impl Into<Option<Bip44>>) -> Self {
+        self.bip_path = bip_path.into();
         self
     }
 
-    pub fn with_bip_path(mut self, bip_path: impl Into<Option<Bip44>>) -> Self {
-        self.bip_path = bip_path.into();
+    pub fn with_alias(mut self, alias: impl Into<Option<String>>) -> Self {
+        self.alias = alias.into();
         self
     }
 
@@ -90,8 +90,8 @@ impl WalletOptions {
         log::debug!("wallet options: {self:?}");
         let mut builder = Wallet::builder()
             .with_address(self.address)
-            .with_alias(self.alias)
             .with_bip_path(self.bip_path)
+            .with_alias(self.alias)
             .with_client_options(self.client_options);
 
         #[cfg(feature = "storage")]

diff --git a/bindings/core/tests/secrets_debug.rs b/bindings/core/tests/secrets_debug.rs
@@ -26,6 +26,6 @@ fn method_interface_secrets_debug() {
     let wallet_options = WalletOptions::default().with_secret_manager(SecretManagerDto::Placeholder);
     assert_eq!(
         format!("{:?}", wallet_options),
-        "WalletOptions { address: None, alias: None, bip_path: None, client_options: None, secret_manager: Some(<omitted>), storage_path: None }"
+        "WalletOptions { address: None, bip_path: None, alias: None, client_options: None, secret_manager: Some(<omitted>), storage_path: None }"
     );
 }
diff --git a/cli/src/cli.rs b/cli/src/cli.rs
@@ -411,17 +411,18 @@ pub async fn init_command(
     } else {
         None
     };
-    let address = init_params
-        .address
-        .map(|addr| Bech32Address::from_str(&addr))
-        .transpose()?;
 
     Ok(Wallet::builder()
         .with_secret_manager(secret_manager)
         .with_client_options(ClientOptions::new().with_node(init_params.node_url.as_str())?)
         .with_storage_path(storage_path.to_str().expect("invalid unicode"))
+        .with_address(
+            init_params
+                .address
+                .map(|addr| Bech32Address::from_str(&addr))
+                .transpose()?,
+        )
         .with_bip_path(init_params.bip_path)
-        .with_address(address)
         .with_alias(alias)
         .finish()
         .await?)

diff --git a/sdk/examples/wallet/offline_signing/1_prepare_transaction.rs b/sdk/examples/wallet/offline_signing/1_prepare_transaction.rs
@@ -44,8 +44,8 @@ async fn main() -> Result<()> {
         .with_secret_manager(SecretManager::Placeholder)
         .with_storage_path(ONLINE_WALLET_DB_PATH)
         .with_client_options(client_options.clone())
-        .with_bip_path(Bip44::new(SHIMMER_COIN_TYPE))
         .with_address(address)
+        .with_bip_path(Bip44::new(SHIMMER_COIN_TYPE))
         .finish()
         .await?;
 

diff --git a/sdk/examples/wallet/spammer.rs b/sdk/examples/wallet/spammer.rs
@@ -60,8 +60,8 @@ async fn main() -> Result<()> {
     let wallet = Wallet::builder()
         .with_secret_manager(SecretManager::Mnemonic(secret_manager))
         .with_client_options(client_options)
-        .with_bip_path(bip_path)
         .with_address(address)
+        .with_bip_path(bip_path)
         .finish()
         .await?;
 

diff --git a/sdk/src/wallet/core/builder.rs b/sdk/src/wallet/core/builder.rs
@@ -17,7 +17,7 @@ use crate::wallet::storage::adapter::memory::Memory;
 use crate::wallet::storage::{StorageManager, StorageOptions};
 use crate::{
     client::secret::{GenerateAddressOptions, SecretManage, SecretManager},
-    types::block::address::{Address, Bech32Address},
+    types::block::address::{Bech32Address, Ed25519Address},
     wallet::{
         core::{operations::background_syncing::BackgroundSyncStatus, Bip44, WalletInner, WalletLedger},
         operations::syncing::SyncOptions,
@@ -29,8 +29,8 @@ use crate::{
 #[derive(Debug, Serialize)]
 #[serde(rename_all = "camelCase")]
 pub struct WalletBuilder<S: SecretManage = SecretManager> {
-    pub(crate) bip_path: Option<Bip44>,
     pub(crate) address: Option<Bech32Address>,
+    pub(crate) bip_path: Option<Bip44>,
     pub(crate) alias: Option<String>,
     pub(crate) client_options: Option<ClientOptions>,
     #[cfg(feature = "storage")]
@@ -42,8 +42,8 @@ pub struct WalletBuilder<S: SecretManage = SecretManager> {
 impl<S: SecretManage> Default for WalletBuilder<S> {
     fn default() -> Self {
         Self {
-            bip_path: Default::default(),
             address: Default::default(),
+            bip_path: Default::default(),
             alias: Default::default(),
             client_options: Default::default(),
             #[cfg(feature = "storage")]
@@ -59,21 +59,18 @@ where
 {
     /// Initialises a new instance of the wallet builder with the default storage adapter.
     pub fn new() -> Self {
-        Self {
-            secret_manager: None,
-            ..Default::default()
-        }
+        Self::default()
     }
 
-    /// Set the BIP44 path of the wallet.
-    pub fn with_bip_path(mut self, bip_path: impl Into<Option<Bip44>>) -> Self {
-        self.bip_path = bip_path.into();
+    /// Set the address of the wallet.
+    pub fn with_address(mut self, address: impl Into<Option<Bech32Address>>) -> Self {
+        self.address = address.into();
         self
     }
 
-    /// Set the wallet address.
-    pub fn with_address(mut self, address: impl Into<Option<Bech32Address>>) -> Self {
-        self.address = address.into();
+    /// Set the BIP44 path of the wallet.
+    pub fn with_bip_path(mut self, bip_path: impl Into<Option<Bip44>>) -> Self {
+        self.bip_path = bip_path.into();
         self
     }
 
@@ -178,46 +175,78 @@ where
             self.secret_manager = secret_manager;
         }
 
-        let restored_bip_path = loaded_wallet_builder.as_ref().and_then(|builder| builder.bip_path);
+        let mut verify_address = false;
+        let loaded_address = loaded_wallet_builder
+            .as_ref()
+            .and_then(|builder| builder.address.clone());
+
+        // May use a previously stored address if it wasn't provided
+        if let Some(address) = &self.address {
+            if let Some(loaded_address) = &loaded_address {
+                if address != loaded_address {
+                    return Err(crate::wallet::Error::WalletAddressMismatch(address.clone()));
+                }
+            } else {
+                verify_address = true;
+            }
+        } else {
+            self.address = loaded_address;
+        }
+
+        let loaded_bip_path = loaded_wallet_builder.as_ref().and_then(|builder| builder.bip_path);
 
         // May use a previously stored BIP path if it wasn't provided
         if let Some(bip_path) = self.bip_path {
-            if let Some(restored_bip_path) = restored_bip_path {
-                if bip_path != restored_bip_path {
+            if let Some(loaded_bip_path) = loaded_bip_path {
+                if bip_path != loaded_bip_path {
                     return Err(crate::wallet::Error::BipPathMismatch {
                         new_bip_path: Some(bip_path),
-                        old_bip_path: Some(restored_bip_path),
+                        old_bip_path: Some(loaded_bip_path),
                     });
                 }
+            } else {
+                verify_address = true;
             }
         } else {
-            self.bip_path = restored_bip_path;
+            self.bip_path = loaded_bip_path;
         }
 
+        // Create the node client.
+        let client = self
+            .client_options
+            .clone()
+            .ok_or(crate::wallet::Error::MissingParameter("client_options"))?
+            .finish()
+            .await?;
+
+        match (self.address.as_ref(), self.bip_path.as_ref()) {
+            (Some(address), Some(bip_path)) => {
+                if verify_address {
+                    // verify that the address is derived from the provided bip path.
+                    if let Some(backing_ed25519_address) = address.inner.backing_ed25519() {
+                        self.verify_ed25519_address(backing_ed25519_address, bip_path).await?;
+                    } else {
+                        return Err(crate::wallet::Error::InvalidParameter("address/bip_path mismatch"));
+                    }
+                }
+            }
+            (Some(_address), None) => {}
+            (None, Some(bip_path)) => {
+                self.address.replace(Bech32Address::new(
+                    client.get_bech32_hrp().await?,
+                    self.generate_ed25519_address(bip_path).await?,
+                ));
+            }
+            (None, None) => {
+                return Err(crate::wallet::Error::MissingParameter("address or bip_path"));
+            }
+        };
+
         // May use a previously stored wallet alias if it wasn't provided
         if self.alias.is_none() {
             self.alias = loaded_wallet_builder.as_ref().and_then(|builder| builder.alias.clone());
         }
 
-        // May use a previously stored wallet address if it wasn't provided
-        if self.address.is_none() {
-            self.address = loaded_wallet_builder
-                .as_ref()
-                .and_then(|builder| builder.address.clone());
-        }
-
-        // May create a default Ed25519 wallet address if there's a secret manager.
-        if self.address.is_none() {
-            if self.secret_manager.is_some() {
-                let address = self.create_default_wallet_address().await?;
-                self.address = Some(address);
-            } else {
-                return Err(crate::wallet::Error::MissingParameter("address"));
-            }
-        }
-        // Panic: can be safely unwrapped now
-        let wallet_address = self.address.as_ref().unwrap().clone();
-
         #[cfg(feature = "storage")]
         let mut wallet_ledger = storage_manager.load_wallet_ledger().await?;
 
@@ -232,14 +261,6 @@ where
             unlock_unused_inputs(wallet_ledger)?;
         }
 
-        // Create the node client.
-        let client = self
-            .client_options
-            .clone()
-            .ok_or(crate::wallet::Error::MissingParameter("client_options"))?
-            .finish()
-            .await?;
-
         let background_syncing_status = tokio::sync::watch::channel(BackgroundSyncStatus::Stopped);
         let background_syncing_status = (Arc::new(background_syncing_status.0), background_syncing_status.1);
 
@@ -263,7 +284,8 @@ where
         let wallet_ledger = WalletLedger::default();
 
         let wallet = Wallet {
-            address: Arc::new(RwLock::new(wallet_address)),
+            // Panic: it's invalid to build a wallet without an address.
+            address: Arc::new(RwLock::new(self.address.unwrap())),
             bip_path: Arc::new(RwLock::new(self.bip_path)),
             alias: Arc::new(RwLock::new(self.alias)),
             inner: Arc::new(wallet_inner),
@@ -279,41 +301,6 @@ where
         Ok(wallet)
     }
 
-    /// Generate the wallet address.
-    ///
-    /// Note: make sure to only call it after `self.secret_manager` and `self.bip_path` has been set.
-    pub(crate) async fn create_default_wallet_address(&self) -> crate::wallet::Result<Bech32Address> {
-        let bech32_hrp = self
-            .client_options
-            .as_ref()
-            .unwrap()
-            .network_info
-            .protocol_parameters
-            .bech32_hrp;
-        let bip_path = self.bip_path.as_ref().unwrap();
-
-        Ok(Bech32Address::new(
-            bech32_hrp,
-            Address::Ed25519(
-                self.secret_manager
-                    .as_ref()
-                    .unwrap()
-                    .read()
-                    .await
-                    .generate_ed25519_addresses(
-                        bip_path.coin_type,
-                        bip_path.account,
-                        bip_path.address_index..bip_path.address_index + 1,
-                        GenerateAddressOptions {
-                            internal: bip_path.change != 0,
-                            ledger_nano_prompt: false,
-                        },
-                    )
-                    .await?[0],
-            ),
-        ))
-    }
-
     #[cfg(feature = "storage")]
     pub(crate) async fn from_wallet(wallet: &Wallet<S>) -> Self {
         Self {
@@ -325,6 +312,37 @@ where
             secret_manager: Some(wallet.secret_manager.clone()),
         }
     }
+
+    #[inline(always)]
+    async fn verify_ed25519_address(
+        &self,
+        ed25519_address: &Ed25519Address,
+        bip_path: &Bip44,
+    ) -> crate::wallet::Result<()> {
+        (ed25519_address == &self.generate_ed25519_address(bip_path).await?)
+            .then_some(())
+            .ok_or(crate::wallet::Error::InvalidParameter("address/bip_path mismatch"))
+    }
+
+    async fn generate_ed25519_address(&self, bip_path: &Bip44) -> crate::wallet::Result<Ed25519Address> {
+        if let Some(secret_manager) = &self.secret_manager {
+            let secret_manager = &*secret_manager.read().await;
+            Ok(secret_manager
+                .generate_ed25519_addresses(
+                    bip_path.coin_type,
+                    bip_path.account,
+                    bip_path.address_index..bip_path.address_index + 1,
+                    GenerateAddressOptions {
+                        internal: bip_path.change != 0,
+                        ledger_nano_prompt: false,
+                    },
+                )
+                // Panic: if it didn't return an Err, then there must be at least one address
+                .await?[0])
+        } else {
+            Err(crate::wallet::Error::MissingParameter("secret_manager"))
+        }
+    }
 }
 
 // Check if any of the locked inputs is not used in a transaction and unlock them, so they get available for new

diff --git a/sdk/src/wallet/error.rs b/sdk/src/wallet/error.rs
@@ -56,6 +56,9 @@ pub enum Error {
     /// Invalid output kind.
     #[error("invalid output kind: {0}")]
     InvalidOutputKind(String),
+    /// Invalid parameter.
+    #[error("invalid parameter: {0}")]
+    InvalidParameter(&'static str),
     /// Invalid Voting Power
     #[cfg(feature = "participation")]
     #[cfg_attr(docsrs, doc(cfg(feature = "participation")))]

diff --git a/sdk/tests/wallet/common/mod.rs b/sdk/tests/wallet/common/mod.rs
@@ -29,10 +29,15 @@ pub use self::constants::*;
 ///
 /// A Wallet
 #[allow(dead_code, unused_variables)]
-pub(crate) async fn make_wallet(storage_path: &str, mnemonic: Option<Mnemonic>, node: Option<&str>) -> Result<Wallet> {
+pub(crate) async fn make_wallet(
+    storage_path: &str,
+    mnemonic: impl Into<Option<Mnemonic>>,
+    node: Option<&str>,
+) -> Result<Wallet> {
     let client_options = ClientOptions::new().with_node(node.unwrap_or(NODE_LOCAL))?;
-    let secret_manager =
-        MnemonicSecretManager::try_from_mnemonic(mnemonic.unwrap_or_else(|| Client::generate_mnemonic().unwrap()))?;
+    let secret_manager = MnemonicSecretManager::try_from_mnemonic(
+        mnemonic.into().unwrap_or_else(|| Client::generate_mnemonic().unwrap()),
+    )?;
 
     #[allow(unused_mut)]
     let mut wallet_builder = Wallet::builder()