fix(sync): always sync newly generated change addresses (#491)

iotaledger · Apr 1, 2021 · 1887007 · 1887007
1 parent 6e3db5c
commit 1887007
Show file tree

Hide file tree

Showing 5 changed files with 31 additions and 6 deletions.
diff --git a/bindings/java/native/src/types/transaction.rs b/bindings/java/native/src/types/transaction.rs
@@ -130,7 +130,7 @@ pub struct TransactionInput {
 impl TransactionInput {
     pub fn kind(&self) -> InputKind {
         match self.input {
-            RustWalletInput::Utxo(_) => InputKind::Utxo,
+            RustWalletInput::UTXO(_) => InputKind::Utxo,
             RustWalletInput::Treasury(_) => InputKind::Treasury,
         }
     }

diff --git a/bindings/python/native/src/types/message.rs b/bindings/python/native/src/types/message.rs
@@ -194,7 +194,7 @@ impl TryFrom<RustWalletTransactionEssence> for Essence {
                     .iter()
                     .cloned()
                     .map(|input| {
-                        if let RustWalletInput::Utxo(input) = input {
+                        if let RustWalletInput::UTXO(input) = input {
                             Input {
                                 transaction_id: input.input.output_id().transaction_id().to_string(),
                                 index: input.input.output_id().index(),

diff --git a/src/account/mod.rs b/src/account/mod.rs
@@ -19,6 +19,7 @@ use serde::{Deserialize, Serialize};
 use tokio::sync::{Mutex, RwLock, RwLockWriteGuard};
 
 use std::{
+    collections::HashSet,
     hash::{Hash, Hasher},
     ops::Deref,
     path::PathBuf,
@@ -389,6 +390,7 @@ pub struct AccountHandle {
     pub(crate) locked_addresses: Arc<Mutex<Vec<AddressWrapper>>>,
     pub(crate) account_options: AccountOptions,
     is_mqtt_enabled: Arc<AtomicBool>,
+    pub(crate) change_addresses_to_sync: Arc<Mutex<HashSet<AddressWrapper>>>,
 }
 
 impl AccountHandle {
@@ -399,6 +401,7 @@ impl AccountHandle {
             locked_addresses: Default::default(),
             account_options,
             is_mqtt_enabled: Arc::new(AtomicBool::new(true)),
+            change_addresses_to_sync: Default::default(),
         }
     }
 

diff --git a/src/account/sync/mod.rs b/src/account/sync/mod.rs
@@ -387,6 +387,7 @@ async fn sync_messages(
     skip_addresses: &[Address],
     options: AccountOptions,
     skip_change_addresses: bool,
+    change_addresses_to_sync: HashSet<AddressWrapper>,
 ) -> crate::Result<(Vec<Address>, Vec<SyncedMessage>)> {
     let mut messages = vec![];
     let client_options = account.client_options().clone();
@@ -404,7 +405,9 @@ async fn sync_messages(
 
     let mut tasks = Vec::new();
     for mut address in account.addresses().to_vec() {
-        if skip_addresses.contains(&address) || (*address.internal() && skip_change_addresses) {
+        if skip_addresses.contains(&address)
+            || (*address.internal() && skip_change_addresses && !change_addresses_to_sync.contains(address.address()))
+        {
             continue;
         }
         let client = client.clone();
@@ -500,6 +503,7 @@ async fn perform_sync(
     address_index: usize,
     gap_limit: usize,
     skip_change_addresses: bool,
+    change_addresses_to_sync: HashSet<AddressWrapper>,
     steps: &[AccountSynchronizeStep],
     options: AccountOptions,
 ) -> crate::Result<SyncedAccountData> {
@@ -562,8 +566,14 @@ async fn perform_sync(
     }
 
     if steps.contains(&AccountSynchronizeStep::SyncMessages) {
-        let (synced_addresses, synced_messages) =
-            sync_messages(&account, &found_addresses, options, skip_change_addresses).await?;
+        let (synced_addresses, synced_messages) = sync_messages(
+            &account,
+            &found_addresses,
+            options,
+            skip_change_addresses,
+            change_addresses_to_sync,
+        )
+        .await?;
         found_addresses.extend(synced_addresses);
         new_messages.extend(synced_messages.into_iter());
     }
@@ -733,11 +743,13 @@ impl AccountSynchronizer {
     }
 
     pub(crate) async fn get_new_history(&self) -> crate::Result<SyncedAccountData> {
+        let change_addresses_to_sync = self.account_handle.change_addresses_to_sync.lock().await.clone();
         perform_sync(
             self.account_handle.read().await.clone(),
             self.address_index,
             self.gap_limit,
             self.skip_change_addresses,
+            change_addresses_to_sync,
             &self.steps,
             self.account_handle.account_options,
         )
@@ -1450,6 +1462,11 @@ async fn perform_transfer(
                     .iter()
                     .find(|a| *a.internal() && a.key_index() == remainder_address.key_index())
                 {
+                    account_handle
+                        .change_addresses_to_sync
+                        .lock()
+                        .await
+                        .insert(address.address().clone());
                     address.address().clone()
                 } else {
                     transfer_obj
@@ -1465,6 +1482,11 @@ async fn perform_transfer(
                     )
                     .await?;
                     let addr = change_address.address().clone();
+                    account_handle
+                        .change_addresses_to_sync
+                        .lock()
+                        .await
+                        .insert(addr.clone());
                     log::debug!(
                         "[TRANSFER] generated new change address as remainder target: {}",
                         addr.to_bech32()

diff --git a/src/address.rs b/src/address.rs
@@ -217,7 +217,7 @@ impl AddressBuilder {
 }
 
 /// An address and its network type.
-#[derive(Debug, Clone, PartialEq, Eq)]
+#[derive(Debug, Clone, PartialEq, Eq, Hash)]
 pub struct AddressWrapper {
     inner: IotaAddress,
     pub(crate) bech32_hrp: String,