From ad25ae1a0716a6b820a9d3951e4d8686c02788b8 Mon Sep 17 00:00:00 2001
From: Olaoluwa Osuntokun <laolu32@gmail.com>
Date: Mon, 13 Aug 2018 19:19:58 -0700
Subject: [PATCH] keychain: ensure we properly set the KeyLocator for keys from
 DeriveNextKey

In this commit, we fix a slight bug in the existing implementation of
DeriveNextKey for btcwallet. Before this commit, we would only set the
public key, and not also the derivation path. It's important that we
also set the path information, as in the near future we'll be using the
KeyDescriptors returned from this method to create static channel back
ups. With these static backups, the key alone may be insufficient to
re-derive the private key as we may need to fallback to brute forcing in
order to re-derive the key as it's possible we add new key families in
the future.
---
 fundingmanager.go      |  7 ++++---
 keychain/btcwallet.go  | 24 +++++++++++++++++++++---
 keychain/derivation.go |  6 +++---
 3 files changed, 28 insertions(+), 9 deletions(-)

diff --git a/fundingmanager.go b/fundingmanager.go
index 4df62b32..68550ab3 100644
--- a/fundingmanager.go
+++ b/fundingmanager.go
@@ -521,13 +521,13 @@ func (f *fundingManager) Start() error {
 				// mined since the channel was initiated reaches
 				// maxWaitNumBlocksFundingConf and we are not the channel
 				// initiator.
-
+				localBalance := ch.LocalCommitment.LocalBalance.ToSatoshis()
 				closeInfo := &channeldb.ChannelCloseSummary{
 					ChainHash:               ch.ChainHash,
 					ChanPoint:               ch.FundingOutpoint,
 					RemotePub:               ch.IdentityPub,
 					Capacity:                ch.Capacity,
-					SettledBalance:          ch.LocalBalance,
+					SettledBalance:          localBalance,
 					CloseType:               channeldb.FundingCanceled,
 					RemoteCurrentRevocation: ch.RemoteCurrentRevocation,
 					RemoteNextRevocation:    ch.RemoteNextRevocation,
@@ -1376,13 +1376,14 @@ func (f *fundingManager) handleFundingCreated(fmsg *fundingCreatedMsg) {
 	// we use this convenience method to delete the pending OpenChannel
 	// from the database.
 	deleteFromDatabase := func() {
+		localBalance := completeChan.LocalCommitment.LocalBalance.ToSatoshis()
 		closeInfo := &channeldb.ChannelCloseSummary{
 			ChanPoint:               completeChan.FundingOutpoint,
 			ChainHash:               completeChan.ChainHash,
 			RemotePub:               completeChan.IdentityPub,
 			CloseType:               channeldb.FundingCanceled,
 			Capacity:                completeChan.Capacity,
-			SettledBalance:          completeChan.LocalBalance,
+			SettledBalance:          localBalance,
 			RemoteCurrentRevocation: completeChan.RemoteCurrentRevocation,
 			RemoteNextRevocation:    completeChan.RemoteNextRevocation,
 			LocalChanConfig:         completeChan.LocalChanCfg,
diff --git a/keychain/btcwallet.go b/keychain/btcwallet.go
index e0f52203..e1feecfc 100644
--- a/keychain/btcwallet.go
+++ b/keychain/btcwallet.go
@@ -139,7 +139,10 @@ func (b *BtcWalletKeyRing) createAccountIfNotExists(
 //
 // NOTE: This is part of the keychain.KeyRing interface.
 func (b *BtcWalletKeyRing) DeriveNextKey(keyFam KeyFamily) (KeyDescriptor, error) {
-	var pubKey *btcec.PublicKey
+	var (
+		pubKey *btcec.PublicKey
+		keyLoc KeyLocator
+	)
 
 	db := b.wallet.Database()
 	err := walletdb.Update(db, func(tx walletdb.ReadWriteTx) error {
@@ -165,7 +168,21 @@ func (b *BtcWalletKeyRing) DeriveNextKey(keyFam KeyFamily) (KeyDescriptor, error
 			return err
 		}
 
-		pubKey = addrs[0].(waddrmgr.ManagedPubKeyAddress).PubKey()
+		// Extract the first address, ensuring that it is of the proper
+		// interface type, otherwise we can't manipulate it below.
+		addr, ok := addrs[0].(waddrmgr.ManagedPubKeyAddress)
+		if !ok {
+			return fmt.Errorf("address is not a managed pubkey " +
+				"addr")
+		}
+
+		pubKey = addr.PubKey()
+
+		_, pathInfo, _ := addr.DerivationInfo()
+		keyLoc = KeyLocator{
+			Family: keyFam,
+			Index:  pathInfo.Index,
+		}
 
 		return nil
 	})
@@ -174,7 +191,8 @@ func (b *BtcWalletKeyRing) DeriveNextKey(keyFam KeyFamily) (KeyDescriptor, error
 	}
 
 	return KeyDescriptor{
-		PubKey: pubKey,
+		PubKey:     pubKey,
+		KeyLocator: keyLoc,
 	}, nil
 }
 
diff --git a/keychain/derivation.go b/keychain/derivation.go
index 241e79d6..c0ad4355 100644
--- a/keychain/derivation.go
+++ b/keychain/derivation.go
@@ -83,7 +83,7 @@ const (
 // will vary per channel and use case) is the final element which allows us to
 // deterministically derive keys.
 type KeyLocator struct {
-	// TODO(roasbeef); add the key scope as well??
+	// TODO(roasbeef): add the key scope as well??
 
 	// Family is the family of key being identified.
 	Family KeyFamily
@@ -92,8 +92,8 @@ type KeyLocator struct {
 	Index uint32
 }
 
-// IsEmpty returns true if a KeyLocator is "empty". This may be the case where we
-// learn of a key from a remote party for a contract, but don't know the
+// IsEmpty returns true if a KeyLocator is "empty". This may be the case where
+// we learn of a key from a remote party for a contract, but don't know the
 // precise details of its derivation (as we don't know the private key!).
 func (k KeyLocator) IsEmpty() bool {
 	return k.Family == 0 && k.Index == 0