From 1ba717c74135942eef42ff77e9fafb51ff9bf2d9 Mon Sep 17 00:00:00 2001
From: Oliver Gugger <gugger@gmail.com>
Date: Tue, 14 Jun 2022 19:07:23 +0200
Subject: [PATCH] account+rpcserver: bump account key index on recovery

Fixes #372.
This commit fixes the issue that if the lnd node the trader client is
connected to was also restored from seed, it is starting at account key
derivation index 0. So when recovering accounts we need to make sure we
re-derive the right number of keys from the wallet in order to allow
properly creating new accounts with the recovered node.
---
 account/recovery.go | 56 +++++++++++++++++++++++++++++++++++++++++++++
 rpcserver.go        | 20 +++++++++++++++-
 2 files changed, 75 insertions(+), 1 deletion(-)

diff --git a/account/recovery.go b/account/recovery.go
index 49a60f0..f46e8ce 100644
--- a/account/recovery.go
+++ b/account/recovery.go
@@ -9,11 +9,13 @@ import (
 
 	"github.com/btcsuite/btcd/btcec/v2"
 	"github.com/btcsuite/btcd/btcutil"
+	"github.com/btcsuite/btcd/chaincfg"
 	"github.com/btcsuite/btcd/rpcclient"
 	"github.com/btcsuite/btcd/wire"
 	"github.com/lightninglabs/lndclient"
 	"github.com/lightninglabs/pool/poolscript"
 	"github.com/lightningnetwork/lnd/keychain"
+	"github.com/lightningnetwork/lnd/lnrpc/walletrpc"
 )
 
 var (
@@ -523,3 +525,57 @@ func GenerateRecoveryKeys(ctx context.Context, accountTarget uint32,
 	}
 	return acctKeys, nil
 }
+
+// AdvanceAccountDerivationIndex asserts that the internal wallet has at least
+// the given minimum index for deriving account keys.
+func AdvanceAccountDerivationIndex(ctx context.Context, minimumIndex uint32,
+	wallet lndclient.WalletKitClient, chainParams *chaincfg.Params) error {
+
+	allAccounts, err := wallet.ListAccounts(
+		ctx, "", walletrpc.AddressType_UNKNOWN,
+	)
+	if err != nil {
+		return fmt.Errorf("error listing accounts: %v", err)
+	}
+
+	poolAccountsPath := fmt.Sprintf("m/%d'/%d'/%d'",
+		keychain.BIP0043Purpose, chainParams.HDCoinType,
+		poolscript.AccountKeyFamily)
+
+	// Before we change anything in the wallet, let's check if maybe we
+	// already have enough keys?
+	for _, acct := range allAccounts {
+		if acct.DerivationPath == poolAccountsPath {
+			// If we did find our account in the list and already
+			// have enough keys derived, we don't need to do
+			// anything here. This is the key _count_, so it must be
+			// strictly greater than the minimum _index_ that we
+			// want.
+			if acct.ExternalKeyCount > minimumIndex {
+				log.Debugf("Account %s already has %d "+
+					"external keys (want minimum index "+
+					"%d), not deriving any keys",
+					poolAccountsPath, acct.ExternalKeyCount,
+					minimumIndex)
+
+				return nil
+			}
+
+			break
+		}
+	}
+
+	for {
+		key, err := wallet.DeriveNextKey(
+			ctx, int32(poolscript.AccountKeyFamily),
+		)
+		if err != nil {
+			return fmt.Errorf("error deriving next key: %v", err)
+		}
+
+		log.Debugf("Derived next account key with index %d", key.Index)
+		if key.Index >= minimumIndex {
+			return nil
+		}
+	}
+}
diff --git a/rpcserver.go b/rpcserver.go
index bc1b6aa..9e29fa5 100644
--- a/rpcserver.go
+++ b/rpcserver.go
@@ -1185,16 +1185,34 @@ func (s *rpcServer) RecoverAccounts(ctx context.Context,
 	// nice since it allows us to try recovery multiple times until it
 	// actually works.
 	numRecovered := len(recoveredAccounts)
+	var maxIndex uint32
 	for _, acct := range recoveredAccounts {
+		// We need to know the highest index we ever used so we can make
+		// sure lnd's wallet is also at that index and the next account
+		// will be derived from the proper index.
+		if acct.TraderKey.Index > maxIndex {
+			maxIndex = acct.TraderKey.Index
+		}
+
 		err = s.accountManager.RecoverAccount(ctx, acct)
 		if err != nil {
 			// If something goes wrong for one account we still want
 			// to continue with the others.
 			numRecovered--
-			rpcLog.Errorf("error storing recovered account: %v", err)
+			rpcLog.Errorf("Error storing recovered account: %v", err)
 		}
 	}
 
+	// Try to ratchet forward lnd's derivation index for accounts.
+	err = account.AdvanceAccountDerivationIndex(
+		ctx, maxIndex, s.lndServices.WalletKit,
+		s.lndServices.ChainParams,
+	)
+	if err != nil {
+		rpcLog.Errorf("Error advancing lnd's wallet to index %d: %v",
+			maxIndex, err)
+	}
+
 	return &poolrpc.RecoverAccountsResponse{
 		NumRecoveredAccounts: uint32(numRecovered),
 	}, nil