contractcourt: track historic justice tx variants across rebuild cycles

The breach arbiter rebuilds justice tx variants after each spend
detection (e.g. when an HTLC transitions to second-level). The tx that
ultimately confirms may have been created in an earlier rebuild cycle
and is no longer present in the current justiceTxVariants struct.

Add a historicJusticeTxs map that records every justice tx variant ever
created (keyed by txid) via recordJusticeTxVariants(). The
notifyConfirmedJusticeTx function now falls back to this map when the
confirmed spend doesn't match any current variant, ensuring the aux
sweeper receives NotifyBroadcast for asset proof generation.

Also improve the split-broadcast path: rebuild justice tx variants
from the updated breach info before splitting, and re-attempt the
spendAll variant first (which may now succeed after second-level
spends have been incorporated). Add logging to createJusticeTx for
input counts and variant creation.

Author: GeorgeTsagk

contractcourt/breach_arbitrator.go

@@ -820,7 +820,9 @@ func updateBreachInfo(breachInfo *retributionInfo, spends []spend,
 // haven't already notified about it, we call NotifyBroadcast on the aux sweeper
 // to generate asset-level proofs.
 func (b *BreachArbitrator) notifyConfirmedJusticeTx(spends []spend,
-	justiceTxs *justiceTxVariants, notifiedTxs map[chainhash.Hash]bool) {
+	justiceTxs *justiceTxVariants,
+	historicJusticeTxs map[chainhash.Hash]*justiceTxCtx,
+	notifiedTxs map[chainhash.Hash]bool) {
 
 	// Check each spend to see if it's from one of our justice txs.
 	for _, s := range spends {
@@ -864,6 +866,14 @@ func (b *BreachArbitrator) notifyConfirmedJusticeTx(spends []spend,
 			}
 		}
 
+		// Check the historic map of all justice tx variants ever
+		// created. This handles the case where the confirmed tx
+		// was from a previous rebuild cycle and the current
+		// justiceTxs has been replaced with newer variants.
+		if justiceCtx == nil {
+			justiceCtx = historicJusticeTxs[spendingTxHash]
+		}
+
 		// If this is one of our justice txs, notify the aux sweeper.
 		if justiceCtx != nil {
 			bumpReq := sweep.BumpRequest{
@@ -896,6 +906,28 @@ func (b *BreachArbitrator) notifyConfirmedJusticeTx(spends []spend,
 	}
 }
 
+// recordJusticeTxVariants records all non-nil justice tx variants into the
+// historic map keyed by their txid. This allows notifyConfirmedJusticeTx to
+// match confirmed spends against justice txs from previous rebuild cycles.
+func recordJusticeTxVariants(variants *justiceTxVariants,
+	history map[chainhash.Hash]*justiceTxCtx) {
+
+	record := func(jtx *justiceTxCtx) {
+		if jtx == nil {
+			return
+		}
+		hash := jtx.justiceTx.TxHash()
+		history[hash] = jtx
+	}
+
+	record(variants.spendAll)
+	record(variants.spendCommitOuts)
+	record(variants.spendHTLCs)
+	for _, tx := range variants.spendSecondLevelHTLCs {
+		record(tx)
+	}
+}
+
 // exactRetribution is a goroutine which is executed once a contract breach has
 // been detected by a breachObserver. This function is responsible for
 // punishing a counterparty for violating the channel contract by sweeping ALL
@@ -951,6 +983,13 @@ func (b *BreachArbitrator) exactRetribution(
 	// handle duplicate calls idempotently.
 	notifiedJusticeTxs := make(map[chainhash.Hash]bool)
 
+	// Track all historically created justice tx contexts by their txid.
+	// This is needed because justiceTxs is rebuilt after each spend
+	// detection, and the tx that actually confirmed may have been from
+	// an earlier rebuild cycle. Without this history, we can't match
+	// the confirmed tx to call NotifyBroadcast on the aux sweeper.
+	historicJusticeTxs := make(map[chainhash.Hash]*justiceTxCtx)
+
 	// Compute both the total value of funds being swept and the
 	// amount of funds that were revoked from the counter party.
 	var totalFunds, revokedFunds btcutil.Amount
@@ -964,6 +1003,7 @@ justiceTxBroadcast:
 		brarLog.Errorf("Unable to create justice tx: %v", err)
 		return
 	}
+	recordJusticeTxVariants(justiceTxs, historicJusticeTxs)
 	finalTx := justiceTxs.spendAll
 
 	brarLog.Debugf("Broadcasting justice tx: %v", lnutils.SpewLogClosure(
@@ -1017,7 +1057,9 @@ Loop:
 			// justice transaction, and if so, notify the aux
 			// sweeper.
 			b.notifyConfirmedJusticeTx(
-				spends, justiceTxs, notifiedJusticeTxs,
+				spends, justiceTxs,
+				historicJusticeTxs,
+				notifiedJusticeTxs,
 			)
 
 			// Update the breach info with the new spends.
@@ -1085,6 +1127,44 @@ Loop:
 				"height %v), splitting justice tx.",
 				epoch.Height, breachInfo.breachHeight)
 
+			// Rebuild justice tx variants from the current
+			// breach info, which may have been updated by
+			// spend detection (e.g. second-level HTLC spends
+			// replacing commit-level outputs).
+			justiceTxs, err = b.createJusticeTx(
+				breachInfo.breachedOutputs,
+			)
+			if err != nil {
+				brarLog.Errorf("Unable to recreate "+
+					"justice tx for split: %v", err)
+				continue Loop
+			}
+			recordJusticeTxVariants(
+				justiceTxs, historicJusticeTxs,
+			)
+
+			// Re-attempt the spendAll variant first, in
+			// case the breach info was updated since the
+			// initial broadcast. This avoids splitting into
+			// small txs that can't pay fees when a combined
+			// tx would work.
+			if justiceTxs.spendAll != nil {
+				label := labels.MakeLabel(
+					labels.LabelTypeJusticeTransaction,
+					nil,
+				)
+				err = b.cfg.PublishTransaction(
+					justiceTxs.spendAll.justiceTx,
+					label,
+				)
+				if err != nil {
+					brarLog.Warnf("Unable to "+
+						"broadcast updated "+
+						"spendAll justice "+
+						"tx: %v", err)
+				}
+			}
+
 			// Otherwise we'll attempt to publish the two separate
 			// justice transactions that sweeps the commitment
 			// outputs and the HTLC outputs separately. This is to
@@ -1675,6 +1755,10 @@ func (b *BreachArbitrator) createJusticeTx(
 		}
 	}
 
+	brarLog.Infof("createJusticeTx: %d total inputs (%d commit, "+
+		"%d htlc, %d second-level)", len(allInputs),
+		len(commitInputs), len(htlcInputs), len(secondLevelInputs))
+
 	var (
 		txs = &justiceTxVariants{}
 		err error
@@ -1685,31 +1769,42 @@ func (b *BreachArbitrator) createJusticeTx(
 	if err != nil {
 		return nil, err
 	}
+	brarLog.Infof("createJusticeTx: spendAll created successfully "+
+		"(%d inputs, txid=%v)", len(allInputs),
+		txs.spendAll.justiceTx.TxHash())
 
 	txs.spendCommitOuts, err = b.createSweepTx(commitInputs...)
 	if err != nil {
 		brarLog.Errorf("could not create sweep tx for commitment "+
 			"outputs: %v", err)
+	} else if txs.spendCommitOuts != nil {
+		brarLog.Infof("createJusticeTx: spendCommitOuts created "+
+			"successfully (%d inputs)", len(commitInputs))
 	}
 
 	txs.spendHTLCs, err = b.createSweepTx(htlcInputs...)
 	if err != nil {
 		brarLog.Errorf("could not create sweep tx for HTLC outputs: %v",
 			err)
+	} else if txs.spendHTLCs != nil {
+		brarLog.Infof("createJusticeTx: spendHTLCs created "+
+			"successfully (%d inputs)", len(htlcInputs))
 	}
 
 	// TODO(roasbeef): only register one of them?
 
 	secondLevelSweeps := make([]*justiceTxCtx, 0, len(secondLevelInputs))
-	for _, input := range secondLevelInputs {
+	for i, input := range secondLevelInputs {
 		sweepTx, err := b.createSweepTx(input)
 		if err != nil {
 			brarLog.Errorf("could not create sweep tx for "+
-				"second-level HTLC output: %v", err)
+				"second-level HTLC output %d: %v", i, err)
 
 			continue
 		}
 
+		brarLog.Infof("createJusticeTx: individual second-level "+
+			"sweep %d created successfully", i)
 		secondLevelSweeps = append(secondLevelSweeps, sweepTx)
 	}
 	txs.spendSecondLevelHTLCs = secondLevelSweeps

contractcourt/breach_arbitrator_test.go

@@ -2676,8 +2676,10 @@ func TestNotifyConfirmedJusticeTx(t *testing.T) {
 			}
 
 			// Call the function under test.
+			historicTxs := make(map[chainhash.Hash]*justiceTxCtx)
 			brar.notifyConfirmedJusticeTx(
-				tc.spends, tc.justiceTxs, tc.notifiedTxs,
+				tc.spends, tc.justiceTxs,
+				historicTxs, tc.notifiedTxs,
 			)
 
 			// Verify the number of NotifyBroadcast calls.
@@ -2746,8 +2748,9 @@ func TestNotifyConfirmedJusticeTxNoAuxSweeper(t *testing.T) {
 
 	// Should not panic and should not mark as notified since there's no
 	// aux sweeper to notify.
+	historicTxs := make(map[chainhash.Hash]*justiceTxCtx)
 	brar.notifyConfirmedJusticeTx(
-		spends, justiceTxs, notifiedTxs,
+		spends, justiceTxs, historicTxs, notifiedTxs,
 	)
 
 	// The tx should still be marked as notified even without an aux