From 6a610a8cdfb0b7913278267f4b0df11aadaf1760 Mon Sep 17 00:00:00 2001
From: Roy Lee <roylee17@gmail.com>
Date: Thu, 25 Aug 2022 22:27:21 -0700
Subject: [PATCH] multi-account: update getaddressbyaccount

---
 rpc/legacyrpc/methods.go | 18 ++++++++++++++----
 wallet/createtx_test.go  |  6 +++---
 wallet/wallet.go         | 29 ++++++++++++++++++++++++-----
 3 files changed, 41 insertions(+), 12 deletions(-)

diff --git a/rpc/legacyrpc/methods.go b/rpc/legacyrpc/methods.go
index 51529ad..3f587f8 100644
--- a/rpc/legacyrpc/methods.go
+++ b/rpc/legacyrpc/methods.go
@@ -394,16 +394,25 @@ func dumpPrivKey(icmd interface{}, w *wallet.Wallet) (interface{}, error) {
 
 // getAddressesByAccount handles a getaddressesbyaccount request by returning
 // all addresses for an account, or an error if the requested account does
-// not exist.
-func getAddressesByAccount(icmd interface{}, w *wallet.Wallet) (interface{}, error) {
+// not exist. If addresstype is also specified, only those address types are
+// returned.
+func getAddressesByAccount(icmd interface{}, w *wallet.Wallet) (
+	interface{}, error) {
+
 	cmd := icmd.(*btcjson.GetAddressesByAccountCmd)
 
-	account, err := w.AccountNumber(waddrmgr.KeyScopeBIP0044, cmd.Account)
+	account, err := w.AccountNumber(*cmd.Account)
 	if err != nil {
 		return nil, err
 	}
 
-	addrs, err := w.AccountAddresses(account)
+	// Use specified scope, if provided.
+	scope, err := lookupKeyScope(cmd.AddressType)
+	if err != nil {
+		return nil, err
+	}
+
+	addrs, err := w.AccountAddresses(account, scope)
 	if err != nil {
 		return nil, err
 	}
@@ -412,6 +421,7 @@ func getAddressesByAccount(icmd interface{}, w *wallet.Wallet) (interface{}, err
 	for i, a := range addrs {
 		addrStrs[i] = a.EncodeAddress()
 	}
+
 	return addrStrs, nil
 }
 
diff --git a/wallet/createtx_test.go b/wallet/createtx_test.go
index 5d3b4f8..73bff24 100644
--- a/wallet/createtx_test.go
+++ b/wallet/createtx_test.go
@@ -83,7 +83,7 @@ func TestTxToOutputsDryRun(t *testing.T) {
 	}
 	change := dryRunTx.Tx.TxOut[dryRunTx.ChangeIndex]
 
-	addresses, err := w.AccountAddresses(0)
+	addresses, err := w.AccountAddresses(0, nil)
 	if err != nil {
 		t.Fatalf("unable to get addresses: %v", err)
 	}
@@ -100,7 +100,7 @@ func TestTxToOutputsDryRun(t *testing.T) {
 	}
 	change2 := dryRunTx2.Tx.TxOut[dryRunTx2.ChangeIndex]
 
-	addresses, err = w.AccountAddresses(0)
+	addresses, err = w.AccountAddresses(0, nil)
 	if err != nil {
 		t.Fatalf("unable to get addresses: %v", err)
 	}
@@ -135,7 +135,7 @@ func TestTxToOutputsDryRun(t *testing.T) {
 	}
 	change3 := tx.Tx.TxOut[tx.ChangeIndex]
 
-	addresses, err = w.AccountAddresses(0)
+	addresses, err = w.AccountAddresses(0, nil)
 	if err != nil {
 		t.Fatalf("unable to get addresses: %v", err)
 	}
diff --git a/wallet/wallet.go b/wallet/wallet.go
index f9ffef4..fed3eae 100644
--- a/wallet/wallet.go
+++ b/wallet/wallet.go
@@ -1407,13 +1407,32 @@ func (w *Wallet) ChangePassphrases(publicOld, publicNew, privateOld,
 
 // AccountAddresses returns the addresses for every created address for an
 // account.
-func (w *Wallet) AccountAddresses(account uint32) (addrs []btcutil.Address, err error) {
+func (w *Wallet) AccountAddresses(account uint32, scope *waddrmgr.KeyScope) (
+	addrs []btcutil.Address, err error) {
+
+	// By default, append all addresses under this account.
+	fn := func(maddr waddrmgr.ManagedAddress) error {
+		addrs = append(addrs, maddr.Address())
+		return nil
+	}
+
+	// If scope is set, append only those have the address type under
+	// this scope.
+	if scope != nil {
+		addrSchema := waddrmgr.ScopeAddrMap[*scope]
+
+		fn = func(maddr waddrmgr.ManagedAddress) error {
+			if maddr.AddrType() == addrSchema.InternalAddrType ||
+				maddr.AddrType() == addrSchema.ExternalAddrType {
+				addrs = append(addrs, maddr.Address())
+			}
+			return nil
+		}
+	}
+
 	err = walletdb.View(w.db, func(tx walletdb.ReadTx) error {
 		addrmgrNs := tx.ReadBucket(waddrmgrNamespaceKey)
-		return w.Manager.ForEachAccountAddress(addrmgrNs, account, func(maddr waddrmgr.ManagedAddress) error {
-			addrs = append(addrs, maddr.Address())
-			return nil
-		})
+		return w.Manager.ForEachAccountAddress(addrmgrNs, account, fn)
 	})
 	return
 }