From 12850499231eaf6c38c165ef2e3365d0a8a2a157 Mon Sep 17 00:00:00 2001
From: Wilmer Paulino <wilmer.paulino@gmail.com>
Date: Mon, 30 Mar 2020 15:32:34 -0700
Subject: [PATCH] wallet: include addresses from relevant key scopes in rescan

Due to a no longer existing bug within the wallet, it was possible for
change addresses to be created outside of their intended key scope (the
default), so wallets affected by this now need to ensure they scan the
chain for all addresses within the default key scopes (as expected), and
all _internal_ addresses (branch used for change addresses) within any
other registered key scopes to reflect their proper balance.
---
 waddrmgr/manager.go        | 37 +++++++++++++++++++++++++++++++++++++
 waddrmgr/scoped_manager.go | 27 +++++++++++++++++++++++++++
 wallet/wallet.go           | 10 ++++++----
 3 files changed, 70 insertions(+), 4 deletions(-)

diff --git a/waddrmgr/manager.go b/waddrmgr/manager.go
index b94c6ee..8dc279a 100644
--- a/waddrmgr/manager.go
+++ b/waddrmgr/manager.go
@@ -738,6 +738,43 @@ func (m *Manager) ForEachActiveAddress(ns walletdb.ReadBucket, fn func(addr btcu
 	return nil
 }
 
+// ForEachRelevantActiveAddress invokes the given closure on each active
+// address relevant to the wallet. Ideally, only addresses within the default
+// key scopes would be relevant, but due to a bug (now fixed) in which change
+// addresses could be created outside of the default key scopes, we now need to
+// check for those as well.
+func (m *Manager) ForEachRelevantActiveAddress(ns walletdb.ReadBucket,
+	fn func(addr btcutil.Address) error) error {
+
+	m.mtx.RLock()
+	defer m.mtx.RUnlock()
+
+	for _, scopedMgr := range m.scopedManagers {
+		// If the manager is for a default key scope, we'll return all
+		// addresses, otherwise we'll only return internal addresses, as
+		// that's the branch used for change addresses.
+		isDefaultKeyScope := false
+		for _, defaultKeyScope := range DefaultKeyScopes {
+			if scopedMgr.Scope() == defaultKeyScope {
+				isDefaultKeyScope = true
+				break
+			}
+		}
+
+		var err error
+		if isDefaultKeyScope {
+			err = scopedMgr.ForEachActiveAddress(ns, fn)
+		} else {
+			err = scopedMgr.ForEachInternalActiveAddress(ns, fn)
+		}
+		if err != nil {
+			return err
+		}
+	}
+
+	return nil
+}
+
 // ForEachAccountAddress calls the given function with each address of
 // the given account stored in the manager, breaking early on error.
 func (m *Manager) ForEachAccountAddress(ns walletdb.ReadBucket, account uint32,
diff --git a/waddrmgr/scoped_manager.go b/waddrmgr/scoped_manager.go
index b9acda3..16bd83a 100644
--- a/waddrmgr/scoped_manager.go
+++ b/waddrmgr/scoped_manager.go
@@ -1760,3 +1760,30 @@ func (s *ScopedKeyManager) ForEachActiveAddress(ns walletdb.ReadBucket,
 
 	return nil
 }
+
+// ForEachInternalActiveAddress invokes the given closure on each _internal_
+// active address belonging to the scoped key manager, breaking early on error.
+func (s *ScopedKeyManager) ForEachInternalActiveAddress(ns walletdb.ReadBucket,
+	fn func(addr btcutil.Address) error) error {
+
+	s.mtx.Lock()
+	defer s.mtx.Unlock()
+
+	addrFn := func(rowInterface interface{}) error {
+		managedAddr, err := s.rowInterfaceToManaged(ns, rowInterface)
+		if err != nil {
+			return err
+		}
+		// Skip any non-internal branch addresses.
+		if !managedAddr.Internal() {
+			return nil
+		}
+		return fn(managedAddr.Address())
+	}
+
+	if err := forEachActiveAddress(ns, &s.scope, addrFn); err != nil {
+		return maybeConvertDbError(err)
+	}
+
+	return nil
+}
diff --git a/wallet/wallet.go b/wallet/wallet.go
index cdb1862..3a8167b 100644
--- a/wallet/wallet.go
+++ b/wallet/wallet.go
@@ -311,10 +311,12 @@ func (w *Wallet) activeData(dbtx walletdb.ReadTx) ([]btcutil.Address, []wtxmgr.C
 	txmgrNs := dbtx.ReadBucket(wtxmgrNamespaceKey)
 
 	var addrs []btcutil.Address
-	err := w.Manager.ForEachActiveAddress(addrmgrNs, func(addr btcutil.Address) error {
-		addrs = append(addrs, addr)
-		return nil
-	})
+	err := w.Manager.ForEachRelevantActiveAddress(
+		addrmgrNs, func(addr btcutil.Address) error {
+			addrs = append(addrs, addr)
+			return nil
+		},
+	)
 	if err != nil {
 		return nil, nil, err
 	}