lnwallet+channeldb: add new AuxLeafStore for dynamic aux leaves

In this commit, we add a new AuxLeafStore which can be used to dynamically fetch the latest aux leaves for a given state. This is useful for custom channel types that will store some extra information in the form of a custom blob, then will use that information to derive the new leaf tapscript leaves that may be attached to reach state.

contractcourt/chain_watcher.go

@@ -431,7 +431,7 @@ func (c *chainWatcher) handleUnknownLocalState(
 	}
 	remoteScript, _, err := lnwallet.CommitScriptToRemote(
 		c.cfg.chanState.ChanType, c.cfg.chanState.IsInitiator,
-		commitKeyRing.ToRemoteKey, leaseExpiry,
+		commitKeyRing.ToRemoteKey, leaseExpiry, input.NoneTapLeaf(),
 	)
 	if err != nil {
 		return false, err
@@ -444,6 +444,7 @@ func (c *chainWatcher) handleUnknownLocalState(
 		c.cfg.chanState.ChanType, c.cfg.chanState.IsInitiator,
 		commitKeyRing.ToLocalKey, commitKeyRing.RevocationKey,
 		uint32(c.cfg.chanState.LocalChanCfg.CsvDelay), leaseExpiry,
+		input.NoneTapLeaf(),
 	)
 	if err != nil {
 		return false, err

lnwallet/aux_leaf_store.go

@@ -0,0 +1,239 @@
+package lnwallet
+
+import (
+	"github.com/btcsuite/btcd/btcutil"
+	"github.com/btcsuite/btcd/chaincfg/chainhash"
+	"github.com/btcsuite/btcd/wire"
+	"github.com/lightningnetwork/lnd/channeldb"
+	"github.com/lightningnetwork/lnd/fn"
+	"github.com/lightningnetwork/lnd/input"
+	"github.com/lightningnetwork/lnd/lntypes"
+	"github.com/lightningnetwork/lnd/lnwire"
+	"github.com/lightningnetwork/lnd/tlv"
+)
+
+// CommitSortFunc is a function type alias for a function that sorts the
+// commitment transaction outputs. The second parameter is a list of CLTV
+// timeouts that must correspond to the number of transaction outputs, with the
+// value of 0 for non-HTLC outputs. The HTLC indexes are needed to have a
+// deterministic sort value for HTLCs that have the identical amount, CLTV
+// timeout and payment hash (e.g. multiple MPP shards of the same payment, where
+// the on-chain script would be identical).
+type CommitSortFunc func(tx *wire.MsgTx, cltvs []uint32,
+	indexes []input.HtlcIndex) error
+
+// DefaultCommitSort is the default commitment sort function that sorts the
+// commitment transaction inputs and outputs according to BIP69. The second
+// parameter is a list of CLTV timeouts that must correspond to the number of
+// transaction outputs, with the value of 0 for non-HTLC outputs. The third
+// parameter is unused for the default sort function.
+func DefaultCommitSort(tx *wire.MsgTx, cltvs []uint32,
+	_ []input.HtlcIndex) error {
+
+	InPlaceCommitSort(tx, cltvs)
+	return nil
+}
+
+// CommitAuxLeaves stores two potential auxiliary leaves for the remote and
+// local output that may be used to augment the final tapscript trees of the
+// commitment transaction.
+type CommitAuxLeaves struct {
+	// LocalAuxLeaf is the local party's auxiliary leaf.
+	LocalAuxLeaf input.AuxTapLeaf
+
+	// RemoteAuxLeaf is the remote party's auxiliary leaf.
+	RemoteAuxLeaf input.AuxTapLeaf
+
+	// OutgoingHTLCLeaves is the set of aux leaves for the outgoing HTLCs
+	// on this commitment transaction.
+	OutgoingHtlcLeaves input.HtlcAuxLeaves
+
+	// IncomingHTLCLeaves is the set of aux leaves for the incoming HTLCs
+	// on this commitment transaction.
+	IncomingHtlcLeaves input.HtlcAuxLeaves
+}
+
+// AuxChanState is a struct that holds certain fields of the
+// channeldb.OpenChannel struct that are used by the aux components. The data
+// is copied over to prevent accidental mutation of the original channel state.
+type AuxChanState struct {
+	// ChanType denotes which type of channel this is.
+	ChanType channeldb.ChannelType
+
+	// FundingOutpoint is the outpoint of the final funding transaction.
+	// This value uniquely and globally identifies the channel within the
+	// target blockchain as specified by the chain hash parameter.
+	FundingOutpoint wire.OutPoint
+
+	// ShortChannelID encodes the exact location in the chain in which the
+	// channel was initially confirmed. This includes: the block height,
+	// transaction index, and the output within the target transaction.
+	//
+	// If IsZeroConf(), then this will the "base" (very first) ALIAS scid
+	// and the confirmed SCID will be stored in ConfirmedScid.
+	ShortChannelID lnwire.ShortChannelID
+
+	// IsInitiator is a bool which indicates if we were the original
+	// initiator for the channel. This value may affect how higher levels
+	// negotiate fees, or close the channel.
+	IsInitiator bool
+
+	// Capacity is the total capacity of this channel.
+	Capacity btcutil.Amount
+
+	// LocalChanCfg is the channel configuration for the local node.
+	LocalChanCfg channeldb.ChannelConfig
+
+	// RemoteChanCfg is the channel configuration for the remote node.
+	RemoteChanCfg channeldb.ChannelConfig
+
+	// ThawHeight is the height when a frozen channel once again becomes a
+	// normal channel. If this is zero, then there're no restrictions on
+	// this channel. If the value is lower than 500,000, then it's
+	// interpreted as a relative height, or an absolute height otherwise.
+	ThawHeight uint32
+
+	// TapscriptRoot is an optional tapscript root used to derive the MuSig2
+	// funding output.
+	TapscriptRoot fn.Option[chainhash.Hash]
+
+	// CustomBlob is an optional blob that can be used to store information
+	// specific to a custom channel type. This information is only created
+	// at channel funding time, and after wards is to be considered
+	// immutable.
+	CustomBlob fn.Option[tlv.Blob]
+}
+
+// NewAuxChanState creates a new AuxChanState from the given channel state.
+func NewAuxChanState(chanState *channeldb.OpenChannel) AuxChanState {
+	return AuxChanState{
+		ChanType:        chanState.ChanType,
+		FundingOutpoint: chanState.FundingOutpoint,
+		ShortChannelID:  chanState.ShortChannelID,
+		IsInitiator:     chanState.IsInitiator,
+		Capacity:        chanState.Capacity,
+		LocalChanCfg:    chanState.LocalChanCfg,
+		RemoteChanCfg:   chanState.RemoteChanCfg,
+		ThawHeight:      chanState.ThawHeight,
+		TapscriptRoot:   chanState.TapscriptRoot,
+		CustomBlob:      chanState.CustomBlob,
+	}
+}
+
+// CommitDiffAuxInput is the input required to compute the diff of the auxiliary
+// leaves for a commitment transaction.
+type CommitDiffAuxInput struct {
+	// ChannelState is the static channel information of the channel this
+	// commitment transaction relates to.
+	ChannelState AuxChanState
+
+	// PrevBlob is the blob of the previous commitment transaction.
+	PrevBlob tlv.Blob
+
+	// UnfilteredView is the unfiltered, original HTLC view of the channel.
+	// Unfiltered in this context means that the view contains all HTLCs,
+	// including the canceled ones.
+	UnfilteredView *HtlcView
+
+	// WhoseCommit denotes whose commitment transaction we are computing the
+	// diff for.
+	WhoseCommit lntypes.ChannelParty
+
+	// OurBalance is the balance of the local party.
+	OurBalance lnwire.MilliSatoshi
+
+	// TheirBalance is the balance of the remote party.
+	TheirBalance lnwire.MilliSatoshi
+
+	// KeyRing is the key ring that can be used to derive keys for the
+	// commitment transaction.
+	KeyRing CommitmentKeyRing
+}
+
+// CommitDiffAuxResult is the result of computing the diff of the auxiliary
+// leaves for a commitment transaction.
+type CommitDiffAuxResult struct {
+	// AuxLeaves are the auxiliary leaves for the new commitment
+	// transaction.
+	AuxLeaves fn.Option[CommitAuxLeaves]
+
+	// CommitSortFunc is an optional function that sorts the commitment
+	// transaction inputs and outputs.
+	CommitSortFunc fn.Option[CommitSortFunc]
+}
+
+// AuxLeafStore is used to optionally fetch auxiliary tapscript leaves for the
+// commitment transaction given an opaque blob. This is also used to implement
+// a state transition function for the blobs to allow them to be refreshed with
+// each state.
+type AuxLeafStore interface {
+	// FetchLeavesFromView attempts to fetch the auxiliary leaves that
+	// correspond to the passed aux blob, and pending original (unfiltered)
+	// HTLC view.
+	FetchLeavesFromView(
+		in CommitDiffAuxInput) fn.Result[CommitDiffAuxResult]
+
+	// FetchLeavesFromCommit attempts to fetch the auxiliary leaves that
+	// correspond to the passed aux blob, and an existing channel
+	// commitment.
+	FetchLeavesFromCommit(chanState AuxChanState,
+		commit channeldb.ChannelCommitment,
+		keyRing CommitmentKeyRing) fn.Result[CommitDiffAuxResult]
+
+	// FetchLeavesFromRevocation attempts to fetch the auxiliary leaves
+	// from a channel revocation that stores balance + blob information.
+	FetchLeavesFromRevocation(
+		r *channeldb.RevocationLog) fn.Result[CommitDiffAuxResult]
+
+	// ApplyHtlcView serves as the state transition function for the custom
+	// channel's blob. Given the old blob, and an HTLC view, then a new
+	// blob should be returned that reflects the pending updates.
+	ApplyHtlcView(in CommitDiffAuxInput) fn.Result[fn.Option[tlv.Blob]]
+}
+
+// auxLeavesFromView is used to derive the set of commit aux leaves (if any),
+// that are needed to create a new commitment transaction using the original
+// (unfiltered) htlc view.
+func auxLeavesFromView(leafStore AuxLeafStore, chanState *channeldb.OpenChannel,
+	prevBlob fn.Option[tlv.Blob], originalView *HtlcView,
+	whoseCommit lntypes.ChannelParty, ourBalance,
+	theirBalance lnwire.MilliSatoshi,
+	keyRing CommitmentKeyRing) fn.Result[CommitDiffAuxResult] {
+
+	return fn.MapOptionZ(
+		prevBlob, func(blob tlv.Blob) fn.Result[CommitDiffAuxResult] {
+			return leafStore.FetchLeavesFromView(CommitDiffAuxInput{
+				ChannelState:   NewAuxChanState(chanState),
+				PrevBlob:       blob,
+				UnfilteredView: originalView,
+				WhoseCommit:    whoseCommit,
+				OurBalance:     ourBalance,
+				TheirBalance:   theirBalance,
+				KeyRing:        keyRing,
+			})
+		},
+	)
+}
+
+// updateAuxBlob is a helper function that attempts to update the aux blob
+// given the prior and current state information.
+func updateAuxBlob(leafStore AuxLeafStore, chanState *channeldb.OpenChannel,
+	prevBlob fn.Option[tlv.Blob], nextViewUnfiltered *HtlcView,
+	whoseCommit lntypes.ChannelParty, ourBalance,
+	theirBalance lnwire.MilliSatoshi,
+	keyRing CommitmentKeyRing) fn.Result[fn.Option[tlv.Blob]] {
+
+	return fn.MapOptionZ(
+		prevBlob, func(blob tlv.Blob) fn.Result[fn.Option[tlv.Blob]] {
+			return leafStore.ApplyHtlcView(CommitDiffAuxInput{
+				ChannelState:   NewAuxChanState(chanState),
+				PrevBlob:       blob,
+				UnfilteredView: nextViewUnfiltered,
+				WhoseCommit:    whoseCommit,
+				OurBalance:     ourBalance,
+				TheirBalance:   theirBalance,
+				KeyRing:        keyRing,
+			})
+		},
+	)
+}

lnwallet/channel.go

@@ -804,6 +804,10 @@ type LightningChannel struct {
 	// machine.
 	Signer input.Signer
 
+	// leafStore is used to retrieve extra tapscript leaves for special
+	// custom channel types.
+	leafStore fn.Option[AuxLeafStore]
+
 	// signDesc is the primary sign descriptor that is capable of signing
 	// the commitment transaction that spends the multi-sig output.
 	signDesc *input.SignDescriptor
@@ -879,6 +883,8 @@ type channelOpts struct {
 	localNonce  *musig2.Nonces
 	remoteNonce *musig2.Nonces
 
+	leafStore fn.Option[AuxLeafStore]
+
 	skipNonceInit bool
 }
 
@@ -909,6 +915,13 @@ func WithSkipNonceInit() ChannelOpt {
 	}
 }
 
+// WithLeafStore is used to specify a custom leaf store for the channel.
+func WithLeafStore(store AuxLeafStore) ChannelOpt {
+	return func(o *channelOpts) {
+		o.leafStore = fn.Some[AuxLeafStore](store)
+	}
+}
+
 // defaultChannelOpts returns the set of default options for a new channel.
 func defaultChannelOpts() *channelOpts {
 	return &channelOpts{}
@@ -950,13 +963,16 @@ func NewLightningChannel(signer input.Signer,
 	}
 
 	lc := &LightningChannel{
-		Signer:               signer,
+		Signer:            signer,
-		sigPool:              sigPool,
+		leafStore:         opts.leafStore,
-		currentHeight:        localCommit.CommitHeight,
+		sigPool:           sigPool,
-		remoteCommitChain:    newCommitmentChain(),
+		currentHeight:     localCommit.CommitHeight,
-		localCommitChain:     newCommitmentChain(),
+		remoteCommitChain: newCommitmentChain(),
-		channelState:         state,
+		localCommitChain:  newCommitmentChain(),
-		commitBuilder:        NewCommitmentBuilder(state),
+		channelState:      state,
+		commitBuilder: NewCommitmentBuilder(
+			state, opts.leafStore,
+		),
 		localUpdateLog:       localUpdateLog,
 		remoteUpdateLog:      remoteUpdateLog,
 		Capacity:             state.Capacity,
@@ -1988,12 +2004,14 @@ func NewBreachRetribution(chanState *channeldb.OpenChannel, stateNum uint64,
 		leaseExpiry = chanState.ThawHeight
 	}
 
+	// TODO(roasbeef): Actually fetch aux leaves (later commits in this PR).
+
 	// Since it is the remote breach we are reconstructing, the output
 	// going to us will be a to-remote script with our local params.
 	isRemoteInitiator := !chanState.IsInitiator
 	ourScript, ourDelay, err := CommitScriptToRemote(
 		chanState.ChanType, isRemoteInitiator, keyRing.ToRemoteKey,
-		leaseExpiry,
+		leaseExpiry, input.NoneTapLeaf(),
 	)
 	if err != nil {
 		return nil, err
@@ -2003,6 +2021,7 @@ func NewBreachRetribution(chanState *channeldb.OpenChannel, stateNum uint64,
 	theirScript, err := CommitScriptToSelf(
 		chanState.ChanType, isRemoteInitiator, keyRing.ToLocalKey,
 		keyRing.RevocationKey, theirDelay, leaseExpiry,
+		input.NoneTapLeaf(),
 	)
 	if err != nil {
 		return nil, err
@@ -2560,7 +2579,8 @@ func (lc *LightningChannel) fetchCommitmentView(
 	// Actually generate unsigned commitment transaction for this view.
 	commitTx, err := lc.commitBuilder.createUnsignedCommitmentTx(
 		ourBalance, theirBalance, whoseCommitChain, feePerKw,
-		nextHeight, filteredHTLCView, keyRing,
+		nextHeight, htlcView, filteredHTLCView, keyRing,
+		commitChain.tip(),
 	)
 	if err != nil {
 		return nil, err
@@ -2595,6 +2615,23 @@ func (lc *LightningChannel) fetchCommitmentView(
 			effFeeRate, spew.Sdump(commitTx))
 	}
 
+	// Given the custom blob of the past state, and this new HTLC view,
+	// we'll generate a new blob for the latest commitment.
+	newCommitBlob, err := fn.MapOptionZ(
+		lc.leafStore,
+		func(s AuxLeafStore) fn.Result[fn.Option[tlv.Blob]] {
+			return updateAuxBlob(
+				s, lc.channelState,
+				commitChain.tip().customBlob, htlcView,
+				whoseCommitChain, ourBalance, theirBalance,
+				*keyRing,
+			)
+		},
+	).Unpack()
+	if err != nil {
+		return nil, fmt.Errorf("unable to fetch aux leaves: %w", err)
+	}
+
 	// With the commitment view created, store the resulting balances and
 	// transaction with the other parameters for this height.
 	c := &commitment{
@@ -2610,6 +2647,7 @@ func (lc *LightningChannel) fetchCommitmentView(
 		feePerKw:          feePerKw,
 		dustLimit:         dustLimit,
 		whoseCommit:       whoseCommitChain,
+		customBlob:        newCommitBlob,
 	}
 
 	// In order to ensure _none_ of the HTLC's associated with this new
@@ -2703,6 +2741,7 @@ func (lc *LightningChannel) evaluateHTLCView(view *HtlcView, ourBalance,
 		if mutateState && entry.EntryType == Settle &&
 			whoseCommitChain.IsLocal() &&
 			entry.removeCommitHeightLocal == 0 {
+
 			lc.channelState.TotalMSatReceived += entry.Amount
 		}
 
@@ -5398,7 +5437,7 @@ func (lc *LightningChannel) ReceiveRevocation(revMsg *lnwire.RevokeAndAck) (
 	// before the change since the indexes are meant for the current,
 	// revoked remote commitment.
 	ourOutputIndex, theirOutputIndex, err := findOutputIndexesFromRemote(
-		revocation, lc.channelState,
+		revocation, lc.channelState, lc.leafStore,
 	)
 	if err != nil {
 		return nil, nil, nil, nil, err
@@ -6286,12 +6325,14 @@ func NewUnilateralCloseSummary(chanState *channeldb.OpenChannel, signer input.Si
 
 	commitTxBroadcast := commitSpend.SpendingTx
 
+	// TODO(roasbeef): Actually fetch aux leaves (later commits in this PR).
+
 	// Before we can generate the proper sign descriptor, we'll need to
 	// locate the output index of our non-delayed output on the commitment
 	// transaction.
 	selfScript, maturityDelay, err := CommitScriptToRemote(
 		chanState.ChanType, isRemoteInitiator, keyRing.ToRemoteKey,
-		leaseExpiry,
+		leaseExpiry, input.NoneTapLeaf(),
 	)
 	if err != nil {
 		return nil, fmt.Errorf("unable to create self commit "+
@@ -7236,6 +7277,8 @@ func NewLocalForceCloseSummary(chanState *channeldb.OpenChannel,
 		&chanState.LocalChanCfg, &chanState.RemoteChanCfg,
 	)
 
+	// TODO(roasbeef): Actually fetch aux leaves (later commits in this PR).
+
 	var leaseExpiry uint32
 	if chanState.ChanType.HasLeaseExpiration() {
 		leaseExpiry = chanState.ThawHeight
@@ -7243,6 +7286,7 @@ func NewLocalForceCloseSummary(chanState *channeldb.OpenChannel,
 	toLocalScript, err := CommitScriptToSelf(
 		chanState.ChanType, chanState.IsInitiator, keyRing.ToLocalKey,
 		keyRing.RevocationKey, csvTimeout, leaseExpiry,
+		input.NoneTapLeaf(),
 	)
 	if err != nil {
 		return nil, err

lnwallet/commitment.go

@@ -11,6 +11,7 @@ import (
 	"github.com/btcsuite/btcd/txscript"
 	"github.com/btcsuite/btcd/wire"
 	"github.com/lightningnetwork/lnd/channeldb"
+	"github.com/lightningnetwork/lnd/fn"
 	"github.com/lightningnetwork/lnd/input"
 	"github.com/lightningnetwork/lnd/lntypes"
 	"github.com/lightningnetwork/lnd/lnwallet/chainfee"
@@ -225,8 +226,8 @@ func (w *WitnessScriptDesc) WitnessScriptForPath(
 // party learns of the preimage to the revocation hash, then they can claim all
 // the settled funds in the channel, plus the unsettled funds.
 func CommitScriptToSelf(chanType channeldb.ChannelType, initiator bool,
-	selfKey, revokeKey *btcec.PublicKey, csvDelay,
+	selfKey, revokeKey *btcec.PublicKey, csvDelay, leaseExpiry uint32,
-	leaseExpiry uint32) (input.ScriptDescriptor, error) {
+	auxLeaf input.AuxTapLeaf) (input.ScriptDescriptor, error) {
 
 	switch {
 	// For taproot scripts, we'll need to make a slightly modified script
@@ -236,7 +237,7 @@ func CommitScriptToSelf(chanType channeldb.ChannelType, initiator bool,
 	// Our "redeem" script here is just the taproot witness program.
 	case chanType.IsTaproot():
 		return input.NewLocalCommitScriptTree(
-			csvDelay, selfKey, revokeKey, input.NoneTapLeaf(),
+			csvDelay, selfKey, revokeKey, auxLeaf,
 		)
 
 	// If we are the initiator of a leased channel, then we have an
@@ -290,8 +291,8 @@ func CommitScriptToSelf(chanType channeldb.ChannelType, initiator bool,
 // script for. The second return value is the CSV delay of the output script,
 // what must be satisfied in order to spend the output.
 func CommitScriptToRemote(chanType channeldb.ChannelType, initiator bool,
-	remoteKey *btcec.PublicKey,
+	remoteKey *btcec.PublicKey, leaseExpiry uint32,
-	leaseExpiry uint32) (input.ScriptDescriptor, uint32, error) {
+	auxLeaf input.AuxTapLeaf) (input.ScriptDescriptor, uint32, error) {
 
 	switch {
 	// If we are not the initiator of a leased channel, then the remote
@@ -320,7 +321,7 @@ func CommitScriptToRemote(chanType channeldb.ChannelType, initiator bool,
 	// with the sole tap leaf enforcing the 1 CSV delay.
 	case chanType.IsTaproot():
 		toRemoteScriptTree, err := input.NewRemoteCommitScriptTree(
-			remoteKey, input.NoneTapLeaf(),
+			remoteKey, auxLeaf,
 		)
 		if err != nil {
 			return nil, 0, err
@@ -612,7 +613,7 @@ func CommitScriptAnchors(chanType channeldb.ChannelType,
 // with, and abstracts the various ways of constructing commitment
 // transactions.
 type CommitmentBuilder struct {
-	// chanState is the underlying channels's state struct, used to
+	// chanState is the underlying channel's state struct, used to
 	// determine the type of channel we are dealing with, and relevant
 	// parameters.
 	chanState *channeldb.OpenChannel
@@ -620,18 +621,25 @@ type CommitmentBuilder struct {
 	// obfuscator is a 48-bit state hint that's used to obfuscate the
 	// current state number on the commitment transactions.
 	obfuscator [StateHintSize]byte
+
+	// auxLeafStore is an interface that allows us to fetch auxiliary
+	// tapscript leaves for the commitment output.
+	auxLeafStore fn.Option[AuxLeafStore]
 }
 
 // NewCommitmentBuilder creates a new CommitmentBuilder from chanState.
-func NewCommitmentBuilder(chanState *channeldb.OpenChannel) *CommitmentBuilder {
+func NewCommitmentBuilder(chanState *channeldb.OpenChannel,
+	leafStore fn.Option[AuxLeafStore]) *CommitmentBuilder {
+
 	// The anchor channel type MUST be tweakless.
 	if chanState.ChanType.HasAnchors() && !chanState.ChanType.IsTweakless() {
 		panic("invalid channel type combination")
 	}
 
 	return &CommitmentBuilder{
-		chanState:  chanState,
+		chanState:    chanState,
-		obfuscator: createStateHintObfuscator(chanState),
+		obfuscator:   createStateHintObfuscator(chanState),
+		auxLeafStore: leafStore,
 	}
 }
 
@@ -684,9 +692,9 @@ type unsignedCommitmentTx struct {
 // fees, but after anchor outputs.
 func (cb *CommitmentBuilder) createUnsignedCommitmentTx(ourBalance,
 	theirBalance lnwire.MilliSatoshi, whoseCommit lntypes.ChannelParty,
-	feePerKw chainfee.SatPerKWeight, height uint64,
+	feePerKw chainfee.SatPerKWeight, height uint64, originalHtlcView,
-	filteredHTLCView *HtlcView,
+	filteredHTLCView *HtlcView, keyRing *CommitmentKeyRing,
-	keyRing *CommitmentKeyRing) (*unsignedCommitmentTx, error) {
+	prevCommit *commitment) (*unsignedCommitmentTx, error) {
 
 	dustLimit := cb.chanState.LocalChanCfg.DustLimit
 	if whoseCommit.IsRemote() {
@@ -752,6 +760,23 @@ func (cb *CommitmentBuilder) createUnsignedCommitmentTx(ourBalance,
 		err      error
 	)
 
+	// Before we create the commitment transaction below, we'll try to see
+	// if there're any aux leaves that need to be a part of the tapscript
+	// tree. We'll only do this if we have a custom blob defined though.
+	auxResult, err := fn.MapOptionZ(
+		cb.auxLeafStore,
+		func(s AuxLeafStore) fn.Result[CommitDiffAuxResult] {
+			return auxLeavesFromView(
+				s, cb.chanState, prevCommit.customBlob,
+				originalHtlcView, whoseCommit, ourBalance,
+				theirBalance, *keyRing,
+			)
+		},
+	).Unpack()
+	if err != nil {
+		return nil, fmt.Errorf("unable to fetch aux leaves: %w", err)
+	}
+
 	// Depending on whether the transaction is ours or not, we call
 	// CreateCommitTx with parameters matching the perspective, to generate
 	// a new commitment transaction with all the latest unsettled/un-timed
@@ -766,6 +791,7 @@ func (cb *CommitmentBuilder) createUnsignedCommitmentTx(ourBalance,
 			&cb.chanState.LocalChanCfg, &cb.chanState.RemoteChanCfg,
 			ourBalance.ToSatoshis(), theirBalance.ToSatoshis(),
 			numHTLCs, cb.chanState.IsInitiator, leaseExpiry,
+			auxResult.AuxLeaves,
 		)
 	} else {
 		commitTx, err = CreateCommitTx(
@@ -773,6 +799,7 @@ func (cb *CommitmentBuilder) createUnsignedCommitmentTx(ourBalance,
 			&cb.chanState.RemoteChanCfg, &cb.chanState.LocalChanCfg,
 			theirBalance.ToSatoshis(), ourBalance.ToSatoshis(),
 			numHTLCs, !cb.chanState.IsInitiator, leaseExpiry,
+			auxResult.AuxLeaves,
 		)
 	}
 	if err != nil {
@@ -789,6 +816,7 @@ func (cb *CommitmentBuilder) createUnsignedCommitmentTx(ourBalance,
 	// commitment outputs and should correspond to zero values for the
 	// purposes of sorting.
 	cltvs := make([]uint32, len(commitTx.TxOut))
+	htlcIndexes := make([]input.HtlcIndex, len(commitTx.TxOut))
 	for _, htlc := range filteredHTLCView.OurUpdates {
 		if HtlcIsDust(
 			cb.chanState.ChanType, false, whoseCommit, feePerKw,
@@ -805,7 +833,11 @@ func (cb *CommitmentBuilder) createUnsignedCommitmentTx(ourBalance,
 		if err != nil {
 			return nil, err
 		}
-		cltvs = append(cltvs, htlc.Timeout) // nolint:makezero
+
+		// We want to add the CLTV and HTLC index to their respective
+		// slices, even if we already pre-allocated them.
+		cltvs = append(cltvs, htlc.Timeout)               //nolint
+		htlcIndexes = append(htlcIndexes, htlc.HtlcIndex) //nolint
 	}
 	for _, htlc := range filteredHTLCView.TheirUpdates {
 		if HtlcIsDust(
@@ -823,7 +855,11 @@ func (cb *CommitmentBuilder) createUnsignedCommitmentTx(ourBalance,
 		if err != nil {
 			return nil, err
 		}
-		cltvs = append(cltvs, htlc.Timeout) // nolint:makezero
+
+		// We want to add the CLTV and HTLC index to their respective
+		// slices, even if we already pre-allocated them.
+		cltvs = append(cltvs, htlc.Timeout)               //nolint
+		htlcIndexes = append(htlcIndexes, htlc.HtlcIndex) //nolint
 	}
 
 	// Set the state hint of the commitment transaction to facilitate
@@ -835,9 +871,16 @@ func (cb *CommitmentBuilder) createUnsignedCommitmentTx(ourBalance,
 	}
 
 	// Sort the transactions according to the agreed upon canonical
-	// ordering. This lets us skip sending the entire transaction over,
+	// ordering (which might be customized for custom channel types, but
-	// instead we'll just send signatures.
+	// deterministic and both parties will arrive at the same result). This
-	InPlaceCommitSort(commitTx, cltvs)
+	// lets us skip sending the entire transaction over, instead we'll just
+	// send signatures.
+	commitSort := auxResult.CommitSortFunc.UnwrapOr(DefaultCommitSort)
+	err = commitSort(commitTx, cltvs, htlcIndexes)
+	if err != nil {
+		return nil, fmt.Errorf("unable to sort commitment "+
+			"transaction: %w", err)
+	}
 
 	// Next, we'll ensure that we don't accidentally create a commitment
 	// transaction which would be invalid by consensus.
@@ -879,24 +922,33 @@ func CreateCommitTx(chanType channeldb.ChannelType,
 	fundingOutput wire.TxIn, keyRing *CommitmentKeyRing,
 	localChanCfg, remoteChanCfg *channeldb.ChannelConfig,
 	amountToLocal, amountToRemote btcutil.Amount,
-	numHTLCs int64, initiator bool, leaseExpiry uint32) (*wire.MsgTx, error) {
+	numHTLCs int64, initiator bool, leaseExpiry uint32,
+	auxLeaves fn.Option[CommitAuxLeaves]) (*wire.MsgTx, error) {
 
 	// First, we create the script for the delayed "pay-to-self" output.
 	// This output has 2 main redemption clauses: either we can redeem the
 	// output after a relative block delay, or the remote node can claim
 	// the funds with the revocation key if we broadcast a revoked
 	// commitment transaction.
+	localAuxLeaf := fn.MapOption(func(l CommitAuxLeaves) input.AuxTapLeaf {
+		return l.LocalAuxLeaf
+	})(auxLeaves)
 	toLocalScript, err := CommitScriptToSelf(
 		chanType, initiator, keyRing.ToLocalKey, keyRing.RevocationKey,
 		uint32(localChanCfg.CsvDelay), leaseExpiry,
+		fn.FlattenOption(localAuxLeaf),
 	)
 	if err != nil {
 		return nil, err
 	}
 
 	// Next, we create the script paying to the remote.
+	remoteAuxLeaf := fn.MapOption(func(l CommitAuxLeaves) input.AuxTapLeaf {
+		return l.RemoteAuxLeaf
+	})(auxLeaves)
 	toRemoteScript, _, err := CommitScriptToRemote(
 		chanType, initiator, keyRing.ToRemoteKey, leaseExpiry,
+		fn.FlattenOption(remoteAuxLeaf),
 	)
 	if err != nil {
 		return nil, err
@@ -1201,7 +1253,8 @@ func addHTLC(commitTx *wire.MsgTx, whoseCommit lntypes.ChannelParty,
 // output scripts and compares them against the outputs inside the commitment
 // to find the match.
 func findOutputIndexesFromRemote(revocationPreimage *chainhash.Hash,
-	chanState *channeldb.OpenChannel) (uint32, uint32, error) {
+	chanState *channeldb.OpenChannel,
+	leafStore fn.Option[AuxLeafStore]) (uint32, uint32, error) {
 
 	// Init the output indexes as empty.
 	ourIndex := uint32(channeldb.OutputIndexEmpty)
@@ -1231,26 +1284,51 @@ func findOutputIndexesFromRemote(revocationPreimage *chainhash.Hash,
 		leaseExpiry = chanState.ThawHeight
 	}
 
-	// Map the scripts from our PoV. When facing a local commitment, the to
+	// If we have a custom blob, then we'll attempt to fetch the aux leaves
-	// local output belongs to us and the to remote output belongs to them.
+	// for this state.
-	// When facing a remote commitment, the to local output belongs to them
+	auxResult, err := fn.MapOptionZ(
-	// and the to remote output belongs to us.
+		leafStore, func(a AuxLeafStore) fn.Result[CommitDiffAuxResult] {
+			return a.FetchLeavesFromCommit(
+				NewAuxChanState(chanState), chanCommit,
+				*keyRing,
+			)
+		},
+	).Unpack()
+	if err != nil {
+		return ourIndex, theirIndex, fmt.Errorf("unable to fetch aux "+
+			"leaves: %w", err)
+	}
 
-	// Compute the to local script. From our PoV, when facing a remote
+	// Map the scripts from our PoV. When facing a local commitment, the
-	// commitment, the to local output belongs to them.
+	// to_local output belongs to us and the to_remote output belongs to
+	// them. When facing a remote commitment, the to_local output belongs to
+	// them and the to_remote output belongs to us.
+
+	// Compute the to_local script. From our PoV, when facing a remote
+	// commitment, the to_local output belongs to them.
+	localAuxLeaf := fn.ChainOption(
+		func(l CommitAuxLeaves) input.AuxTapLeaf {
+			return l.LocalAuxLeaf
+		},
+	)(auxResult.AuxLeaves)
 	theirScript, err := CommitScriptToSelf(
 		chanState.ChanType, isRemoteInitiator, keyRing.ToLocalKey,
-		keyRing.RevocationKey, theirDelay, leaseExpiry,
+		keyRing.RevocationKey, theirDelay, leaseExpiry, localAuxLeaf,
 	)
 	if err != nil {
 		return ourIndex, theirIndex, err
 	}
 
-	// Compute the to remote script. From our PoV, when facing a remote
+	// Compute the to_remote script. From our PoV, when facing a remote
-	// commitment, the to remote output belongs to us.
+	// commitment, the to_remote output belongs to us.
+	remoteAuxLeaf := fn.ChainOption(
+		func(l CommitAuxLeaves) input.AuxTapLeaf {
+			return l.RemoteAuxLeaf
+		},
+	)(auxResult.AuxLeaves)
 	ourScript, _, err := CommitScriptToRemote(
 		chanState.ChanType, isRemoteInitiator, keyRing.ToRemoteKey,
-		leaseExpiry,
+		leaseExpiry, remoteAuxLeaf,
 	)
 	if err != nil {
 		return ourIndex, theirIndex, err

lnwallet/transactions_test.go

@@ -21,6 +21,7 @@ import (
 	"github.com/btcsuite/btcd/txscript"
 	"github.com/btcsuite/btcd/wire"
 	"github.com/lightningnetwork/lnd/channeldb"
+	"github.com/lightningnetwork/lnd/fn"
 	"github.com/lightningnetwork/lnd/input"
 	"github.com/lightningnetwork/lnd/keychain"
 	"github.com/lightningnetwork/lnd/lntypes"
@@ -631,6 +632,7 @@ func testSpendValidation(t *testing.T, tweakless bool) {
 	commitmentTx, err := CreateCommitTx(
 		channelType, *fakeFundingTxIn, keyRing, aliceChanCfg,
 		bobChanCfg, channelBalance, channelBalance, 0, true, 0,
+		fn.None[CommitAuxLeaves](),
 	)
 	if err != nil {
 		t.Fatalf("unable to create commitment transaction: %v", nil)

lnwallet/wallet.go

@@ -1470,6 +1470,21 @@ func (l *LightningWallet) handleFundingCancelRequest(req *fundingReserveCancelMs
 	req.err <- nil
 }
 
+// createCommitOpts is a struct that holds the options for creating a new
+// commitment transaction.
+type createCommitOpts struct {
+	auxLeaves fn.Option[CommitAuxLeaves]
+}
+
+// defaultCommitOpts returns a new createCommitOpts with default values.
+func defaultCommitOpts() createCommitOpts {
+	return createCommitOpts{}
+}
+
+// CreateCommitOpt is a functional option that can be used to modify the way a
+// new commitment transaction is created.
+type CreateCommitOpt func(*createCommitOpts)
+
 // CreateCommitmentTxns is a helper function that creates the initial
 // commitment transaction for both parties. This function is used during the
 // initial funding workflow as both sides must generate a signature for the
@@ -1479,7 +1494,13 @@ func CreateCommitmentTxns(localBalance, remoteBalance btcutil.Amount,
 	ourChanCfg, theirChanCfg *channeldb.ChannelConfig,
 	localCommitPoint, remoteCommitPoint *btcec.PublicKey,
 	fundingTxIn wire.TxIn, chanType channeldb.ChannelType, initiator bool,
-	leaseExpiry uint32) (*wire.MsgTx, *wire.MsgTx, error) {
+	leaseExpiry uint32, opts ...CreateCommitOpt) (*wire.MsgTx, *wire.MsgTx,
+	error) {
+
+	options := defaultCommitOpts()
+	for _, optFunc := range opts {
+		optFunc(&options)
+	}
 
 	localCommitmentKeys := DeriveCommitmentKeys(
 		localCommitPoint, lntypes.Local, chanType, ourChanCfg,
@@ -1493,7 +1514,7 @@ func CreateCommitmentTxns(localBalance, remoteBalance btcutil.Amount,
 	ourCommitTx, err := CreateCommitTx(
 		chanType, fundingTxIn, localCommitmentKeys, ourChanCfg,
 		theirChanCfg, localBalance, remoteBalance, 0, initiator,
-		leaseExpiry,
+		leaseExpiry, options.auxLeaves,
 	)
 	if err != nil {
 		return nil, nil, err
@@ -1507,7 +1528,7 @@ func CreateCommitmentTxns(localBalance, remoteBalance btcutil.Amount,
 	theirCommitTx, err := CreateCommitTx(
 		chanType, fundingTxIn, remoteCommitmentKeys, theirChanCfg,
 		ourChanCfg, remoteBalance, localBalance, 0, !initiator,
-		leaseExpiry,
+		leaseExpiry, options.auxLeaves,
 	)
 	if err != nil {
 		return nil, nil, err