lnwallet: thread thru input.AuxTapleaf to all relevant areas

In this commit, we start to thread thru the new aux tap leaf structures to all relevant areas. This includes: commitment outputs, resolution creation, breach handling, and also HTLC scripts.

contractcourt/chain_watcher.go

@@ -424,15 +424,30 @@ func (c *chainWatcher) handleUnknownLocalState(
 		&c.cfg.chanState.LocalChanCfg, &c.cfg.chanState.RemoteChanCfg,
 	)
 
+	auxLeaves, err := lnwallet.AuxLeavesFromCommit(
+		c.cfg.chanState, c.cfg.chanState.LocalCommitment,
+		c.cfg.auxLeafStore, *commitKeyRing,
+	)
+	if err != nil {
+		return false, fmt.Errorf("unable to fetch aux leaves: %w", err)
+	}
+
 	// With the keys derived, we'll construct the remote script that'll be
 	// present if they have a non-dust balance on the commitment.
 	var leaseExpiry uint32
 	if c.cfg.chanState.ChanType.HasLeaseExpiration() {
 		leaseExpiry = c.cfg.chanState.ThawHeight
 	}
+
+	remoteAuxLeaf := fn.MapOption(
+		func(l lnwallet.CommitAuxLeaves) input.AuxTapLeaf {
+			return l.RemoteAuxLeaf
+		},
+	)(auxLeaves)
 	remoteScript, _, err := lnwallet.CommitScriptToRemote(
 		c.cfg.chanState.ChanType, c.cfg.chanState.IsInitiator,
-		commitKeyRing.ToRemoteKey, leaseExpiry, input.NoneTapLeaf(),
+		commitKeyRing.ToRemoteKey, leaseExpiry,
+		fn.FlattenOption(remoteAuxLeaf),
 	)
 	if err != nil {
 		return false, err
@@ -441,11 +456,16 @@ func (c *chainWatcher) handleUnknownLocalState(
 	// Next, we'll derive our script that includes the revocation base for
 	// the remote party allowing them to claim this output before the CSV
 	// delay if we breach.
+	localAuxLeaf := fn.MapOption(
+		func(l lnwallet.CommitAuxLeaves) input.AuxTapLeaf {
+			return l.LocalAuxLeaf
+		},
+	)(auxLeaves)
 	localScript, err := lnwallet.CommitScriptToSelf(
 		c.cfg.chanState.ChanType, c.cfg.chanState.IsInitiator,
 		commitKeyRing.ToLocalKey, commitKeyRing.RevocationKey,
 		uint32(c.cfg.chanState.LocalChanCfg.CsvDelay), leaseExpiry,
-		input.NoneTapLeaf(),
+		fn.FlattenOption(localAuxLeaf),
 	)
 	if err != nil {
 		return false, err

input/size_test.go

@@ -1386,7 +1386,7 @@ func genTimeoutTx(t *testing.T,
 	// Create the unsigned timeout tx.
 	timeoutTx, err := lnwallet.CreateHtlcTimeoutTx(
 		chanType, false, testOutPoint, testAmt, testCLTVExpiry,
-		testCSVDelay, 0, testPubkey, testPubkey,
+		testCSVDelay, 0, testPubkey, testPubkey, input.NoneTapLeaf(),
 	)
 	require.NoError(t, err)
 
@@ -1455,7 +1455,7 @@ func genSuccessTx(t *testing.T, chanType channeldb.ChannelType) *wire.MsgTx {
 	// Create the unsigned success tx.
 	successTx, err := lnwallet.CreateHtlcSuccessTx(
 		chanType, false, testOutPoint, testAmt, testCSVDelay, 0,
-		testPubkey, testPubkey,
+		testPubkey, testPubkey, input.NoneTapLeaf(),
 	)
 	require.NoError(t, err)
 

lnwallet/channel.go

@@ -810,8 +810,8 @@ func (c *commitment) toDiskCommit(ourCommit bool) *channeldb.ChannelCommitment {
 // restart a channel session.
 func (lc *LightningChannel) diskHtlcToPayDesc(feeRate chainfee.SatPerKWeight,
 	commitHeight uint64, htlc *channeldb.HTLC, localCommitKeys,
-	remoteCommitKeys *CommitmentKeyRing, isLocal bool) (PaymentDescriptor,
+	remoteCommitKeys *CommitmentKeyRing, isLocal bool,
-	error) {
+	auxLeaf input.AuxTapLeaf) (PaymentDescriptor, error) {
 
 	// The proper pkScripts for this PaymentDescriptor must be
 	// generated so we can easily locate them within the commitment
@@ -835,7 +835,7 @@ func (lc *LightningChannel) diskHtlcToPayDesc(feeRate chainfee.SatPerKWeight,
 	if !isDustLocal && localCommitKeys != nil {
 		scriptInfo, err := genHtlcScript(
 			chanType, htlc.Incoming, true, htlc.RefundTimeout,
-			htlc.RHash, localCommitKeys,
+			htlc.RHash, localCommitKeys, auxLeaf,
 		)
 		if err != nil {
 			return pd, err
@@ -850,7 +850,7 @@ func (lc *LightningChannel) diskHtlcToPayDesc(feeRate chainfee.SatPerKWeight,
 	if !isDustRemote && remoteCommitKeys != nil {
 		scriptInfo, err := genHtlcScript(
 			chanType, htlc.Incoming, false, htlc.RefundTimeout,
-			htlc.RHash, remoteCommitKeys,
+			htlc.RHash, remoteCommitKeys, auxLeaf,
 		)
 		if err != nil {
 			return pd, err
@@ -907,7 +907,8 @@ func (lc *LightningChannel) diskHtlcToPayDesc(feeRate chainfee.SatPerKWeight,
 // for each side.
 func (lc *LightningChannel) extractPayDescs(commitHeight uint64,
 	feeRate chainfee.SatPerKWeight, htlcs []channeldb.HTLC, localCommitKeys,
-	remoteCommitKeys *CommitmentKeyRing, isLocal bool) ([]PaymentDescriptor,
+	remoteCommitKeys *CommitmentKeyRing, isLocal bool,
+	auxLeaves fn.Option[CommitAuxLeaves]) ([]PaymentDescriptor,
 	[]PaymentDescriptor, error) {
 
 	var (
@@ -925,10 +926,21 @@ func (lc *LightningChannel) extractPayDescs(commitHeight uint64,
 
 		htlc := htlc
 
+		auxLeaf := fn.MapOption(
+			func(l CommitAuxLeaves) input.AuxTapLeaf {
+				leaves := l.OutgoingHtlcLeaves
+				if htlc.Incoming {
+					leaves = l.IncomingHtlcLeaves
+				}
+
+				return leaves[htlc.HtlcIndex].AuxTapLeaf
+			},
+		)(auxLeaves)
+
 		payDesc, err := lc.diskHtlcToPayDesc(
 			feeRate, commitHeight, &htlc,
 			localCommitKeys, remoteCommitKeys,
-			isLocal,
+			isLocal, fn.FlattenOption(auxLeaf),
 		)
 		if err != nil {
 			return incomingHtlcs, outgoingHtlcs, err
@@ -972,14 +984,28 @@ func (lc *LightningChannel) diskCommitToMemCommit(isLocal bool,
 		)
 	}
 
+	auxLeaves, err := AuxLeavesFromCommit(
+		lc.channelState, *diskCommit, lc.leafStore,
+		func() CommitmentKeyRing {
+			if isLocal {
+				return *localCommitKeys
+			}
+
+			return *remoteCommitKeys
+		}(),
+	)
+	if err != nil {
+		return nil, fmt.Errorf("unable to fetch aux leaves: %w", err)
+	}
+
 	// With the key rings re-created, we'll now convert all the on-disk
 	// HTLC"s into PaymentDescriptor's so we can re-insert them into our
 	// update log.
 	incomingHtlcs, outgoingHtlcs, err := lc.extractPayDescs(
 		diskCommit.CommitHeight,
 		chainfee.SatPerKWeight(diskCommit.FeePerKw),
-		diskCommit.Htlcs, localCommitKeys, remoteCommitKeys,
+		diskCommit.Htlcs, localCommitKeys, remoteCommitKeys, isLocal,
-		isLocal,
+		auxLeaves,
 	)
 	if err != nil {
 		return nil, err
@@ -1583,7 +1609,8 @@ func (lc *LightningChannel) ResetState() {
 func (lc *LightningChannel) logUpdateToPayDesc(logUpdate *channeldb.LogUpdate,
 	remoteUpdateLog *updateLog, commitHeight uint64,
 	feeRate chainfee.SatPerKWeight, remoteCommitKeys *CommitmentKeyRing,
-	remoteDustLimit btcutil.Amount) (*PaymentDescriptor, error) {
+	remoteDustLimit btcutil.Amount,
+	auxLeaves fn.Option[CommitAuxLeaves]) (*PaymentDescriptor, error) {
 
 	// Depending on the type of update message we'll map that to a distinct
 	// PaymentDescriptor instance.
@@ -1619,10 +1646,17 @@ func (lc *LightningChannel) logUpdateToPayDesc(logUpdate *channeldb.LogUpdate,
 			wireMsg.Amount.ToSatoshis(), remoteDustLimit,
 		)
 		if !isDustRemote {
+			auxLeaf := fn.MapOption(
+				func(l CommitAuxLeaves) input.AuxTapLeaf {
+					leaves := l.OutgoingHtlcLeaves
+					return leaves[pd.HtlcIndex].AuxTapLeaf
+				},
+			)(auxLeaves)
+
 			scriptInfo, err := genHtlcScript(
 				lc.channelState.ChanType, false, false,
 				wireMsg.Expiry, wireMsg.PaymentHash,
-				remoteCommitKeys,
+				remoteCommitKeys, fn.FlattenOption(auxLeaf),
 			)
 			if err != nil {
 				return nil, err
@@ -2290,6 +2324,14 @@ func (lc *LightningChannel) restorePendingLocalUpdates(
 	pendingCommit := pendingRemoteCommitDiff.Commitment
 	pendingHeight := pendingCommit.CommitHeight
 
+	auxLeaves, err := AuxLeavesFromCommit(
+		lc.channelState, pendingCommit, lc.leafStore,
+		*pendingRemoteKeys,
+	)
+	if err != nil {
+		return fmt.Errorf("unable to fetch aux leaves: %w", err)
+	}
+
 	// If we did have a dangling commit, then we'll examine which updates
 	// we included in that state and re-insert them into our update log.
 	for _, logUpdate := range pendingRemoteCommitDiff.LogUpdates {
@@ -2299,7 +2341,7 @@ func (lc *LightningChannel) restorePendingLocalUpdates(
 			&logUpdate, lc.remoteUpdateLog, pendingHeight,
 			chainfee.SatPerKWeight(pendingCommit.FeePerKw),
 			pendingRemoteKeys,
-			lc.channelState.RemoteChanCfg.DustLimit,
+			lc.channelState.RemoteChanCfg.DustLimit, auxLeaves,
 		)
 		if err != nil {
 			return err
@@ -2506,24 +2548,35 @@ func NewBreachRetribution(chanState *channeldb.OpenChannel, stateNum uint64,
 		leaseExpiry = chanState.ThawHeight
 	}
 
-	// TODO(roasbeef): fetch aux leave
+	auxLeaves, err := auxLeavesFromRevocation(
+		chanState, revokedLog, leafStore, *keyRing,
+	)
+	if err != nil {
+		return nil, fmt.Errorf("unable to fetch aux leaves: %w", err)
+	}
 
 	// Since it is the remote breach we are reconstructing, the output
 	// going to us will be a to-remote script with our local params.
+	localAuxLeaf := fn.MapOption(func(l CommitAuxLeaves) input.AuxTapLeaf {
+		return l.LocalAuxLeaf
+	})(auxLeaves)
 	isRemoteInitiator := !chanState.IsInitiator
 	ourScript, ourDelay, err := CommitScriptToRemote(
 		chanState.ChanType, isRemoteInitiator, keyRing.ToRemoteKey,
-		leaseExpiry, input.NoneTapLeaf(),
+		leaseExpiry, fn.FlattenOption(localAuxLeaf),
 	)
 	if err != nil {
 		return nil, err
 	}
 
+	remoteAuxLeaf := fn.MapOption(func(l CommitAuxLeaves) input.AuxTapLeaf {
+		return l.RemoteAuxLeaf
+	})(auxLeaves)
 	theirDelay := uint32(chanState.RemoteChanCfg.CsvDelay)
 	theirScript, err := CommitScriptToSelf(
 		chanState.ChanType, isRemoteInitiator, keyRing.ToLocalKey,
 		keyRing.RevocationKey, theirDelay, leaseExpiry,
-		input.NoneTapLeaf(),
+		fn.FlattenOption(remoteAuxLeaf),
 	)
 	if err != nil {
 		return nil, err
@@ -2541,7 +2594,7 @@ func NewBreachRetribution(chanState *channeldb.OpenChannel, stateNum uint64,
 	if revokedLog != nil {
 		br, ourAmt, theirAmt, err = createBreachRetribution(
 			revokedLog, spendTx, chanState, keyRing,
-			commitmentSecret, leaseExpiry,
+			commitmentSecret, leaseExpiry, auxLeaves,
 		)
 		if err != nil {
 			return nil, err
@@ -2675,7 +2728,8 @@ func NewBreachRetribution(chanState *channeldb.OpenChannel, stateNum uint64,
 func createHtlcRetribution(chanState *channeldb.OpenChannel,
 	keyRing *CommitmentKeyRing, commitHash chainhash.Hash,
 	commitmentSecret *btcec.PrivateKey, leaseExpiry uint32,
-	htlc *channeldb.HTLCEntry) (HtlcRetribution, error) {
+	htlc *channeldb.HTLCEntry,
+	auxLeaves fn.Option[CommitAuxLeaves]) (HtlcRetribution, error) {
 
 	var emptyRetribution HtlcRetribution
 
@@ -2685,10 +2739,21 @@ func createHtlcRetribution(chanState *channeldb.OpenChannel,
 	// We'll generate the original second level witness script now, as
 	// we'll need it if we're revoking an HTLC output on the remote
 	// commitment transaction, and *they* go to the second level.
+	secondLevelAuxLeaf := fn.MapOption(
+		func(l CommitAuxLeaves) input.AuxTapLeaf {
+			idx := input.HtlcIndex(htlc.HtlcIndex.Val)
+
+			if htlc.Incoming.Val {
+				return l.IncomingHtlcLeaves[idx].SecondLevelLeaf
+			}
+
+			return l.OutgoingHtlcLeaves[idx].SecondLevelLeaf
+		},
+	)(auxLeaves)
 	secondLevelScript, err := SecondLevelHtlcScript(
 		chanState.ChanType, isRemoteInitiator,
 		keyRing.RevocationKey, keyRing.ToLocalKey, theirDelay,
-		leaseExpiry,
+		leaseExpiry, fn.FlattenOption(secondLevelAuxLeaf),
 	)
 	if err != nil {
 		return emptyRetribution, err
@@ -2699,9 +2764,19 @@ func createHtlcRetribution(chanState *channeldb.OpenChannel,
 	// HTLC script. Otherwise, is this was an outgoing HTLC that we sent,
 	// then from the PoV of the remote commitment state, they're the
 	// receiver of this HTLC.
+	htlcLeaf := fn.MapOption(func(l CommitAuxLeaves) input.AuxTapLeaf {
+		idx := input.HtlcIndex(htlc.HtlcIndex.Val)
+
+		if htlc.Incoming.Val {
+			return l.IncomingHtlcLeaves[idx].AuxTapLeaf
+		}
+
+		return l.OutgoingHtlcLeaves[idx].AuxTapLeaf
+	})(auxLeaves)
 	scriptInfo, err := genHtlcScript(
 		chanState.ChanType, htlc.Incoming.Val, false,
 		htlc.RefundTimeout.Val, htlc.RHash.Val, keyRing,
+		fn.FlattenOption(htlcLeaf),
 	)
 	if err != nil {
 		return emptyRetribution, err
@@ -2765,7 +2840,9 @@ func createHtlcRetribution(chanState *channeldb.OpenChannel,
 func createBreachRetribution(revokedLog *channeldb.RevocationLog,
 	spendTx *wire.MsgTx, chanState *channeldb.OpenChannel,
 	keyRing *CommitmentKeyRing, commitmentSecret *btcec.PrivateKey,
-	leaseExpiry uint32) (*BreachRetribution, int64, int64, error) {
+	leaseExpiry uint32,
+	auxLeaves fn.Option[CommitAuxLeaves]) (*BreachRetribution, int64, int64,
+	error) {
 
 	commitHash := revokedLog.CommitTxHash
 
@@ -2774,7 +2851,7 @@ func createBreachRetribution(revokedLog *channeldb.RevocationLog,
 	for i, htlc := range revokedLog.HTLCEntries {
 		hr, err := createHtlcRetribution(
 			chanState, keyRing, commitHash.Val,
-			commitmentSecret, leaseExpiry, htlc,
+			commitmentSecret, leaseExpiry, htlc, auxLeaves,
 		)
 		if err != nil {
 			return nil, 0, 0, err
@@ -2922,6 +2999,7 @@ func createBreachRetributionLegacy(revokedLog *channeldb.ChannelCommitment,
 		hr, err := createHtlcRetribution(
 			chanState, keyRing, commitHash,
 			commitmentSecret, leaseExpiry, entry,
+			fn.None[CommitAuxLeaves](),
 		)
 		if err != nil {
 			return nil, 0, 0, err
@@ -3400,6 +3478,9 @@ func processAddEntry(htlc *PaymentDescriptor, ourBalance, theirBalance *lnwire.M
 		*ourBalance -= htlc.Amount
 	}
 
+	// TODO(roasbef): also have it modify balances here
+	//  * obtain for HTLC as well?
+
 	if mutateState {
 		*addHeight = nextHeight
 	}
@@ -3520,6 +3601,15 @@ func genRemoteHtlcSigJobs(keyRing *CommitmentKeyRing,
 	var err error
 	cancelChan := make(chan struct{})
 
+	auxLeaves, err := AuxLeavesFromCommit(
+		chanState, *remoteCommitView.toDiskCommit(false), leafStore,
+		*keyRing,
+	)
+	if err != nil {
+		return nil, nil, fmt.Errorf("unable to fetch aux leaves: "+
+			"%w", err)
+	}
+
 	// For each outgoing and incoming HTLC, if the HTLC isn't considered a
 	// dust output after taking into account second-level HTLC fees, then a
 	// sigJob will be generated and appended to the current batch.
@@ -3546,6 +3636,14 @@ func genRemoteHtlcSigJobs(keyRing *CommitmentKeyRing,
 		htlcFee := HtlcTimeoutFee(chanType, feePerKw)
 		outputAmt := htlc.Amount.ToSatoshis() - htlcFee
 
+		auxLeaf := fn.MapOption(func(
+			l CommitAuxLeaves) input.AuxTapLeaf {
+
+			leaves := l.IncomingHtlcLeaves
+			return leaves[htlc.HtlcIndex].SecondLevelLeaf
+		},
+		)(auxLeaves)
+
 		// With the fee calculate, we can properly create the HTLC
 		// timeout transaction using the HTLC amount minus the fee.
 		op := wire.OutPoint{
@@ -3556,11 +3654,14 @@ func genRemoteHtlcSigJobs(keyRing *CommitmentKeyRing,
 			chanType, isRemoteInitiator, op, outputAmt,
 			htlc.Timeout, uint32(remoteChanCfg.CsvDelay),
 			leaseExpiry, keyRing.RevocationKey, keyRing.ToLocalKey,
+			fn.FlattenOption(auxLeaf),
 		)
 		if err != nil {
 			return nil, nil, err
 		}
 
+		// TODO(roasbeef): hook up signer interface here
+
 		// Construct a full hash cache as we may be signing a segwit v1
 		// sighash.
 		txOut := remoteCommitView.txn.TxOut[htlc.remoteOutputIndex]
@@ -3587,7 +3688,8 @@ func genRemoteHtlcSigJobs(keyRing *CommitmentKeyRing,
 		// If this is a taproot channel, then we'll need to set the
 		// method type to ensure we generate a valid signature.
 		if chanType.IsTaproot() {
-			sigJob.SignDesc.SignMethod = input.TaprootScriptSpendSignMethod //nolint:lll
+			//nolint:lll
+			sigJob.SignDesc.SignMethod = input.TaprootScriptSpendSignMethod
 		}
 
 		sigBatch = append(sigBatch, sigJob)
@@ -3613,6 +3715,14 @@ func genRemoteHtlcSigJobs(keyRing *CommitmentKeyRing,
 		htlcFee := HtlcSuccessFee(chanType, feePerKw)
 		outputAmt := htlc.Amount.ToSatoshis() - htlcFee
 
+		auxLeaf := fn.MapOption(func(
+			l CommitAuxLeaves) input.AuxTapLeaf {
+
+			leaves := l.OutgoingHtlcLeaves
+			return leaves[htlc.HtlcIndex].SecondLevelLeaf
+		},
+		)(auxLeaves)
+
 		// With the proper output amount calculated, we can now
 		// generate the success transaction using the remote party's
 		// CSV delay.
@@ -3624,6 +3734,7 @@ func genRemoteHtlcSigJobs(keyRing *CommitmentKeyRing,
 			chanType, isRemoteInitiator, op, outputAmt,
 			uint32(remoteChanCfg.CsvDelay), leaseExpiry,
 			keyRing.RevocationKey, keyRing.ToLocalKey,
+			fn.FlattenOption(auxLeaf),
 		)
 		if err != nil {
 			return nil, nil, err
@@ -4871,6 +4982,9 @@ func (lc *LightningChannel) computeView(view *HtlcView, remoteChain bool,
 	// rate.
 	view.FeePerKw = commitChain.tip().feePerKw
 
+	// TODO(roasbeef): also need to pass blob here as well for final
+	// balances?
+
 	// We evaluate the view at this stage, meaning settled and failed HTLCs
 	// will remove their corresponding added HTLCs.  The resulting filtered
 	// view will only have Add entries left, making it easy to compare the
@@ -4959,6 +5073,15 @@ func genHtlcSigValidationJobs(chanState *channeldb.OpenChannel,
 		len(localCommitmentView.outgoingHTLCs))
 	verifyJobs := make([]VerifyJob, 0, numHtlcs)
 
+	auxLeaves, err := AuxLeavesFromCommit(
+		chanState, *localCommitmentView.toDiskCommit(true), leafStore,
+		*keyRing,
+	)
+	if err != nil {
+		return nil, fmt.Errorf("unable to fetch aux leaves: %w",
+			err)
+	}
+
 	// We'll iterate through each output in the commitment transaction,
 	// populating the sigHash closure function if it's detected to be an
 	// HLTC output. Given the sighash, and the signing key, we'll be able
@@ -4992,11 +5115,20 @@ func genHtlcSigValidationJobs(chanState *channeldb.OpenChannel,
 				htlcFee := HtlcSuccessFee(chanType, feePerKw)
 				outputAmt := htlc.Amount.ToSatoshis() - htlcFee
 
+				auxLeaf := fn.MapOption(func(
+					l CommitAuxLeaves) input.AuxTapLeaf {
+
+					leaves := l.IncomingHtlcLeaves
+					idx := htlc.HtlcIndex
+					return leaves[idx].SecondLevelLeaf
+				})(auxLeaves)
+
 				successTx, err := CreateHtlcSuccessTx(
 					chanType, isLocalInitiator, op,
 					outputAmt, uint32(localChanCfg.CsvDelay),
 					leaseExpiry, keyRing.RevocationKey,
 					keyRing.ToLocalKey,
+					fn.FlattenOption(auxLeaf),
 				)
 				if err != nil {
 					return nil, err
@@ -5076,12 +5208,21 @@ func genHtlcSigValidationJobs(chanState *channeldb.OpenChannel,
 				htlcFee := HtlcTimeoutFee(chanType, feePerKw)
 				outputAmt := htlc.Amount.ToSatoshis() - htlcFee
 
+				auxLeaf := fn.MapOption(func(
+					l CommitAuxLeaves) input.AuxTapLeaf {
+
+					leaves := l.OutgoingHtlcLeaves
+					idx := htlc.HtlcIndex
+					return leaves[idx].SecondLevelLeaf
+				})(auxLeaves)
+
 				timeoutTx, err := CreateHtlcTimeoutTx(
 					chanType, isLocalInitiator, op,
 					outputAmt, htlc.Timeout,
 					uint32(localChanCfg.CsvDelay),
 					leaseExpiry, keyRing.RevocationKey,
 					keyRing.ToLocalKey,
+					fn.FlattenOption(auxLeaf),
 				)
 				if err != nil {
 					return nil, err
@@ -6799,6 +6940,13 @@ func NewUnilateralCloseSummary(chanState *channeldb.OpenChannel,
 		&chanState.LocalChanCfg, &chanState.RemoteChanCfg,
 	)
 
+	auxLeaves, err := AuxLeavesFromCommit(
+		chanState, remoteCommit, leafStore, *keyRing,
+	)
+	if err != nil {
+		return nil, fmt.Errorf("unable to fetch aux leaves: %w", err)
+	}
+
 	// Next, we'll obtain HTLC resolutions for all the outgoing HTLC's we
 	// had on their commitment transaction.
 	var leaseExpiry uint32
@@ -6810,7 +6958,7 @@ func NewUnilateralCloseSummary(chanState *channeldb.OpenChannel,
 		chainfee.SatPerKWeight(remoteCommit.FeePerKw), isOurCommit,
 		signer, remoteCommit.Htlcs, keyRing, &chanState.LocalChanCfg,
 		&chanState.RemoteChanCfg, commitSpend.SpendingTx,
-		chanState.ChanType, isRemoteInitiator, leaseExpiry,
+		chanState.ChanType, isRemoteInitiator, leaseExpiry, auxLeaves,
 	)
 	if err != nil {
 		return nil, fmt.Errorf("unable to create htlc "+
@@ -6819,14 +6967,17 @@ func NewUnilateralCloseSummary(chanState *channeldb.OpenChannel,
 
 	commitTxBroadcast := commitSpend.SpendingTx
 
-	// TODO(roasbeef): fetch aux leave
-
 	// 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.
+	//
+	// TODO(roasbeef): helper func to hide flatten
+	remoteAuxLeaf := fn.MapOption(func(l CommitAuxLeaves) input.AuxTapLeaf {
+		return l.RemoteAuxLeaf
+	})(auxLeaves)
 	selfScript, maturityDelay, err := CommitScriptToRemote(
 		chanState.ChanType, isRemoteInitiator, keyRing.ToRemoteKey,
-		leaseExpiry, input.NoneTapLeaf(),
+		leaseExpiry, fn.FlattenOption(remoteAuxLeaf),
 	)
 	if err != nil {
 		return nil, fmt.Errorf("unable to create self commit "+
@@ -7065,8 +7216,8 @@ func newOutgoingHtlcResolution(signer input.Signer,
 	localChanCfg *channeldb.ChannelConfig, commitTx *wire.MsgTx,
 	htlc *channeldb.HTLC, keyRing *CommitmentKeyRing,
 	feePerKw chainfee.SatPerKWeight, csvDelay, leaseExpiry uint32,
-	localCommit, isCommitFromInitiator bool,
+	localCommit, isCommitFromInitiator bool, chanType channeldb.ChannelType,
-	chanType channeldb.ChannelType) (*OutgoingHtlcResolution, error) {
+	auxLeaves fn.Option[CommitAuxLeaves]) (*OutgoingHtlcResolution, error) {
 
 	op := wire.OutPoint{
 		Hash:  commitTx.TxHash(),
@@ -7075,9 +7226,12 @@ func newOutgoingHtlcResolution(signer input.Signer,
 
 	// First, we'll re-generate the script used to send the HTLC to the
 	// remote party within their commitment transaction.
+	auxLeaf := fn.MapOption(func(l CommitAuxLeaves) input.AuxTapLeaf {
+		return l.OutgoingHtlcLeaves[htlc.HtlcIndex].AuxTapLeaf
+	})(auxLeaves)
 	htlcScriptInfo, err := genHtlcScript(
 		chanType, false, localCommit, htlc.RefundTimeout, htlc.RHash,
-		keyRing,
+		keyRing, fn.FlattenOption(auxLeaf),
 	)
 	if err != nil {
 		return nil, err
@@ -7150,10 +7304,17 @@ func newOutgoingHtlcResolution(signer input.Signer,
 
 	// With the fee calculated, re-construct the second level timeout
 	// transaction.
+	secondLevelAuxLeaf := fn.MapOption(
+		func(l CommitAuxLeaves) input.AuxTapLeaf {
+			leaves := l.OutgoingHtlcLeaves
+			return leaves[htlc.HtlcIndex].SecondLevelLeaf
+		},
+	)(auxLeaves)
 	timeoutTx, err := CreateHtlcTimeoutTx(
 		chanType, isCommitFromInitiator, op, secondLevelOutputAmt,
-		htlc.RefundTimeout, csvDelay, leaseExpiry, keyRing.RevocationKey,
+		htlc.RefundTimeout, csvDelay, leaseExpiry,
-		keyRing.ToLocalKey,
+		keyRing.RevocationKey, keyRing.ToLocalKey,
+		fn.FlattenOption(secondLevelAuxLeaf),
 	)
 	if err != nil {
 		return nil, err
@@ -7236,6 +7397,7 @@ func newOutgoingHtlcResolution(signer input.Signer,
 		htlcSweepScript, err = SecondLevelHtlcScript(
 			chanType, isCommitFromInitiator, keyRing.RevocationKey,
 			keyRing.ToLocalKey, csvDelay, leaseExpiry,
+			fn.FlattenOption(secondLevelAuxLeaf),
 		)
 		if err != nil {
 			return nil, err
@@ -7244,7 +7406,7 @@ func newOutgoingHtlcResolution(signer input.Signer,
 		//nolint:lll
 		secondLevelScriptTree, err := input.TaprootSecondLevelScriptTree(
 			keyRing.RevocationKey, keyRing.ToLocalKey, csvDelay,
-			fn.None[txscript.TapLeaf](),
+			fn.FlattenOption(secondLevelAuxLeaf),
 		)
 		if err != nil {
 			return nil, err
@@ -7318,8 +7480,8 @@ func newIncomingHtlcResolution(signer input.Signer,
 	localChanCfg *channeldb.ChannelConfig, commitTx *wire.MsgTx,
 	htlc *channeldb.HTLC, keyRing *CommitmentKeyRing,
 	feePerKw chainfee.SatPerKWeight, csvDelay, leaseExpiry uint32,
-	localCommit, isCommitFromInitiator bool, chanType channeldb.ChannelType) (
+	localCommit, isCommitFromInitiator bool, chanType channeldb.ChannelType,
-	*IncomingHtlcResolution, error) {
+	auxLeaves fn.Option[CommitAuxLeaves]) (*IncomingHtlcResolution, error) {
 
 	op := wire.OutPoint{
 		Hash:  commitTx.TxHash(),
@@ -7328,9 +7490,12 @@ func newIncomingHtlcResolution(signer input.Signer,
 
 	// First, we'll re-generate the script the remote party used to
 	// send the HTLC to us in their commitment transaction.
+	auxLeaf := fn.MapOption(func(l CommitAuxLeaves) input.AuxTapLeaf {
+		return l.IncomingHtlcLeaves[htlc.HtlcIndex].AuxTapLeaf
+	})(auxLeaves)
 	scriptInfo, err := genHtlcScript(
 		chanType, true, localCommit, htlc.RefundTimeout, htlc.RHash,
-		keyRing,
+		keyRing, fn.FlattenOption(auxLeaf),
 	)
 	if err != nil {
 		return nil, err
@@ -7390,6 +7555,13 @@ func newIncomingHtlcResolution(signer input.Signer,
 		}, nil
 	}
 
+	secondLevelAuxLeaf := fn.MapOption(
+		func(l CommitAuxLeaves) input.AuxTapLeaf {
+			leaves := l.IncomingHtlcLeaves
+			return leaves[htlc.HtlcIndex].SecondLevelLeaf
+		},
+	)(auxLeaves)
+
 	// Otherwise, we'll need to go to the second level to sweep this HTLC.
 	//
 	// First, we'll reconstruct the original HTLC success transaction,
@@ -7399,7 +7571,7 @@ func newIncomingHtlcResolution(signer input.Signer,
 	successTx, err := CreateHtlcSuccessTx(
 		chanType, isCommitFromInitiator, op, secondLevelOutputAmt,
 		csvDelay, leaseExpiry, keyRing.RevocationKey,
-		keyRing.ToLocalKey,
+		keyRing.ToLocalKey, fn.FlattenOption(secondLevelAuxLeaf),
 	)
 	if err != nil {
 		return nil, err
@@ -7482,6 +7654,7 @@ func newIncomingHtlcResolution(signer input.Signer,
 		htlcSweepScript, err = SecondLevelHtlcScript(
 			chanType, isCommitFromInitiator, keyRing.RevocationKey,
 			keyRing.ToLocalKey, csvDelay, leaseExpiry,
+			fn.FlattenOption(secondLevelAuxLeaf),
 		)
 		if err != nil {
 			return nil, err
@@ -7490,7 +7663,7 @@ func newIncomingHtlcResolution(signer input.Signer,
 		//nolint:lll
 		secondLevelScriptTree, err := input.TaprootSecondLevelScriptTree(
 			keyRing.RevocationKey, keyRing.ToLocalKey, csvDelay,
-			fn.None[txscript.TapLeaf](),
+			fn.FlattenOption(secondLevelAuxLeaf),
 		)
 		if err != nil {
 			return nil, err
@@ -7582,7 +7755,8 @@ func extractHtlcResolutions(feePerKw chainfee.SatPerKWeight, ourCommit bool,
 	signer input.Signer, htlcs []channeldb.HTLC, keyRing *CommitmentKeyRing,
 	localChanCfg, remoteChanCfg *channeldb.ChannelConfig,
 	commitTx *wire.MsgTx, chanType channeldb.ChannelType,
-	isCommitFromInitiator bool, leaseExpiry uint32) (*HtlcResolutions, error) {
+	isCommitFromInitiator bool, leaseExpiry uint32,
+	auxLeaves fn.Option[CommitAuxLeaves]) (*HtlcResolutions, error) {
 
 	// TODO(roasbeef): don't need to swap csv delay?
 	dustLimit := remoteChanCfg.DustLimit
@@ -7617,6 +7791,7 @@ func extractHtlcResolutions(feePerKw chainfee.SatPerKWeight, ourCommit bool,
 				signer, localChanCfg, commitTx, &htlc,
 				keyRing, feePerKw, uint32(csvDelay), leaseExpiry,
 				ourCommit, isCommitFromInitiator, chanType,
+				auxLeaves,
 			)
 			if err != nil {
 				return nil, fmt.Errorf("incoming resolution "+
@@ -7630,7 +7805,7 @@ func extractHtlcResolutions(feePerKw chainfee.SatPerKWeight, ourCommit bool,
 		ohr, err := newOutgoingHtlcResolution(
 			signer, localChanCfg, commitTx, &htlc, keyRing,
 			feePerKw, uint32(csvDelay), leaseExpiry, ourCommit,
-			isCommitFromInitiator, chanType,
+			isCommitFromInitiator, chanType, auxLeaves,
 		)
 		if err != nil {
 			return nil, fmt.Errorf("outgoing resolution "+
@@ -7865,11 +8040,19 @@ func NewLocalForceCloseSummary(chanState *channeldb.OpenChannel,
 	// recovery there is not much we can do with HTLCs, so we'll always
 	// use what we have in our latest state when extracting resolutions.
 	localCommit := chanState.LocalCommitment
+
+	auxLeaves, err := AuxLeavesFromCommit(
+		chanState, localCommit, leafStore, *keyRing,
+	)
+	if err != nil {
+		return nil, fmt.Errorf("unable to fetch aux leaves: %w", err)
+	}
+
 	htlcResolutions, err := extractHtlcResolutions(
 		chainfee.SatPerKWeight(localCommit.FeePerKw), true, signer,
 		localCommit.Htlcs, keyRing, &chanState.LocalChanCfg,
 		&chanState.RemoteChanCfg, commitTx, chanState.ChanType,
-		chanState.IsInitiator, leaseExpiry,
+		chanState.IsInitiator, leaseExpiry, auxLeaves,
 	)
 	if err != nil {
 		return nil, fmt.Errorf("unable to gen htlc resolution: %w", err)

lnwallet/channel_test.go

@@ -9965,7 +9965,7 @@ func TestCreateHtlcRetribution(t *testing.T) {
 	// Create the htlc retribution.
 	hr, err := createHtlcRetribution(
 		aliceChannel.channelState, keyRing, commitHash,
-		dummyPrivate, leaseExpiry, htlc,
+		dummyPrivate, leaseExpiry, htlc, fn.None[CommitAuxLeaves](),
 	)
 	// Expect no error.
 	require.NoError(t, err)
@@ -10171,6 +10171,7 @@ func TestCreateBreachRetribution(t *testing.T) {
 				tc.revocationLog, tx,
 				aliceChannel.channelState, keyRing,
 				dummyPrivate, leaseExpiry,
+				fn.None[CommitAuxLeaves](),
 			)
 
 			// Check the error if expected.
@@ -10410,7 +10411,7 @@ func TestExtractPayDescs(t *testing.T) {
 	// NOTE: we use nil commitment key rings to avoid checking the htlc
 	// scripts(`genHtlcScript`) as it should be tested independently.
 	incomingPDs, outgoingPDs, err := lnChan.extractPayDescs(
-		0, 0, htlcs, nil, nil, true,
+		0, 0, htlcs, nil, nil, true, fn.None[CommitAuxLeaves](),
 	)
 	require.NoError(t, err)
 

lnwallet/commitment.go

@@ -420,14 +420,14 @@ func sweepSigHash(chanType channeldb.ChannelType) txscript.SigHashType {
 // argument should correspond to the owner of the commitment transaction which
 // we are generating the to_local script for.
 func SecondLevelHtlcScript(chanType channeldb.ChannelType, initiator bool,
-	revocationKey, delayKey *btcec.PublicKey,
+	revocationKey, delayKey *btcec.PublicKey, csvDelay, leaseExpiry uint32,
-	csvDelay, leaseExpiry uint32) (input.ScriptDescriptor, error) {
+	auxLeaf input.AuxTapLeaf) (input.ScriptDescriptor, error) {
 
 	switch {
 	// For taproot channels, the pkScript is a segwit v1 p2tr output.
 	case chanType.IsTaproot():
 		return input.TaprootSecondLevelScriptTree(
-			revocationKey, delayKey, csvDelay, input.NoneTapLeaf(),
+			revocationKey, delayKey, csvDelay, auxLeaf,
 		)
 
 	// If we are the initiator of a leased channel, then we have an
@@ -902,10 +902,7 @@ func (cb *CommitmentBuilder) createUnsignedCommitmentTx(ourBalance,
 		theirBalance -= commitFeeMSat
 	}
 
-	var (
+	var commitTx *wire.MsgTx
-		commitTx *wire.MsgTx
-		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
@@ -947,6 +944,19 @@ func (cb *CommitmentBuilder) createUnsignedCommitmentTx(ourBalance,
 		return nil, err
 	}
 
+	// Similarly, we'll now attempt to extract the set of aux leaves for
+	// the set of incoming and outgoing HTLCs.
+	incomingAuxLeaves := fn.MapOption(
+		func(leaves CommitAuxLeaves) input.AuxTapLeaves {
+			return leaves.IncomingHtlcLeaves
+		},
+	)(auxLeaves)
+	outgoingAuxLeaves := fn.MapOption(
+		func(leaves CommitAuxLeaves) input.AuxTapLeaves {
+			return leaves.OutgoingHtlcLeaves
+		},
+	)(auxLeaves)
+
 	// We'll now add all the HTLC outputs to the commitment transaction.
 	// Each output includes an off-chain 2-of-2 covenant clause, so we'll
 	// need the objective local/remote keys for this particular commitment
@@ -966,9 +976,15 @@ func (cb *CommitmentBuilder) createUnsignedCommitmentTx(ourBalance,
 			continue
 		}
 
+		auxLeaf := fn.MapOption(
+			func(leaves input.AuxTapLeaves) input.AuxTapLeaf {
+				return leaves[htlc.HtlcIndex].AuxTapLeaf
+			},
+		)(outgoingAuxLeaves)
+
 		err := addHTLC(
 			commitTx, isOurs, false, htlc, keyRing,
-			cb.chanState.ChanType,
+			cb.chanState.ChanType, fn.FlattenOption(auxLeaf),
 		)
 		if err != nil {
 			return nil, err
@@ -984,9 +1000,15 @@ func (cb *CommitmentBuilder) createUnsignedCommitmentTx(ourBalance,
 			continue
 		}
 
+		auxLeaf := fn.MapOption(
+			func(leaves input.AuxTapLeaves) input.AuxTapLeaf {
+				return leaves[htlc.HtlcIndex].AuxTapLeaf
+			},
+		)(incomingAuxLeaves)
+
 		err := addHTLC(
 			commitTx, isOurs, true, htlc, keyRing,
-			cb.chanState.ChanType,
+			cb.chanState.ChanType, fn.FlattenOption(auxLeaf),
 		)
 		if err != nil {
 			return nil, err
@@ -1266,8 +1288,8 @@ func genSegwitV0HtlcScript(chanType channeldb.ChannelType,
 // genTaprootHtlcScript generates the HTLC scripts for a taproot+musig2
 // channel.
 func genTaprootHtlcScript(isIncoming, ourCommit bool, timeout uint32,
-	rHash [32]byte,
+	rHash [32]byte, keyRing *CommitmentKeyRing,
-	keyRing *CommitmentKeyRing) (*input.HtlcScriptTree, error) {
+	auxLeaf input.AuxTapLeaf) (*input.HtlcScriptTree, error) {
 
 	var (
 		htlcScriptTree *input.HtlcScriptTree
@@ -1284,8 +1306,7 @@ func genTaprootHtlcScript(isIncoming, ourCommit bool, timeout uint32,
 	case isIncoming && ourCommit:
 		htlcScriptTree, err = input.ReceiverHTLCScriptTaproot(
 			timeout, keyRing.RemoteHtlcKey, keyRing.LocalHtlcKey,
-			keyRing.RevocationKey, rHash[:], ourCommit,
+			keyRing.RevocationKey, rHash[:], ourCommit, auxLeaf,
-			input.NoneTapLeaf(),
 		)
 
 	// We're being paid via an HTLC by the remote party, and the HTLC is
@@ -1294,8 +1315,7 @@ func genTaprootHtlcScript(isIncoming, ourCommit bool, timeout uint32,
 	case isIncoming && !ourCommit:
 		htlcScriptTree, err = input.SenderHTLCScriptTaproot(
 			keyRing.RemoteHtlcKey, keyRing.LocalHtlcKey,
-			keyRing.RevocationKey, rHash[:], ourCommit,
+			keyRing.RevocationKey, rHash[:], ourCommit, auxLeaf,
-			input.NoneTapLeaf(),
 		)
 
 	// We're sending an HTLC which is being added to our commitment
@@ -1304,8 +1324,7 @@ func genTaprootHtlcScript(isIncoming, ourCommit bool, timeout uint32,
 	case !isIncoming && ourCommit:
 		htlcScriptTree, err = input.SenderHTLCScriptTaproot(
 			keyRing.LocalHtlcKey, keyRing.RemoteHtlcKey,
-			keyRing.RevocationKey, rHash[:], ourCommit,
+			keyRing.RevocationKey, rHash[:], ourCommit, auxLeaf,
-			input.NoneTapLeaf(),
 		)
 
 	// Finally, we're paying the remote party via an HTLC, which is being
@@ -1314,8 +1333,7 @@ func genTaprootHtlcScript(isIncoming, ourCommit bool, timeout uint32,
 	case !isIncoming && !ourCommit:
 		htlcScriptTree, err = input.ReceiverHTLCScriptTaproot(
 			timeout, keyRing.LocalHtlcKey, keyRing.RemoteHtlcKey,
-			keyRing.RevocationKey, rHash[:], ourCommit,
+			keyRing.RevocationKey, rHash[:], ourCommit, auxLeaf,
-			input.NoneTapLeaf(),
 		)
 	}
 
@@ -1329,8 +1347,8 @@ func genTaprootHtlcScript(isIncoming, ourCommit bool, timeout uint32,
 // we need to sign for the remote party (2nd level HTLCs) is also returned
 // along side the multiplexer.
 func genHtlcScript(chanType channeldb.ChannelType, isIncoming, ourCommit bool,
-	timeout uint32, rHash [32]byte,
+	timeout uint32, rHash [32]byte, keyRing *CommitmentKeyRing,
-	keyRing *CommitmentKeyRing) (input.ScriptDescriptor, error) {
+	auxLeaf input.AuxTapLeaf) (input.ScriptDescriptor, error) {
 
 	if !chanType.IsTaproot() {
 		return genSegwitV0HtlcScript(
@@ -1340,7 +1358,7 @@ func genHtlcScript(chanType channeldb.ChannelType, isIncoming, ourCommit bool,
 	}
 
 	return genTaprootHtlcScript(
-		isIncoming, ourCommit, timeout, rHash, keyRing,
+		isIncoming, ourCommit, timeout, rHash, keyRing, auxLeaf,
 	)
 }
 
@@ -1353,13 +1371,15 @@ func genHtlcScript(chanType channeldb.ChannelType, isIncoming, ourCommit bool,
 // the descriptor itself.
 func addHTLC(commitTx *wire.MsgTx, ourCommit bool,
 	isIncoming bool, paymentDesc *PaymentDescriptor,
-	keyRing *CommitmentKeyRing, chanType channeldb.ChannelType) error {
+	keyRing *CommitmentKeyRing, chanType channeldb.ChannelType,
+	auxLeaf input.AuxTapLeaf) error {
 
 	timeout := paymentDesc.Timeout
 	rHash := paymentDesc.RHash
 
 	scriptInfo, err := genHtlcScript(
 		chanType, isIncoming, ourCommit, timeout, rHash, keyRing,
+		auxLeaf,
 	)
 	if err != nil {
 		return err

lnwallet/transactions.go

@@ -8,6 +8,7 @@ import (
 	"github.com/btcsuite/btcd/btcutil"
 	"github.com/btcsuite/btcd/wire"
 	"github.com/lightningnetwork/lnd/channeldb"
+	"github.com/lightningnetwork/lnd/input"
 )
 
 const (
@@ -50,8 +51,8 @@ var (
 //   - <sender sig> <receiver sig> <preimage> <success_script> <control_block>
 func CreateHtlcSuccessTx(chanType channeldb.ChannelType, initiator bool,
 	htlcOutput wire.OutPoint, htlcAmt btcutil.Amount, csvDelay,
-	leaseExpiry uint32, revocationKey, delayKey *btcec.PublicKey) (
+	leaseExpiry uint32, revocationKey, delayKey *btcec.PublicKey,
-	*wire.MsgTx, error) {
+	auxLeaf input.AuxTapLeaf) (*wire.MsgTx, error) {
 
 	// Create a version two transaction (as the success version of this
 	// spends an output with a CSV timeout).
@@ -71,7 +72,7 @@ func CreateHtlcSuccessTx(chanType channeldb.ChannelType, initiator bool,
 	// HTLC outputs.
 	scriptInfo, err := SecondLevelHtlcScript(
 		chanType, initiator, revocationKey, delayKey, csvDelay,
-		leaseExpiry,
+		leaseExpiry, auxLeaf,
 	)
 	if err != nil {
 		return nil, err
@@ -110,7 +111,8 @@ func CreateHtlcSuccessTx(chanType channeldb.ChannelType, initiator bool,
 func CreateHtlcTimeoutTx(chanType channeldb.ChannelType, initiator bool,
 	htlcOutput wire.OutPoint, htlcAmt btcutil.Amount,
 	cltvExpiry, csvDelay, leaseExpiry uint32,
-	revocationKey, delayKey *btcec.PublicKey) (*wire.MsgTx, error) {
+	revocationKey, delayKey *btcec.PublicKey,
+	auxLeaf input.AuxTapLeaf) (*wire.MsgTx, error) {
 
 	// Create a version two transaction (as the success version of this
 	// spends an output with a CSV timeout), and set the lock-time to the
@@ -134,7 +136,7 @@ func CreateHtlcTimeoutTx(chanType channeldb.ChannelType, initiator bool,
 	// HTLC outputs.
 	scriptInfo, err := SecondLevelHtlcScript(
 		chanType, initiator, revocationKey, delayKey, csvDelay,
-		leaseExpiry,
+		leaseExpiry, auxLeaf,
 	)
 	if err != nil {
 		return nil, err