From 8c1181af3b4665e69d147f267d0f03849cde4567 Mon Sep 17 00:00:00 2001
From: Olaoluwa Osuntokun <laolu32@gmail.com>
Date: Tue, 19 Feb 2019 19:16:39 -0800
Subject: [PATCH] lnwallet: add LastUnusedAddress to WalletController interface

In this commit, we add a new `LastUnusedAddress` method to the
`WalletController` interface. Callers can use this new method to graph
the last unused address, which can be useful for UIs that want to
refresh the address, but not cause nearly unbounded address generation.

The implementation for `btcwallet` uses the existing `CurrentAddress`
method. We've also added a new integration tests to exercise the new
functionality.
---
 lnwallet/btcwallet/btcwallet.go | 23 +++++++++++++
 lnwallet/interface.go           |  8 +++++
 lnwallet/interface_test.go      | 60 +++++++++++++++++++++++++++++++++
 3 files changed, 91 insertions(+)

diff --git a/lnwallet/btcwallet/btcwallet.go b/lnwallet/btcwallet/btcwallet.go
index 2449d968b..963ca002c 100644
--- a/lnwallet/btcwallet/btcwallet.go
+++ b/lnwallet/btcwallet/btcwallet.go
@@ -254,6 +254,29 @@ func (b *BtcWallet) NewAddress(t lnwallet.AddressType, change bool) (btcutil.Add
 	return b.wallet.NewAddress(defaultAccount, keyScope)
 }
 
+// LastUnusedAddress returns the last *unused* address known by the wallet. An
+// address is unused if it hasn't received any payments. This can be useful in
+// UIs in order to continually show the "freshest" address without having to
+// worry about "address inflation" caused by continual refreshing. Similar to
+// NewAddress it can derive a specified address type, and also optionally a
+// change address.
+func (b *BtcWallet) LastUnusedAddress(addrType lnwallet.AddressType) (
+	btcutil.Address, error) {
+
+	var keyScope waddrmgr.KeyScope
+
+	switch addrType {
+	case lnwallet.WitnessPubKey:
+		keyScope = waddrmgr.KeyScopeBIP0084
+	case lnwallet.NestedWitnessPubKey:
+		keyScope = waddrmgr.KeyScopeBIP0049Plus
+	default:
+		return nil, fmt.Errorf("unknown address type")
+	}
+
+	return b.wallet.CurrentAddress(defaultAccount, keyScope)
+}
+
 // IsOurAddress checks if the passed address belongs to this wallet
 //
 // This is a part of the WalletController interface.
diff --git a/lnwallet/interface.go b/lnwallet/interface.go
index 18c353be0..907bda58e 100644
--- a/lnwallet/interface.go
+++ b/lnwallet/interface.go
@@ -150,6 +150,14 @@ type WalletController interface {
 	// p2wsh, etc.
 	NewAddress(addrType AddressType, change bool) (btcutil.Address, error)
 
+	// LastUnusedAddress returns the last *unused* address known by the
+	// wallet. An address is unused if it hasn't received any payments.
+	// This can be useful in UIs in order to continually show the
+	// "freshest" address without having to worry about "address inflation"
+	// caused by continual refreshing. Similar to NewAddress it can derive
+	// a specified address type. By default, this is a non-change address.
+	LastUnusedAddress(addrType AddressType) (btcutil.Address, error)
+
 	// IsOurAddress checks if the passed address belongs to this wallet
 	IsOurAddress(a btcutil.Address) bool
 
diff --git a/lnwallet/interface_test.go b/lnwallet/interface_test.go
index 7d5ad9049..4f486670e 100644
--- a/lnwallet/interface_test.go
+++ b/lnwallet/interface_test.go
@@ -2165,6 +2165,62 @@ func testChangeOutputSpendConfirmation(r *rpctest.Harness,
 	}
 }
 
+// testLastUnusedAddr tests that the LastUnusedAddress returns the address if
+// it isn't used, and also that once the address becomes used, then it's
+// properly rotated.
+func testLastUnusedAddr(miner *rpctest.Harness,
+	alice, bob *lnwallet.LightningWallet, t *testing.T) {
+
+	if _, err := miner.Node.Generate(1); err != nil {
+		t.Fatalf("unable to generate block: %v", err)
+	}
+
+	// We'll repeat this test for each address type to ensure they're all
+	// rotated properly.
+	addrTypes := []lnwallet.AddressType{
+		lnwallet.WitnessPubKey, lnwallet.NestedWitnessPubKey,
+	}
+	for _, addrType := range addrTypes {
+		addr1, err := alice.LastUnusedAddress(addrType)
+		if err != nil {
+			t.Fatalf("unable to get addr: %v", err)
+		}
+		addr2, err := alice.LastUnusedAddress(addrType)
+		if err != nil {
+			t.Fatalf("unable to get addr: %v", err)
+		}
+
+		// If we generate two addresses back to back, then we should
+		// get the same addr, as none of them have been used yet.
+		if addr1.String() != addr2.String() {
+			t.Fatalf("addresses changed w/o use: %v vs %v", addr1, addr2)
+		}
+
+		// Next, we'll have Bob pay to Alice's new address. This should
+		// trigger address rotation at the backend wallet.
+		addrScript, err := txscript.PayToAddrScript(addr1)
+		if err != nil {
+			t.Fatalf("unable to convert addr to script: %v", err)
+		}
+		feeRate := lnwallet.SatPerKWeight(2500)
+		output := &wire.TxOut{
+			Value:    1000000,
+			PkScript: addrScript,
+		}
+		sendCoins(t, miner, bob, alice, output, feeRate)
+
+		// If we make a new address, then it should be brand new, as
+		// the prior address has been used.
+		addr3, err := alice.LastUnusedAddress(addrType)
+		if err != nil {
+			t.Fatalf("unable to get addr: %v", err)
+		}
+		if addr1.String() == addr3.String() {
+			t.Fatalf("address should have changed but didn't")
+		}
+	}
+}
+
 type walletTestCase struct {
 	name string
 	test func(miner *rpctest.Harness, alice, bob *lnwallet.LightningWallet,
@@ -2219,6 +2275,10 @@ var walletTests = []walletTestCase{
 		name: "test cancel non-existent reservation",
 		test: testCancelNonExistentReservation,
 	},
+	{
+		name: "last unused addr",
+		test: testLastUnusedAddr,
+	},
 	{
 		name: "reorg wallet balance",
 		test: testReorgWalletBalance,