From 87b4991bb62799ebb254d555af22396b2b593bbf Mon Sep 17 00:00:00 2001
From: Olaoluwa Osuntokun <laolu32@gmail.com>
Date: Thu, 7 Nov 2024 18:42:42 -0800
Subject: [PATCH] lnwallet: add whoseCommit to FetchLeavesFromCommit

This is useful for additional context to know which commit the
AuxLeafStore is fetching the leaves for.
---
 contractcourt/chain_watcher.go |  1 +
 lnwallet/aux_leaf_store.go     |  3 ++-
 lnwallet/channel.go            | 10 ++++++----
 lnwallet/commitment.go         |  2 +-
 lnwallet/mock.go               |  5 +++--
 5 files changed, 13 insertions(+), 8 deletions(-)

diff --git a/contractcourt/chain_watcher.go b/contractcourt/chain_watcher.go
index f25ca7e0a..64307dd02 100644
--- a/contractcourt/chain_watcher.go
+++ b/contractcourt/chain_watcher.go
@@ -436,6 +436,7 @@ func (c *chainWatcher) handleUnknownLocalState(
 			return s.FetchLeavesFromCommit(
 				lnwallet.NewAuxChanState(c.cfg.chanState),
 				c.cfg.chanState.LocalCommitment, *commitKeyRing,
+				lntypes.Local,
 			)
 		},
 	).Unpack()
diff --git a/lnwallet/aux_leaf_store.go b/lnwallet/aux_leaf_store.go
index 4558c2f81..c457a9250 100644
--- a/lnwallet/aux_leaf_store.go
+++ b/lnwallet/aux_leaf_store.go
@@ -178,7 +178,8 @@ type AuxLeafStore interface {
 	// commitment.
 	FetchLeavesFromCommit(chanState AuxChanState,
 		commit channeldb.ChannelCommitment,
-		keyRing CommitmentKeyRing) fn.Result[CommitDiffAuxResult]
+		keyRing CommitmentKeyRing, whoseCommit lntypes.ChannelParty,
+	) fn.Result[CommitDiffAuxResult]
 
 	// FetchLeavesFromRevocation attempts to fetch the auxiliary leaves
 	// from a channel revocation that stores balance + blob information.
diff --git a/lnwallet/channel.go b/lnwallet/channel.go
index 0f167bf6a..e7ec60772 100644
--- a/lnwallet/channel.go
+++ b/lnwallet/channel.go
@@ -666,6 +666,7 @@ func (lc *LightningChannel) diskCommitToMemCommit(
 			return s.FetchLeavesFromCommit(
 				NewAuxChanState(lc.channelState), *diskCommit,
 				*commitKeys.GetForParty(whoseCommit),
+				whoseCommit,
 			)
 		},
 	).Unpack()
@@ -1834,7 +1835,7 @@ func (lc *LightningChannel) restorePendingLocalUpdates(
 		func(s AuxLeafStore) fn.Result[CommitDiffAuxResult] {
 			return s.FetchLeavesFromCommit(
 				NewAuxChanState(lc.channelState), pendingCommit,
-				*pendingRemoteKeys,
+				*pendingRemoteKeys, lntypes.Remote,
 			)
 		},
 	).Unpack()
@@ -3154,7 +3155,7 @@ func genRemoteHtlcSigJobs(keyRing *CommitmentKeyRing,
 		leafStore, func(s AuxLeafStore) fn.Result[CommitDiffAuxResult] {
 			return s.FetchLeavesFromCommit(
 				NewAuxChanState(chanState), *diskCommit,
-				*keyRing,
+				*keyRing, lntypes.Remote,
 			)
 		},
 	).Unpack()
@@ -4740,7 +4741,7 @@ func genHtlcSigValidationJobs(chanState *channeldb.OpenChannel,
 		leafStore, func(s AuxLeafStore) fn.Result[CommitDiffAuxResult] {
 			return s.FetchLeavesFromCommit(
 				NewAuxChanState(chanState), *diskCommit,
-				*keyRing,
+				*keyRing, lntypes.Local,
 			)
 		},
 	).Unpack()
@@ -6727,7 +6728,7 @@ func NewUnilateralCloseSummary(chanState *channeldb.OpenChannel, //nolint:funlen
 		leafStore, func(s AuxLeafStore) fn.Result[CommitDiffAuxResult] {
 			return s.FetchLeavesFromCommit(
 				NewAuxChanState(chanState), remoteCommit,
-				*keyRing,
+				*keyRing, lntypes.Remote,
 			)
 		},
 	).Unpack()
@@ -7949,6 +7950,7 @@ func NewLocalForceCloseSummary(chanState *channeldb.OpenChannel,
 			return s.FetchLeavesFromCommit(
 				NewAuxChanState(chanState),
 				chanState.LocalCommitment, *keyRing,
+				lntypes.Local,
 			)
 		},
 	).Unpack()
diff --git a/lnwallet/commitment.go b/lnwallet/commitment.go
index 6d61729a4..36ff75ede 100644
--- a/lnwallet/commitment.go
+++ b/lnwallet/commitment.go
@@ -1307,7 +1307,7 @@ func findOutputIndexesFromRemote(revocationPreimage *chainhash.Hash,
 		leafStore, func(a AuxLeafStore) fn.Result[CommitDiffAuxResult] {
 			return a.FetchLeavesFromCommit(
 				NewAuxChanState(chanState), chanCommit,
-				*keyRing,
+				*keyRing, lntypes.Remote,
 			)
 		},
 	).Unpack()
diff --git a/lnwallet/mock.go b/lnwallet/mock.go
index d71b292a7..6623d8014 100644
--- a/lnwallet/mock.go
+++ b/lnwallet/mock.go
@@ -19,6 +19,7 @@ import (
 	"github.com/lightningnetwork/lnd/chainntnfs"
 	"github.com/lightningnetwork/lnd/channeldb"
 	"github.com/lightningnetwork/lnd/fn"
+	"github.com/lightningnetwork/lnd/lntypes"
 	"github.com/lightningnetwork/lnd/lnwallet/chainfee"
 	"github.com/lightningnetwork/lnd/tlv"
 	"github.com/stretchr/testify/mock"
@@ -420,8 +421,8 @@ func (*MockAuxLeafStore) FetchLeavesFromView(
 // correspond to the passed aux blob, and an existing channel
 // commitment.
 func (*MockAuxLeafStore) FetchLeavesFromCommit(_ AuxChanState,
-	_ channeldb.ChannelCommitment,
-	_ CommitmentKeyRing) fn.Result[CommitDiffAuxResult] {
+	_ channeldb.ChannelCommitment, _ CommitmentKeyRing,
+	_ lntypes.ChannelParty) fn.Result[CommitDiffAuxResult] {
 
 	return fn.Ok(CommitDiffAuxResult{})
 }