From 5aee8e9ebfa7cc514d66e621d1e77422f37bf4d5 Mon Sep 17 00:00:00 2001
From: Alexander Schmidt <alexander.schmidt@iota.org>
Date: Mon, 11 Dec 2023 11:24:54 +0100
Subject: [PATCH] handle builder address cases

---
 bindings/core/src/method/wallet.rs         |  1 -
 bindings/core/src/method_handler/wallet.rs |  2 +-
 sdk/src/wallet/core/builder.rs             | 46 +++++++++++-----------
 3 files changed, 23 insertions(+), 26 deletions(-)

diff --git a/bindings/core/src/method/wallet.rs b/bindings/core/src/method/wallet.rs
index fa428a83bb..9126ad217c 100644
--- a/bindings/core/src/method/wallet.rs
+++ b/bindings/core/src/method/wallet.rs
@@ -18,7 +18,6 @@ use iota_sdk::{
     client::{
         api::{input_selection::Burn, PreparedTransactionDataDto, SignedTransactionDataDto},
         node_manager::node::NodeAuth,
-        secret::GenerateAddressOptions,
     },
     types::block::{
         address::{Bech32Address, Hrp},
diff --git a/bindings/core/src/method_handler/wallet.rs b/bindings/core/src/method_handler/wallet.rs
index 287c882433..8306d68b20 100644
--- a/bindings/core/src/method_handler/wallet.rs
+++ b/bindings/core/src/method_handler/wallet.rs
@@ -8,7 +8,7 @@ use iota_sdk::{
     client::api::{
         PreparedTransactionData, PreparedTransactionDataDto, SignedTransactionData, SignedTransactionDataDto,
     },
-    types::{block::address::ToBech32Ext, TryFromDto},
+    types::TryFromDto,
     wallet::{types::TransactionWithMetadataDto, PreparedCreateNativeTokenTransactionDto, Wallet},
 };
 
diff --git a/sdk/src/wallet/core/builder.rs b/sdk/src/wallet/core/builder.rs
index bef0b9a042..9ce97f26ec 100644
--- a/sdk/src/wallet/core/builder.rs
+++ b/sdk/src/wallet/core/builder.rs
@@ -16,11 +16,8 @@ use crate::wallet::storage::adapter::memory::Memory;
 #[cfg(feature = "storage")]
 use crate::wallet::storage::{StorageManager, StorageOptions};
 use crate::{
-    client::{
-        api::GetAddressesOptions,
-        secret::{GenerateAddressOptions, SecretManage, SecretManager},
-    },
-    types::block::address::{Address, Bech32Address},
+    client::secret::{GenerateAddressOptions, SecretManage, SecretManager},
+    types::block::address::Bech32Address,
     wallet::{
         core::{Bip44, WalletData, WalletInner},
         operations::syncing::SyncOptions,
@@ -74,12 +71,11 @@ where
         self
     }
 
-    // TODO: only add this fn if the SecretManager is optional. Otherwise the address should be generated.
-    // /// Set the wallet address.
-    // 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();
+        self
+    }
 
     /// Set the alias of the wallet.
     pub fn with_alias(mut self, alias: impl Into<Option<String>>) -> Self {
@@ -204,23 +200,25 @@ where
             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
-        // Since we don't allow setting it `address` must be None at this point.
-        self.address = loaded_wallet_builder
+        let provided_address = self.address;
+        let restored_address = loaded_wallet_builder
             .as_ref()
             .and_then(|builder| builder.address.clone());
 
-        let generated_address = Self::generate_wallet_address(client_options, bip_path, secret_manager).await?;
-        if let Some(address) = &self.address {
-            // Make sure the provided or loaded address doesn't conflict with the secret manager (private key) and the
-            // bip path.
-            if address != &generated_address {
-                // TODO: change to appropriate error, e.g. Error::ConflictingWalletAddress or smth.
-                return Err(crate::wallet::Error::MissingParameter("address"));
+        self.address = Some(match (provided_address, restored_address) {
+            (Some(provided), Some(restored)) => {
+                // error early if they conflict
+                if provided != restored {
+                    // TODO: change to appropriate error, e.g. Error::ConflictingWalletAddress or smth.
+                    return Err(crate::wallet::Error::MissingParameter("address"));
+                } else {
+                    restored
+                }
             }
-        } else {
-            self.address = Some(generated_address);
-        }
+            (Some(provided), None) => provided,
+            (None, Some(restored)) => restored,
+            (None, None) => Self::generate_wallet_address(client_options, bip_path, secret_manager).await?,
+        });
 
         // Panic: an address must exist at this point
         let address = self.address.as_ref().unwrap().clone();