fix(HD-Wallet): load correct HD Accounts that truly belong to the HD Wallet

mm2src/coins/coin_balance.rs

@@ -483,59 +483,6 @@ pub mod common_impl {
             accounts: Vec::with_capacity(accounts.len() + 1),
         };
 
-        if accounts.get(&path_to_address.account_id).is_none() {
-            // Is seems that we couldn't find any HD account from the HD wallet storage.
-            drop(accounts);
-            info!(
-                "{} HD wallet hasn't been enabled before. Create default HD account",
-                coin.ticker()
-            );
-
-            // Create new HD account.
-            let mut new_account = create_new_account(coin, hd_wallet, xpub_extractor, Some(path_to_address.account_id))
-                .await
-                .map_mm_err()?;
-            let scan_new_addresses = matches!(
-                params.scan_policy,
-                EnableCoinScanPolicy::ScanIfNewWallet | EnableCoinScanPolicy::Scan
-            );
-
-            let account_balance = enable_hd_account(
-                coin,
-                hd_wallet,
-                &mut new_account,
-                path_to_address.chain,
-                &address_scanner,
-                scan_new_addresses,
-                params.min_addresses_number.max(Some(path_to_address.address_id + 1)),
-            )
-            .await?;
-            drop(new_account);
-
-            if coin.is_trezor() {
-                let enabled_address =
-                    hd_wallet
-                        .get_enabled_address()
-                        .await
-                        .ok_or(EnableCoinBalanceError::NewAddressDerivingError(
-                            NewAddressDerivingError::Internal(
-                                "Couldn't find enabled address after it has already been enabled".to_string(),
-                            ),
-                        ))?;
-                coin.received_enabled_address_from_hw_wallet(enabled_address)
-                    .await
-                    .map_err(|e| {
-                        EnableCoinBalanceError::NewAddressDerivingError(NewAddressDerivingError::Internal(format!(
-                            "Coin rejected the enabled address derived from the hardware wallet: {}",
-                            e
-                        )))
-                    })?;
-            }
-            // Todo: The enabled address should be indicated in the response.
-            result.accounts.push(account_balance);
-            return Ok(result);
-        }
-
         debug!(
             "{} HD accounts were found on {} coin activation",
             accounts.len(),
@@ -563,8 +510,41 @@ pub mod common_impl {
             .await?;
             result.accounts.push(account_balance);
         }
+        let account_of_enabled_address_exists = accounts.contains_key(&path_to_address.account_id);
         drop(accounts);
 
+        // Make sure that the account of the enabled address is available, or create it otherwise.
+        if !account_of_enabled_address_exists {
+            // Is seems that we couldn't find the account for the enabled address in the HD wallet storage.
+            info!(
+                "account={} hasn't been enabled before, enabling it now for coin={}",
+                path_to_address.account_id,
+                coin.ticker()
+            );
+
+            // Create new HD account.
+            let mut new_account = create_new_account(coin, hd_wallet, xpub_extractor, Some(path_to_address.account_id))
+                .await
+                .map_mm_err()?;
+            let scan_new_addresses = matches!(
+                params.scan_policy,
+                EnableCoinScanPolicy::ScanIfNewWallet | EnableCoinScanPolicy::Scan
+            );
+
+            let account_balance = enable_hd_account(
+                coin,
+                hd_wallet,
+                &mut new_account,
+                path_to_address.chain,
+                &address_scanner,
+                scan_new_addresses,
+                params.min_addresses_number.max(Some(path_to_address.address_id + 1)),
+            )
+            .await?;
+
+            result.accounts.push(account_balance);
+        }
+
         if coin.is_trezor() {
             let enabled_address =
                 hd_wallet
@@ -585,6 +565,7 @@ pub mod common_impl {
                 })?;
         }
 
+        // Todo: The enabled address should be indicated in the response.
         Ok(result)
     }
 

mm2src/coins/eth/v2_activation.rs

@@ -767,7 +767,7 @@ pub(crate) async fn build_address_and_priv_key_policy(
             let hd_wallet_storage = HDWalletCoinStorage::init_with_rmd160(ctx, ticker.to_string(), hd_wallet_rmd160)
                 .await
                 .mm_err(EthActivationV2Error::from)?;
-            let accounts = load_hd_accounts_from_storage(&hd_wallet_storage, &path_to_coin)
+            let accounts = load_hd_accounts_from_storage(&hd_wallet_storage, &path_to_coin, &global_hd_ctx)
                 .await
                 .map_mm_err()?;
             let gap_limit = gap_limit.unwrap_or(DEFAULT_GAP_LIMIT);
@@ -805,15 +805,12 @@ pub(crate) async fn build_address_and_priv_key_policy(
             let hd_wallet_storage = HDWalletCoinStorage::init_with_rmd160(ctx, ticker.to_string(), hd_wallet_rmd160)
                 .await
                 .mm_err(EthActivationV2Error::from)?;
-            let accounts = load_hd_accounts_from_storage(&hd_wallet_storage, &path_to_coin)
-                .await
-                .map_mm_err()?;
             let gap_limit = gap_limit.unwrap_or(DEFAULT_GAP_LIMIT);
             let hd_wallet = EthHDWallet {
                 hd_wallet_rmd160,
                 hd_wallet_storage,
-                derivation_path: path_to_coin.clone(),
-                accounts: HDAccountsMutex::new(accounts),
+                derivation_path: path_to_coin,
+                accounts: HDAccountsMutex::new(Default::default()),
                 enabled_address: *path_to_address,
                 gap_limit,
             };

mm2src/coins/hd_wallet/coin_ops.rs

@@ -218,13 +218,14 @@ pub trait HDWalletCoinOps {
         if new_known_addresses_number >= max_addresses_number {
             return MmError::err(AccountUpdatingError::AddressLimitReached { max_addresses_number });
         }
+        let account_xpub = hd_account.extended_pubkey().to_string(bip32::Prefix::XPUB);
         match chain {
             Bip44Chain::External => hd_wallet
-                .update_external_addresses_number(hd_account.account_id(), new_known_addresses_number)
+                .update_external_addresses_number(account_xpub, new_known_addresses_number)
                 .await
                 .map_mm_err()?,
             Bip44Chain::Internal => hd_wallet
-                .update_internal_addresses_number(hd_account.account_id(), new_known_addresses_number)
+                .update_internal_addresses_number(account_xpub, new_known_addresses_number)
                 .await
                 .map_mm_err()?,
         }

mm2src/coins/hd_wallet/mod.rs

@@ -1,7 +1,8 @@
 use async_trait::async_trait;
-use common::log::warn;
-use crypto::{Bip32DerPathOps, Bip32Error, Bip44Chain, ChildNumber, DerivationPath, HDPathToAccount, HDPathToCoin,
-             Secp256k1ExtendedPublicKey, StandardHDPath, StandardHDPathError};
+use bip32::ExtendedPublicKey;
+use common::log::{error, warn, LogOnError};
+use crypto::{Bip32DerPathOps, Bip32Error, Bip44Chain, ChildNumber, DerivationPath, GlobalHDAccountArc,
+             HDPathToAccount, HDPathToCoin, Secp256k1ExtendedPublicKey, StandardHDPath, StandardHDPathError};
 use futures::lock::{MappedMutexGuard as AsyncMappedMutexGuard, Mutex as AsyncMutex, MutexGuard as AsyncMutexGuard};
 use mm2_err_handle::prelude::*;
 use primitives::hash::H160;
@@ -242,25 +243,77 @@ where
     }
 }
 
-pub async fn load_hd_accounts_from_storage<HDAddress, ExtendedPublicKey>(
+// TODO: Ideally this should be general and not specific for secp256k1 keys only. i.e. we should have this function work with different types
+//       of extended public and private keys (like Solana's and Sia's ed25519 keys).
+//       This was generic in the past, but changed here: https://github.com/KomodoPlatform/komodo-defi-framework/pull/2482/
+pub async fn load_hd_accounts_from_storage<HDAddress>(
     hd_wallet_storage: &HDWalletCoinStorage,
     derivation_path: &HDPathToCoin,
-) -> HDWalletStorageResult<HDAccountsMap<HDAccount<HDAddress, ExtendedPublicKey>>>
+    global_hd_ctx: &GlobalHDAccountArc,
+) -> HDWalletStorageResult<HDAccountsMap<HDAccount<HDAddress, ExtendedPublicKey<secp256k1::PublicKey>>>>
 where
     HDAddress: HDAddressOps + Send,
-    ExtendedPublicKey: ExtendedPublicKeyOps,
-    <ExtendedPublicKey as FromStr>::Err: Display,
 {
+    let coin_der_path = derivation_path.to_derivation_path();
     let accounts = hd_wallet_storage.load_all_accounts().await?;
-    let res: HDWalletStorageResult<HDAccountsMap<HDAccount<HDAddress, ExtendedPublicKey>>> = accounts
-        .iter()
+
+    // Filter accounts that when derived from the current HD wallet, they yield the correct xpub.
+    // If they don't yield the correct xpub, this means that account belong to another HD wallet with
+    // different purpose'/coin_type' fields but same account_id' field. This happens when someone changes
+    // the purpose'/coin_type' fields in the coins cnofig without changing the ticker.
+    // TODO: This is a temporary solution that fixes the issue. But we might wanna re-design how accounts
+    //       are loaded from the database (loading only the single account the user asked for and not all accounts).
+    let mut bad_accounts = Vec::new();
+    let accounts = accounts.into_iter().filter(|account| {
+        let mut account_der_path = coin_der_path.clone();
+        let child_number = match ChildNumber::new(account.account_id, true) {
+            Ok(num) => num,
+            Err(e) => {
+                error!(
+                    "Failed to create ChildNumber for account_id={} found in storage: {}",
+                    account.account_id, e
+                );
+                return false;
+            },
+        };
+        account_der_path.push(child_number);
+        let derived_xpub = match global_hd_ctx.derive_secp256k1_xpub(&account_der_path) {
+            Ok(xpub) => xpub.to_string(bip32::Prefix::XPUB),
+            Err(e) => {
+                error!(
+                    "Failed to derive xpub for account_id={} found in storage: {}",
+                    account.account_id, e
+                );
+                return false;
+            },
+        };
+        if derived_xpub != account.account_xpub {
+            error!(
+                "Account with id={} has an xpub={} that doesn't match the one derived from the current HD wallet. Skipping.",
+                account.account_id, account.account_xpub
+            );
+            bad_accounts.push(account.account_xpub.clone());
+            return false;
+        }
+        true
+    });
+
+    let res: HDWalletStorageResult<_> = accounts
         .map(|account_info| {
-            let account = HDAccount::try_from_storage_item(derivation_path, account_info)?;
+            let account = HDAccount::try_from_storage_item(derivation_path, &account_info)?;
             Ok((account.account_id, account))
         })
         .collect();
+
     match res {
-        Ok(accounts) => Ok(accounts),
+        Ok(accounts) => {
+            // If there are any bad accounts, delete them from the storage.
+            if !bad_accounts.is_empty() {
+                // Ignore and log errors that occur during the deletion of bad accounts.
+                hd_wallet_storage.delete_accounts(bad_accounts).await.error_log();
+            }
+            Ok(accounts)
+        },
         Err(e) if e.get_inner().is_deserializing_err() => {
             warn!("Error loading HD accounts from the storage: '{}'. Clear accounts", e);
             hd_wallet_storage.clear_accounts().await?;
@@ -394,10 +447,19 @@ where
 {
     const INIT_ACCOUNT_ID: u32 = 0;
     let new_account_id = match account_id {
+        // TODO: We should never let the user choose to create a new account with arbitrary account_id. account_ids must be instantiated
+        //       in order and to instantiate account x+1, account x must have been already instantiated and used (i.e. has tx history).
+        //       ref. https://github.com/bitcoin/bips/blob/master/bip-0044.mediawiki#account
         Some(account_id) => account_id,
         None => {
             let accounts = hd_wallet.get_accounts_mut().await;
+            // Get's the last account ID we have.
             let last_account_id = accounts.iter().last().map(|(account_id, _account)| *account_id);
+            // This selects an available account ID, between 0..=last_id, and if all IDs are taken already, we use `last_id + 1`.
+            // TODO: This is inherently wrong. If account 0 is used, and account last_id is also used, then there should not exist
+            //       any accounts with ID such that "0 < ID < last_id" and they aren't used.
+            //       It is safe to just use `last_id + 1` but we gotta make sure that `last_id` was used and has some tx history.
+            //       ref. https://github.com/bitcoin/bips/blob/master/bip-0044.mediawiki#account
             last_account_id.map_or(INIT_ACCOUNT_ID, |last_id| {
                 (INIT_ACCOUNT_ID..=last_id)
                     .find(|id| !accounts.contains_key(id))

mm2src/coins/hd_wallet/storage/mock_storage.rs

@@ -1,5 +1,6 @@
 use crate::hd_wallet::{HDAccountStorageItem, HDWalletId, HDWalletStorageInternalOps, HDWalletStorageResult};
 use async_trait::async_trait;
+use crypto::XPub;
 use mm2_core::mm_ctx::MmArc;
 #[cfg(test)] use mocktopus::macros::*;
 
@@ -19,18 +20,10 @@ impl HDWalletStorageInternalOps for HDWalletMockStorage {
         unimplemented!()
     }
 
-    async fn load_account(
-        &self,
-        _wallet_id: HDWalletId,
-        _account_id: u32,
-    ) -> HDWalletStorageResult<Option<HDAccountStorageItem>> {
-        unimplemented!()
-    }
-
     async fn update_external_addresses_number(
         &self,
         _wallet_id: HDWalletId,
-        _account_id: u32,
+        _account_xpub: XPub,
         _new_external_addresses_number: u32,
     ) -> HDWalletStorageResult<()> {
         unimplemented!()
@@ -39,7 +32,7 @@ impl HDWalletStorageInternalOps for HDWalletMockStorage {
     async fn update_internal_addresses_number(
         &self,
         _wallet_id: HDWalletId,
-        _account_id: u32,
+        _account_xpub: XPub,
         _new_internal_addresses_number: u32,
     ) -> HDWalletStorageResult<()> {
         unimplemented!()
@@ -53,5 +46,9 @@ impl HDWalletStorageInternalOps for HDWalletMockStorage {
         unimplemented!()
     }
 
+    async fn delete_accounts(&self, _wallet_id: HDWalletId, _account_xpubs: Vec<XPub>) -> HDWalletStorageResult<()> {
+        unimplemented!()
+    }
+
     async fn clear_accounts(&self, _wallet_id: HDWalletId) -> HDWalletStorageResult<()> { unimplemented!() }
 }