sweep: update sweeper to use AuxSweeper to add extra change addr

In this commit, we start to use the AuxSweeper (if present) to obtain a
new extra change addr we should add to the sweeping transaction. With
this, we'll take the set of inputs and our change addr, and then maybe
gain a new change addr to add to the sweep transaction.

The extra change addr will be treated as an extra required tx out,
shared across all the relevant inputs. This'll also be used in
NeedWalletInput to make sure that we add an extra input if needed to be
able to pay for the change addr.

sweep/fee_bumper.go

@@ -132,6 +132,8 @@ type BumpRequest struct {
 func (r *BumpRequest) MaxFeeRateAllowed() (chainfee.SatPerKWeight, error) {
 	// Get the size of the sweep tx, which will be used to calculate the
 	// budget fee rate.
+	//
+	// TODO(roasbeef): also wants the extra change output?
 	size, err := calcSweepTxWeight(
 		r.Inputs, r.DeliveryAddress.DeliveryAddress,
 	)
@@ -176,7 +178,7 @@ func calcSweepTxWeight(inputs []input.Input,
 	// TODO(yy): we should refactor the weight estimator to not require a
 	// fee rate and max fee rate and make it a pure tx weight calculator.
 	_, estimator, err := getWeightEstimate(
-		inputs, nil, feeRate, 0, outputPkScript,
+		inputs, nil, feeRate, 0, [][]byte{outputPkScript},
 	)
 	if err != nil {
 		return 0, err
@@ -1162,9 +1164,9 @@ func (t *TxPublisher) createSweepTx(inputs []input.Input,
 	feeRate chainfee.SatPerKWeight) (*sweepTxCtx, error) {
 
 	// Validate and calculate the fee and change amount.
-	txFee, changeAmtOpt, locktimeOpt, err := prepareSweepTx(
-		inputs, changePkScript.DeliveryAddress, feeRate,
-		t.currentHeight.Load(),
+	txFee, changeOutputsOpt, locktimeOpt, err := prepareSweepTx(
+		inputs, changePkScript, feeRate, t.currentHeight.Load(),
+		t.cfg.AuxSweeper,
 	)
 	if err != nil {
 		return nil, err
@@ -1210,12 +1212,12 @@ func (t *TxPublisher) createSweepTx(inputs []input.Input,
 		})
 	}
 
-	// If there's a change amount, add it to the transaction.
-	changeAmtOpt.WhenSome(func(changeAmt btcutil.Amount) {
-		sweepTx.AddTxOut(&wire.TxOut{
-			PkScript: changePkScript.DeliveryAddress,
-			Value:    int64(changeAmt),
-		})
+	// If we have change outputs to add, then add it the sweep transaction
+	// here.
+	changeOutputsOpt.WhenSome(func(changeOuts []SweepOutput) {
+		for i := range changeOuts {
+			sweepTx.AddTxOut(&changeOuts[i].TxOut)
+		}
 	})
 
 	// We'll default to using the current block height as locktime, if none
@@ -1259,31 +1261,64 @@ func (t *TxPublisher) createSweepTx(inputs []input.Input,
 	log.Debugf("Created sweep tx %v for inputs:\n%v", sweepTx.TxHash(),
 		inputTypeSummary(inputs))
 
+	// Try to locate the extra change output, though there might be None.
+	extraTxOut := fn.MapOption(func(sweepOuts []SweepOutput) fn.Option[SweepOutput] { //nolint:lll
+		for _, sweepOut := range sweepOuts {
+			if sweepOut.IsExtra {
+				log.Infof("Sweep produced extra_sweep_out=%v",
+					spew.Sdump(sweepOut))
+
+				return fn.Some(sweepOut)
+			}
+		}
+
+		return fn.None[SweepOutput]()
+	})(changeOutputsOpt)
+
 	return &sweepTxCtx{
-		tx:  sweepTx,
-		fee: txFee,
+		tx:         sweepTx,
+		fee:        txFee,
+		extraTxOut: fn.FlattenOption(extraTxOut),
 	}, nil
 }
 
-// prepareSweepTx returns the tx fee, an optional change amount and an optional
-// locktime after a series of validations:
+// prepareSweepTx returns the tx fee, a set of optional change outputs and an
+// optional locktime after a series of validations:
 // 1. check the locktime has been reached.
 // 2. check the locktimes are the same.
 // 3. check the inputs cover the outputs.
 //
 // NOTE: if the change amount is below dust, it will be added to the tx fee.
-func prepareSweepTx(inputs []input.Input, changePkScript []byte,
-	feeRate chainfee.SatPerKWeight, currentHeight int32) (
-	btcutil.Amount, fn.Option[btcutil.Amount], fn.Option[int32], error) {
+func prepareSweepTx(inputs []input.Input, changePkScript lnwallet.AddrWithKey,
+	feeRate chainfee.SatPerKWeight, currentHeight int32,
+	auxSweeper fn.Option[AuxSweeper]) (
+	btcutil.Amount, fn.Option[[]SweepOutput], fn.Option[int32], error) {
 
-	noChange := fn.None[btcutil.Amount]()
+	noChange := fn.None[[]SweepOutput]()
 	noLocktime := fn.None[int32]()
 
+	// Given the set of inputs we have, if we have an aux sweeper, then
+	// we'll attempt to see if we have any other change outputs we'll need
+	// to add to the sweep transaction.
+	changePkScripts := [][]byte{changePkScript.DeliveryAddress}
+	extraChangeOut := fn.MapOptionZ(
+		auxSweeper,
+		func(aux AuxSweeper) fn.Result[SweepOutput] {
+			return aux.DeriveSweepAddr(inputs, changePkScript)
+		},
+	)
+	if err := extraChangeOut.Err(); err != nil {
+		return 0, noChange, noLocktime, err
+	}
+	extraChangeOut.WhenResult(func(o SweepOutput) {
+		changePkScripts = append(changePkScripts, o.PkScript)
+	})
+
 	// Creating a weight estimator with nil outputs and zero max fee rate.
 	// We don't allow adding customized outputs in the sweeping tx, and the
 	// fee rate is already being managed before we get here.
 	inputs, estimator, err := getWeightEstimate(
-		inputs, nil, feeRate, 0, changePkScript,
+		inputs, nil, feeRate, 0, changePkScripts,
 	)
 	if err != nil {
 		return 0, noChange, noLocktime, err
@@ -1301,6 +1336,12 @@ func prepareSweepTx(inputs []input.Input, changePkScript []byte,
 		requiredOutput btcutil.Amount
 	)
 
+	// If we have an extra change output, then we'll add it as a required
+	// output amt.
+	extraChangeOut.WhenResult(func(o SweepOutput) {
+		requiredOutput += btcutil.Amount(o.Value)
+	})
+
 	// Go through each input and check if the required lock times have
 	// reached and are the same.
 	for _, o := range inputs {
@@ -1348,14 +1389,21 @@ func prepareSweepTx(inputs []input.Input, changePkScript []byte,
 	// change output.
 	changeAmt := totalInput - requiredOutput - txFee
 
+	changeOuts := make([]SweepOutput, 0, 2)
+
+	extraChangeOut.WhenResult(func(o SweepOutput) {
+		changeOuts = append(changeOuts, o)
+	})
+
 	// We'll calculate the dust limit for the given changePkScript since it
 	// is variable.
-	changeFloor := lnwallet.DustLimitForSize(len(changePkScript))
+	changeFloor := lnwallet.DustLimitForSize(
+		len(changePkScript.DeliveryAddress),
+	)
 
-	changeAmtOpt := fn.Some(changeAmt)
-
-	// If the change amount is dust, we'll move it into the fees.
 	switch {
+	// If the change amount is dust, we'll move it into the fees, and
+	// ignore it.
 	case changeAmt < changeFloor:
 		log.Infof("Change amt %v below dustlimit %v, not adding "+
 			"change output", changeAmt, changeFloor)
@@ -1371,8 +1419,16 @@ func prepareSweepTx(inputs []input.Input, changePkScript []byte,
 		// The dust amount is added to the fee.
 		txFee += changeAmt
 
-		// Set the change amount to none.
-		changeAmtOpt = fn.None[btcutil.Amount]()
+	// Otherwise, we'll actually recognize it as a change output.
+	default:
+		changeOuts = append(changeOuts, SweepOutput{
+			TxOut: wire.TxOut{
+				Value:    int64(changeAmt),
+				PkScript: changePkScript.DeliveryAddress,
+			},
+			IsExtra:     false,
+			InternalKey: changePkScript.InternalKey,
+		})
 	}
 
 	// Optionally set the locktime.
@@ -1381,6 +1437,11 @@ func prepareSweepTx(inputs []input.Input, changePkScript []byte,
 		locktimeOpt = noLocktime
 	}
 
+	var changeOutsOpt fn.Option[[]SweepOutput]
+	if len(changeOuts) > 0 {
+		changeOutsOpt = fn.Some(changeOuts)
+	}
+
 	log.Debugf("Creating sweep tx for %v inputs (%s) using %v, "+
 		"tx_weight=%v, tx_fee=%v, locktime=%v, parents_count=%v, "+
 		"parents_fee=%v, parents_weight=%v, current_height=%v",
@@ -1388,5 +1449,5 @@ func prepareSweepTx(inputs []input.Input, changePkScript []byte,
 		estimator.weight(), txFee, locktimeOpt, len(estimator.parents),
 		estimator.parentsFee, estimator.parentsWeight, currentHeight)
 
-	return txFee, changeAmtOpt, locktimeOpt, nil
+	return txFee, changeOutsOpt, locktimeOpt, nil
 }

sweep/tx_input_set.go

@@ -220,6 +220,26 @@ func (b *BudgetInputSet) NeedWalletInput() bool {
 		budgetBorrowable btcutil.Amount
 	)
 
+	// If any of the outputs in the set have a resolution blob, then this
+	// means we'll end up needing an extra change output. We'll tack this
+	// on now as an extra portion of the budget.
+	extraRequiredTxOut := fn.Any(func(i *SweeperInput) bool {
+		// If there's a required txout, then we don't count this as
+		// it'll be a second level HTLC.
+		if i.RequiredTxOut() != nil {
+			return false
+		}
+
+		// Otherwise, we need one if we have a resolution blob.
+		return i.ResolutionBlob().IsSome()
+	}, b.inputs)
+
+	if extraRequiredTxOut {
+		// TODO(roasbeef): aux sweeper ext to ask for extra output
+		// params and value?
+		budgetNeeded += 1_000
+	}
+
 	for _, inp := range b.inputs {
 		// If this input has a required output, we can assume it's a
 		// second-level htlc txns input. Although this input must have

sweep/txgenerator.go

@@ -38,7 +38,7 @@ func createSweepTx(inputs []input.Input, outputs []*wire.TxOut,
 	signer input.Signer) (*wire.MsgTx, btcutil.Amount, error) {
 
 	inputs, estimator, err := getWeightEstimate(
-		inputs, outputs, feeRate, maxFeeRate, changePkScript,
+		inputs, outputs, feeRate, maxFeeRate, [][]byte{changePkScript},
 	)
 	if err != nil {
 		return nil, 0, err
@@ -221,7 +221,7 @@ func createSweepTx(inputs []input.Input, outputs []*wire.TxOut,
 // Additionally, it returns counts for the number of csv and cltv inputs.
 func getWeightEstimate(inputs []input.Input, outputs []*wire.TxOut,
 	feeRate, maxFeeRate chainfee.SatPerKWeight,
-	outputPkScript []byte) ([]input.Input, *weightEstimator, error) {
+	outputPkScripts [][]byte) ([]input.Input, *weightEstimator, error) {
 
 	// We initialize a weight estimator so we can accurately asses the
 	// amount of fees we need to pay for this sweep transaction.
@@ -237,31 +237,33 @@ func getWeightEstimate(inputs []input.Input, outputs []*wire.TxOut,
 
 	// If there is any leftover change after paying to the given outputs
 	// and required outputs, it will go to a single segwit p2wkh or p2tr
-	// address. This will be our change address, so ensure it contributes to
-	// our weight estimate. Note that if we have other outputs, we might end
-	// up creating a sweep tx without a change output. It is okay to add the
-	// change output to the weight estimate regardless, since the estimated
-	// fee will just be subtracted from this already dust output, and
-	// trimmed.
-	switch {
-	case txscript.IsPayToTaproot(outputPkScript):
-		weightEstimate.addP2TROutput()
+	// address. This will be our change address, so ensure it contributes
+	// to our weight estimate. Note that if we have other outputs, we might
+	// end up creating a sweep tx without a change output. It is okay to
+	// add the change output to the weight estimate regardless, since the
+	// estimated fee will just be subtracted from this already dust output,
+	// and trimmed.
+	for _, outputPkScript := range outputPkScripts {
+		switch {
+		case txscript.IsPayToTaproot(outputPkScript):
+			weightEstimate.addP2TROutput()
 
-	case txscript.IsPayToWitnessScriptHash(outputPkScript):
-		weightEstimate.addP2WSHOutput()
+		case txscript.IsPayToWitnessScriptHash(outputPkScript):
+			weightEstimate.addP2WSHOutput()
 
-	case txscript.IsPayToWitnessPubKeyHash(outputPkScript):
-		weightEstimate.addP2WKHOutput()
+		case txscript.IsPayToWitnessPubKeyHash(outputPkScript):
+			weightEstimate.addP2WKHOutput()
 
-	case txscript.IsPayToPubKeyHash(outputPkScript):
-		weightEstimate.estimator.AddP2PKHOutput()
+		case txscript.IsPayToPubKeyHash(outputPkScript):
+			weightEstimate.estimator.AddP2PKHOutput()
 
-	case txscript.IsPayToScriptHash(outputPkScript):
-		weightEstimate.estimator.AddP2SHOutput()
+		case txscript.IsPayToScriptHash(outputPkScript):
+			weightEstimate.estimator.AddP2SHOutput()
 
-	default:
-		// Unknown script type.
-		return nil, nil, errors.New("unknown script type")
+		default:
+			// Unknown script type.
+			return nil, nil, errors.New("unknown script type")
+		}
 	}
 
 	// For each output, use its witness type to determine the estimate

sweep/txgenerator_test.go

@@ -51,7 +51,7 @@ func TestWeightEstimate(t *testing.T) {
 	}
 
 	_, estimator, err := getWeightEstimate(
-		inputs, nil, 0, 0, changePkScript,
+		inputs, nil, 0, 0, [][]byte{changePkScript},
 	)
 	require.NoError(t, err)
 
@@ -153,7 +153,7 @@ func testUnknownScriptInner(t *testing.T, pkscript []byte, expectFail bool) {
 		))
 	}
 
-	_, _, err := getWeightEstimate(inputs, nil, 0, 0, pkscript)
+	_, _, err := getWeightEstimate(inputs, nil, 0, 0, [][]byte{pkscript})
 	if expectFail {
 		require.Error(t, err)
 	} else {