contractcourt: integration aux sweeper to breach arb

Similar to the sweeper, when we're about to make a new breach
transaction, we ask the sweeper for a new change address, if it has one.
Then when we go to publish, we notify broadcast.

Author: Roasbeef

contractcourt/breach_arbitrator.go

@@ -23,6 +23,7 @@ import (
 	"github.com/lightningnetwork/lnd/lnutils"
 	"github.com/lightningnetwork/lnd/lnwallet"
 	"github.com/lightningnetwork/lnd/lnwallet/chainfee"
+	"github.com/lightningnetwork/lnd/sweep"
 	"github.com/lightningnetwork/lnd/tlv"
 )
 
@@ -174,6 +175,10 @@ type BreachConfig struct {
 	// breached channels. This is used in conjunction with DB to recover
 	// from crashes, restarts, or other failures.
 	Store RetributionStorer
+
+	// AuxSweeper is an optional interface that can be used to modify the
+	// way sweep transaction are generated.
+	AuxSweeper fn.Option[sweep.AuxSweeper]
 }
 
 // BreachArbitrator is a special subsystem which is responsible for watching and
@@ -737,10 +742,28 @@ justiceTxBroadcast:
 	brarLog.Debugf("Broadcasting justice tx: %v", lnutils.SpewLogClosure(
 		finalTx))
 
+	// As we're about to broadcast our breach transaction, we'll notify the
+	// aux sweeper of our broadcast attempt first.
+	err = fn.MapOptionZ(b.cfg.AuxSweeper, func(aux sweep.AuxSweeper) error {
+		bumpReq := sweep.BumpRequest{
+			Inputs:          finalTx.inputs,
+			DeliveryAddress: finalTx.sweepAddr,
+			ExtraTxOut:      finalTx.extraTxOut,
+		}
+
+		return aux.NotifyBroadcast(
+			&bumpReq, finalTx.justiceTx, finalTx.fee,
+		)
+	})
+	if err != nil {
+		brarLog.Errorf("unable to notify broadcast: %w", err)
+		return
+	}
+
 	// We'll now attempt to broadcast the transaction which finalized the
 	// channel's retribution against the cheating counter party.
 	label := labels.MakeLabel(labels.LabelTypeJusticeTransaction, nil)
-	err = b.cfg.PublishTransaction(finalTx, label)
+	err = b.cfg.PublishTransaction(finalTx.justiceTx, label)
 	if err != nil {
 		brarLog.Errorf("Unable to broadcast justice tx: %v", err)
 	}
@@ -860,7 +883,9 @@ Loop:
 					"spending commitment outs: %v",
 					lnutils.SpewLogClosure(tx))
 
-				err = b.cfg.PublishTransaction(tx, label)
+				err = b.cfg.PublishTransaction(
+					tx.justiceTx, label,
+				)
 				if err != nil {
 					brarLog.Warnf("Unable to broadcast "+
 						"commit out spending justice "+
@@ -875,7 +900,9 @@ Loop:
 					"spending HTLC outs: %v",
 					lnutils.SpewLogClosure(tx))
 
-				err = b.cfg.PublishTransaction(tx, label)
+				err = b.cfg.PublishTransaction(
+					tx.justiceTx, label,
+				)
 				if err != nil {
 					brarLog.Warnf("Unable to broadcast "+
 						"HTLC out spending justice "+
@@ -890,7 +917,9 @@ Loop:
 					"spending second-level HTLC output: %v",
 					lnutils.SpewLogClosure(tx))
 
-				err = b.cfg.PublishTransaction(tx, label)
+				err = b.cfg.PublishTransaction(
+					tx.justiceTx, label,
+				)
 				if err != nil {
 					brarLog.Warnf("Unable to broadcast "+
 						"second-level HTLC out "+
@@ -1371,10 +1400,10 @@ func newRetributionInfo(chanPoint *wire.OutPoint,
 // spend the to_local output and commitment level HTLC outputs separately,
 // before the CSV locks expire.
 type justiceTxVariants struct {
-	spendAll              *wire.MsgTx
-	spendCommitOuts       *wire.MsgTx
-	spendHTLCs            *wire.MsgTx
-	spendSecondLevelHTLCs []*wire.MsgTx
+	spendAll              *justiceTxCtx
+	spendCommitOuts       *justiceTxCtx
+	spendHTLCs            *justiceTxCtx
+	spendSecondLevelHTLCs []*justiceTxCtx
 }
 
 // createJusticeTx creates transactions which exacts "justice" by sweeping ALL
@@ -1438,7 +1467,9 @@ func (b *BreachArbitrator) createJusticeTx(
 			err)
 	}
 
-	secondLevelSweeps := make([]*wire.MsgTx, 0, len(secondLevelInputs))
+	// TODO(roasbeef): only register one of them?
+
+	secondLevelSweeps := make([]*justiceTxCtx, 0, len(secondLevelInputs))
 	for _, input := range secondLevelInputs {
 		sweepTx, err := b.createSweepTx(input)
 		if err != nil {
@@ -1455,9 +1486,23 @@ func (b *BreachArbitrator) createJusticeTx(
 	return txs, nil
 }
 
+// justiceTxCtx contains the justice transaction along with other related meta
+// data.
+type justiceTxCtx struct {
+	justiceTx *wire.MsgTx
+
+	sweepAddr lnwallet.AddrWithKey
+
+	extraTxOut fn.Option[sweep.SweepOutput]
+
+	fee btcutil.Amount
+
+	inputs []input.Input
+}
+
 // createSweepTx creates a tx that sweeps the passed inputs back to our wallet.
-func (b *BreachArbitrator) createSweepTx(inputs ...input.Input) (*wire.MsgTx,
-	error) {
+func (b *BreachArbitrator) createSweepTx(
+	inputs ...input.Input) (*justiceTxCtx, error) {
 
 	if len(inputs) == 0 {
 		return nil, nil
@@ -1480,6 +1525,15 @@ func (b *BreachArbitrator) createSweepTx(inputs ...input.Input) (*wire.MsgTx,
 	// nLockTime, and output are already included in the TxWeightEstimator.
 	weightEstimate.AddP2TROutput()
 
+	// If any of our inputs has a resolution blob, then we'll add another
+	// P2TR output.
+	hasBlobs := fn.Any(func(i input.Input) bool {
+		return i.ResolutionBlob().IsSome()
+	}, inputs)
+	if hasBlobs {
+		weightEstimate.AddP2TROutput()
+	}
+
 	// Next, we iterate over the breached outputs contained in the
 	// retribution info.  For each, we switch over the witness type such
 	// that we contribute the appropriate weight for each input and
@@ -1513,7 +1567,7 @@ func (b *BreachArbitrator) createSweepTx(inputs ...input.Input) (*wire.MsgTx,
 // sweepSpendableOutputsTxn creates a signed transaction from a sequence of
 // spendable outputs by sweeping the funds into a single p2wkh output.
 func (b *BreachArbitrator) sweepSpendableOutputsTxn(txWeight lntypes.WeightUnit,
-	inputs ...input.Input) (*wire.MsgTx, error) {
+	inputs ...input.Input) (*justiceTxCtx, error) {
 
 	// First, we obtain a new public key script from the wallet which we'll
 	// sweep the funds to.
@@ -1538,19 +1592,41 @@ func (b *BreachArbitrator) sweepSpendableOutputsTxn(txWeight lntypes.WeightUnit,
 	}
 	txFee := feePerKw.FeeForWeight(txWeight)
 
+	// At this point, we'll check to see if we have any extra outputs to
+	// add from the aux sweeper.
+	extraChangeOut := fn.MapOptionZ(
+		b.cfg.AuxSweeper,
+		func(aux sweep.AuxSweeper) fn.Result[sweep.SweepOutput] {
+			return aux.DeriveSweepAddr(inputs, pkScript)
+		},
+	)
+	if err := extraChangeOut.Err(); err != nil {
+		return nil, err
+	}
+
 	// TODO(roasbeef): already start to siphon their funds into fees
 	sweepAmt := int64(totalAmt - txFee)
 
 	// With the fee calculated, we can now create the transaction using the
 	// information gathered above and the provided retribution information.
 	txn := wire.NewMsgTx(2)
 
-	// We begin by adding the output to which our funds will be deposited.
+	// First, we'll add the extra swep output if it exists, subtracting the
+	// amount from the sweep amt.
+	extraChangeOut.WhenResult(func(o sweep.SweepOutput) {
+		sweepAmt -= o.Value
+
+		txn.AddTxOut(&o.TxOut)
+	})
+
+	// Next, we'll add the output to which our funds will be deposited.
 	txn.AddTxOut(&wire.TxOut{
 		PkScript: pkScript.DeliveryAddress,
 		Value:    sweepAmt,
 	})
 
+	// TODO(roasbeef): add other output change modify sweep amt
+
 	// Next, we add all of the spendable outputs as inputs to the
 	// transaction.
 	for _, inp := range inputs {
@@ -1606,7 +1682,13 @@ func (b *BreachArbitrator) sweepSpendableOutputsTxn(txWeight lntypes.WeightUnit,
 		}
 	}
 
-	return txn, nil
+	return &justiceTxCtx{
+		justiceTx:  txn,
+		sweepAddr:  pkScript,
+		extraTxOut: extraChangeOut.Option(),
+		fee:        txFee,
+		inputs:     inputs,
+	}, nil
 }
 
 // RetributionStore handles persistence of retribution states to disk and is

contractcourt/breach_arbitrator_test.go

@@ -1230,16 +1230,16 @@ func TestBreachCreateJusticeTx(t *testing.T) {
 
 	// The spendAll tx should be spending all the outputs. This is the
 	// "regular" justice transaction type.
-	require.Len(t, justiceTxs.spendAll.TxIn, len(breachedOutputs))
+	require.Len(t, justiceTxs.spendAll.justiceTx.TxIn, len(breachedOutputs))
 
 	// The spendCommitOuts tx should be spending the 4 types of commit outs
 	// (note that in practice there will be at most two commit outputs per
 	// commit, but we test all 4 types here).
-	require.Len(t, justiceTxs.spendCommitOuts.TxIn, 4)
+	require.Len(t, justiceTxs.spendCommitOuts.justiceTx.TxIn, 4)
 
 	// Check that the spendHTLCs tx is spending the two revoked commitment
 	// level HTLC output types.
-	require.Len(t, justiceTxs.spendHTLCs.TxIn, 2)
+	require.Len(t, justiceTxs.spendHTLCs.justiceTx.TxIn, 2)
 
 	// Finally, check that the spendSecondLevelHTLCs txs are spending the
 	// second level type.

server.go

@@ -1181,6 +1181,7 @@ func newServer(cfg *Config, listenAddrs []net.Addr,
 			Store: contractcourt.NewRetributionStore(
 				dbs.ChanStateDB,
 			),
+			AuxSweeper: s.implCfg.AuxSweeper,
 		},
 	)