sweep: add new AuxSweeper interface

In this commit, we add a new AuxSweeper interface. This'll take a set of
inputs, and a change addr for the sweep transaction, then optionally
return a new sweep output to be added to the sweep transaction.

We also add a new NotifyBroadcast method.  This'll be used to notify
that we're _about_ to broadcast a sweeping transaction. The set of
inputs is passed in, which allows the caller to prepare for the ultimate
broadcast of the sweeping transaction.

We also add ExtraTxOut to BumpRequest pass fees to NotifyBroadcast. This
allows the callee to know the total fee of the sweeping transaction.

lnwallet/interface.go

@@ -600,7 +600,7 @@ type MessageSigner interface {
 type AddrWithKey struct {
 	lnwire.DeliveryAddress
 
-	InternalKey fn.Option[btcec.PublicKey]
+	InternalKey fn.Option[keychain.KeyDescriptor]
 
 	// TODO(roasbeef): consolidate w/ instance in chan closer
 }

server.go

@@ -4962,12 +4962,7 @@ func newSweepPkScriptGen(
 
 		return fn.Ok(lnwallet.AddrWithKey{
 			DeliveryAddress: addr,
-			InternalKey: fn.MapOption(func(
+			InternalKey:     internalKeyDesc,
-				desc keychain.KeyDescriptor) btcec.PublicKey {
-
-				return *desc.PubKey
-			},
-			)(internalKeyDesc),
 		})
 	}
 }

sweep/fee_bumper.go

@@ -110,7 +110,7 @@ type BumpRequest struct {
 	DeadlineHeight int32
 
 	// DeliveryAddress is the script to send the change output to.
-	DeliveryAddress []byte
+	DeliveryAddress lnwallet.AddrWithKey
 
 	// MaxFeeRate is the maximum fee rate that can be used for fee bumping.
 	MaxFeeRate chainfee.SatPerKWeight
@@ -118,6 +118,10 @@ type BumpRequest struct {
 	// StartingFeeRate is an optional parameter that can be used to specify
 	// the initial fee rate to use for the fee function.
 	StartingFeeRate fn.Option[chainfee.SatPerKWeight]
+
+	// ExtraTxOut tracks if this bump request has an optional set of extra
+	// outputs to add to the transaction.
+	ExtraTxOut fn.Option[SweepOutput]
 }
 
 // MaxFeeRateAllowed returns the maximum fee rate allowed for the given
@@ -127,7 +131,9 @@ 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.
-	size, err := calcSweepTxWeight(r.Inputs, r.DeliveryAddress)
+	size, err := calcSweepTxWeight(
+		r.Inputs, r.DeliveryAddress.DeliveryAddress,
+	)
 	if err != nil {
 		return 0, err
 	}
@@ -248,6 +254,10 @@ type TxPublisherConfig struct {
 
 	// Notifier is used to monitor the confirmation status of the tx.
 	Notifier chainntnfs.ChainNotifier
+
+	// AuxSweeper is an optional interface that can be used to modify the
+	// way sweep transaction are generated.
+	AuxSweeper fn.Option[AuxSweeper]
 }
 
 // TxPublisher is an implementation of the Bumper interface. It utilizes the
@@ -401,16 +411,18 @@ func (t *TxPublisher) createRBFCompliantTx(req *BumpRequest,
 	for {
 		// Create a new tx with the given fee rate and check its
 		// mempool acceptance.
-		tx, fee, err := t.createAndCheckTx(req, f)
+		sweepCtx, err := t.createAndCheckTx(req, f)
 
 		switch {
 		case err == nil:
 			// The tx is valid, return the request ID.
-			requestID := t.storeRecord(tx, req, f, fee)
+			requestID := t.storeRecord(
+				sweepCtx.tx, req, f, sweepCtx.fee,
+			)
 
 			log.Infof("Created tx %v for %v inputs: feerate=%v, "+
-				"fee=%v, inputs=%v", tx.TxHash(),
+				"fee=%v, inputs=%v", sweepCtx.tx.TxHash(),
-				len(req.Inputs), f.FeeRate(), fee,
+				len(req.Inputs), f.FeeRate(), sweepCtx.fee,
 				inputTypeSummary(req.Inputs))
 
 			return requestID, nil
@@ -421,8 +433,8 @@ func (t *TxPublisher) createRBFCompliantTx(req *BumpRequest,
 			// We should at least start with a feerate above the
 			// mempool min feerate, so if we get this error, it
 			// means something is wrong earlier in the pipeline.
-			log.Errorf("Current fee=%v, feerate=%v, %v", fee,
+			log.Errorf("Current fee=%v, feerate=%v, %v",
-				f.FeeRate(), err)
+				sweepCtx.fee, f.FeeRate(), err)
 
 			fallthrough
 
@@ -434,8 +446,8 @@ func (t *TxPublisher) createRBFCompliantTx(req *BumpRequest,
 			// increased or maxed out.
 			for !increased {
 				log.Debugf("Increasing fee for next round, "+
-					"current fee=%v, feerate=%v", fee,
+					"current fee=%v, feerate=%v",
-					f.FeeRate())
+					sweepCtx.fee, f.FeeRate())
 
 				// If the fee function tells us that we have
 				// used up the budget, we will return an error
@@ -484,30 +496,34 @@ func (t *TxPublisher) storeRecord(tx *wire.MsgTx, req *BumpRequest,
 // script, and the fee rate. In addition, it validates the tx's mempool
 // acceptance before returning a tx that can be published directly, along with
 // its fee.
-func (t *TxPublisher) createAndCheckTx(req *BumpRequest, f FeeFunction) (
+func (t *TxPublisher) createAndCheckTx(req *BumpRequest,
-	*wire.MsgTx, btcutil.Amount, error) {
+	f FeeFunction) (*sweepTxCtx, error) {
 
 	// Create the sweep tx with max fee rate of 0 as the fee function
 	// guarantees the fee rate used here won't exceed the max fee rate.
-	tx, fee, err := t.createSweepTx(
+	sweepCtx, err := t.createSweepTx(
 		req.Inputs, req.DeliveryAddress, f.FeeRate(),
 	)
 	if err != nil {
-		return nil, fee, fmt.Errorf("create sweep tx: %w", err)
+		return sweepCtx, fmt.Errorf("create sweep tx: %w", err)
 	}
 
 	// Sanity check the budget still covers the fee.
-	if fee > req.Budget {
+	if sweepCtx.fee > req.Budget {
-		return nil, fee, fmt.Errorf("%w: budget=%v, fee=%v",
+		return sweepCtx, fmt.Errorf("%w: budget=%v, fee=%v",
-			ErrNotEnoughBudget, req.Budget, fee)
+			ErrNotEnoughBudget, req.Budget, sweepCtx.fee)
 	}
 
+	// If we had an extra txOut, then we'll update the result to include
+	// it.
+	req.ExtraTxOut = sweepCtx.extraTxOut
+
 	// Validate the tx's mempool acceptance.
-	err = t.cfg.Wallet.CheckMempoolAcceptance(tx)
+	err = t.cfg.Wallet.CheckMempoolAcceptance(sweepCtx.tx)
 
 	// Exit early if the tx is valid.
 	if err == nil {
-		return tx, fee, nil
+		return sweepCtx, nil
 	}
 
 	// Print an error log if the chain backend doesn't support the mempool
@@ -515,18 +531,18 @@ func (t *TxPublisher) createAndCheckTx(req *BumpRequest, f FeeFunction) (
 	if errors.Is(err, rpcclient.ErrBackendVersion) {
 		log.Errorf("TestMempoolAccept not supported by backend, " +
 			"consider upgrading it to a newer version")
-		return tx, fee, nil
+		return sweepCtx, nil
 	}
 
 	// We are running on a backend that doesn't implement the RPC
 	// testmempoolaccept, eg, neutrino, so we'll skip the check.
 	if errors.Is(err, chain.ErrUnimplemented) {
 		log.Debug("Skipped testmempoolaccept due to not implemented")
-		return tx, fee, nil
+		return sweepCtx, nil
 	}
 
-	return nil, fee, fmt.Errorf("tx=%v failed mempool check: %w",
+	return sweepCtx, fmt.Errorf("tx=%v failed mempool check: %w",
-		tx.TxHash(), err)
+		sweepCtx.tx.TxHash(), err)
 }
 
 // broadcast takes a monitored tx and publishes it to the network. Prior to the
@@ -547,6 +563,15 @@ func (t *TxPublisher) broadcast(requestID uint64) (*BumpResult, error) {
 	log.Debugf("Publishing sweep tx %v, num_inputs=%v, height=%v",
 		txid, len(tx.TxIn), t.currentHeight.Load())
 
+	// Before we go to broadcast, we'll notify the aux sweeper, if it's
+	// present of this new broadcast attempt.
+	err := fn.MapOptionZ(t.cfg.AuxSweeper, func(aux AuxSweeper) error {
+		return aux.NotifyBroadcast(record.req, tx, record.fee)
+	})
+	if err != nil {
+		return nil, fmt.Errorf("unable to notify aux sweeper: %w", err)
+	}
+
 	// Set the event, and change it to TxFailed if the wallet fails to
 	// publish it.
 	event := TxPublished
@@ -554,7 +579,7 @@ func (t *TxPublisher) broadcast(requestID uint64) (*BumpResult, error) {
 	// Publish the sweeping tx with customized label. If the publish fails,
 	// this error will be saved in the `BumpResult` and it will be removed
 	// from being monitored.
-	err := t.cfg.Wallet.PublishTransaction(
+	err = t.cfg.Wallet.PublishTransaction(
 		tx, labels.MakeLabel(labels.LabelTypeSweepTransaction, nil),
 	)
 	if err != nil {
@@ -933,7 +958,7 @@ func (t *TxPublisher) createAndPublishTx(requestID uint64,
 	// NOTE: The fee function is expected to have increased its returned
 	// fee rate after calling the SkipFeeBump method. So we can use it
 	// directly here.
-	tx, fee, err := t.createAndCheckTx(r.req, r.feeFunction)
+	sweepCtx, err := t.createAndCheckTx(r.req, r.feeFunction)
 
 	// If the error is fee related, we will return no error and let the fee
 	// bumper retry it at next block.
@@ -980,17 +1005,17 @@ func (t *TxPublisher) createAndPublishTx(requestID uint64,
 	// The tx has been created without any errors, we now register a new
 	// record by overwriting the same requestID.
 	t.records.Store(requestID, &monitorRecord{
-		tx:          tx,
+		tx:          sweepCtx.tx,
 		req:         r.req,
 		feeFunction: r.feeFunction,
-		fee:         fee,
+		fee:         sweepCtx.fee,
 	})
 
 	// Attempt to broadcast this new tx.
 	result, err := t.broadcast(requestID)
 	if err != nil {
 		log.Infof("Failed to broadcast replacement tx %v: %v",
-			tx.TxHash(), err)
+			sweepCtx.tx.TxHash(), err)
 
 		return fn.None[BumpResult]()
 	}
@@ -1016,7 +1041,8 @@ func (t *TxPublisher) createAndPublishTx(requestID uint64,
 		return fn.Some(*result)
 	}
 
-	log.Infof("Replaced tx=%v with new tx=%v", oldTx.TxHash(), tx.TxHash())
+	log.Infof("Replaced tx=%v with new tx=%v", oldTx.TxHash(),
+		sweepCtx.tx.TxHash())
 
 	// Otherwise, it's a successful RBF, set the event and return.
 	result.Event = TxReplaced
@@ -1129,17 +1155,28 @@ func calcCurrentConfTarget(currentHeight, deadline int32) uint32 {
 	return confTarget
 }
 
+// sweepTxCtx houses a sweep transaction with additional context.
+type sweepTxCtx struct {
+	tx *wire.MsgTx
+
+	fee btcutil.Amount
+
+	extraTxOut fn.Option[SweepOutput]
+}
+
 // createSweepTx creates a sweeping tx based on the given inputs, change
 // address and fee rate.
-func (t *TxPublisher) createSweepTx(inputs []input.Input, changePkScript []byte,
+func (t *TxPublisher) createSweepTx(inputs []input.Input,
-	feeRate chainfee.SatPerKWeight) (*wire.MsgTx, btcutil.Amount, error) {
+	changePkScript lnwallet.AddrWithKey,
+	feeRate chainfee.SatPerKWeight) (*sweepTxCtx, error) {
 
 	// Validate and calculate the fee and change amount.
 	txFee, changeAmtOpt, locktimeOpt, err := prepareSweepTx(
-		inputs, changePkScript, feeRate, t.currentHeight.Load(),
+		inputs, changePkScript.DeliveryAddress, feeRate,
+		t.currentHeight.Load(),
 	)
 	if err != nil {
-		return nil, 0, err
+		return nil, err
 	}
 
 	var (
@@ -1185,7 +1222,7 @@ func (t *TxPublisher) createSweepTx(inputs []input.Input, changePkScript []byte,
 	// If there's a change amount, add it to the transaction.
 	changeAmtOpt.WhenSome(func(changeAmt btcutil.Amount) {
 		sweepTx.AddTxOut(&wire.TxOut{
-			PkScript: changePkScript,
+			PkScript: changePkScript.DeliveryAddress,
 			Value:    int64(changeAmt),
 		})
 	})
@@ -1196,7 +1233,7 @@ func (t *TxPublisher) createSweepTx(inputs []input.Input, changePkScript []byte,
 
 	prevInputFetcher, err := input.MultiPrevOutFetcher(inputs)
 	if err != nil {
-		return nil, 0, fmt.Errorf("error creating prev input fetcher "+
+		return nil, fmt.Errorf("error creating prev input fetcher "+
 			"for hash cache: %v", err)
 	}
 	hashCache := txscript.NewTxSigHashes(sweepTx, prevInputFetcher)
@@ -1224,14 +1261,17 @@ func (t *TxPublisher) createSweepTx(inputs []input.Input, changePkScript []byte,
 
 	for idx, inp := range idxs {
 		if err := addInputScript(idx, inp); err != nil {
-			return nil, 0, err
+			return nil, err
 		}
 	}
 
 	log.Debugf("Created sweep tx %v for inputs:\n%v", sweepTx.TxHash(),
 		inputTypeSummary(inputs))
 
-	return sweepTx, txFee, nil
+	return &sweepTxCtx{
+		tx:  sweepTx,
+		fee: txFee,
+	}, nil
 }
 
 // prepareSweepTx returns the tx fee, an optional change amount and an optional
@@ -1323,7 +1363,8 @@ func prepareSweepTx(inputs []input.Input, changePkScript []byte,
 	changeFloor := lnwallet.DustLimitForSize(len(changePkScript))
 
 	// If the change amount is dust, we'll move it into the fees.
-	if changeAmt < changeFloor {
+	switch {
+	case changeAmt < changeFloor:
 		log.Infof("Change amt %v below dustlimit %v, not adding "+
 			"change output", changeAmt, changeFloor)
 
@@ -1340,6 +1381,10 @@ func prepareSweepTx(inputs []input.Input, changePkScript []byte,
 
 		// Set the change amount to none.
 		changeAmtOpt = fn.None[btcutil.Amount]()
+
+	// Otherwise, we'll actually recognize it as a change output.
+	default:
+		// TODO(roasbeef): Implement (later commit in this PR).
 	}
 
 	// Optionally set the locktime.

sweep/fee_bumper_test.go

@@ -11,6 +11,7 @@ import (
 	"github.com/btcsuite/btcd/wire"
 	"github.com/btcsuite/btcwallet/chain"
 	"github.com/lightningnetwork/lnd/chainntnfs"
+	"github.com/lightningnetwork/lnd/fn"
 	"github.com/lightningnetwork/lnd/input"
 	"github.com/lightningnetwork/lnd/keychain"
 	"github.com/lightningnetwork/lnd/lnwallet"
@@ -21,12 +22,14 @@ import (
 
 var (
 	// Create  a taproot change script.
-	changePkScript = []byte{
+	changePkScript = lnwallet.AddrWithKey{
+		DeliveryAddress: []byte{
 			0x51, 0x20,
 			0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
 			0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
 			0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
 			0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
+		},
 	}
 
 	testInputCount atomic.Uint64
@@ -117,7 +120,9 @@ func TestCalcSweepTxWeight(t *testing.T) {
 	require.Zero(t, weight)
 
 	// Use a correct change script to test the success case.
-	weight, err = calcSweepTxWeight([]input.Input{&inp}, changePkScript)
+	weight, err = calcSweepTxWeight(
+		[]input.Input{&inp}, changePkScript.DeliveryAddress,
+	)
 	require.NoError(t, err)
 
 	// BaseTxSize 8 bytes
@@ -137,7 +142,9 @@ func TestBumpRequestMaxFeeRateAllowed(t *testing.T) {
 	inp := createTestInput(100, input.WitnessKeyHash)
 
 	// The weight is 487.
-	weight, err := calcSweepTxWeight([]input.Input{&inp}, changePkScript)
+	weight, err := calcSweepTxWeight(
+		[]input.Input{&inp}, changePkScript.DeliveryAddress,
+	)
 	require.NoError(t, err)
 
 	// Define a test budget and calculates its fee rate.
@@ -154,8 +161,10 @@ func TestBumpRequestMaxFeeRateAllowed(t *testing.T) {
 			// Use a wrong change script to test the error case.
 			name: "error calc weight",
 			req: &BumpRequest{
+				DeliveryAddress: lnwallet.AddrWithKey{
 					DeliveryAddress: []byte{1},
 				},
+			},
 			expectedMaxFeeRate: 0,
 			expectedErr:        true,
 		},
@@ -240,6 +249,7 @@ func TestInitializeFeeFunction(t *testing.T) {
 	// Create a publisher using the mocks.
 	tp := NewTxPublisher(TxPublisherConfig{
 		Estimator:  estimator,
+		AuxSweeper: fn.Some[AuxSweeper](&MockAuxSweeper{}),
 	})
 
 	// Create a test feerate.
@@ -304,7 +314,9 @@ func TestStoreRecord(t *testing.T) {
 	tx := &wire.MsgTx{}
 
 	// Create a publisher using the mocks.
-	tp := NewTxPublisher(TxPublisherConfig{})
+	tp := NewTxPublisher(TxPublisherConfig{
+		AuxSweeper: fn.Some[AuxSweeper](&MockAuxSweeper{}),
+	})
 
 	// Get the current counter and check it's increased later.
 	initialCounter := tp.requestCounter.Load()
@@ -373,6 +385,7 @@ func createTestPublisher(t *testing.T) (*TxPublisher, *mockers) {
 		Signer:     m.signer,
 		Wallet:     m.wallet,
 		Notifier:   m.notifier,
+		AuxSweeper: fn.Some[AuxSweeper](&MockAuxSweeper{}),
 	})
 
 	return tp, m
@@ -451,7 +464,7 @@ func TestCreateAndCheckTx(t *testing.T) {
 
 		t.Run(tc.name, func(t *testing.T) {
 			// Call the method under test.
-			_, _, err := tp.createAndCheckTx(tc.req, m.feeFunc)
+			_, err := tp.createAndCheckTx(tc.req, m.feeFunc)
 
 			// Check the result is as expected.
 			require.ErrorIs(t, err, tc.expectedErr)

sweep/interface.go

@@ -1,8 +1,12 @@
 package sweep
 
 import (
+	"github.com/btcsuite/btcd/btcutil"
 	"github.com/btcsuite/btcd/chaincfg/chainhash"
 	"github.com/btcsuite/btcd/wire"
+	"github.com/lightningnetwork/lnd/fn"
+	"github.com/lightningnetwork/lnd/input"
+	"github.com/lightningnetwork/lnd/keychain"
 	"github.com/lightningnetwork/lnd/lnwallet"
 )
 
@@ -57,3 +61,31 @@ type Wallet interface {
 	// service.
 	BackEnd() string
 }
+
+// SweepOutput is an output used to sweep funds from a channel output.
+type SweepOutput struct { //nolint:revive
+	wire.TxOut
+
+	// IsExtra indicates whether this output is an extra output that was
+	// added by a party other than the sweeper.
+	IsExtra bool
+
+	// InternalKey is the taproot internal key of the extra output. This is
+	// None, if this isn't a taproot output.
+	InternalKey fn.Option[keychain.KeyDescriptor]
+}
+
+// AuxSweeper is used to enable a 3rd party to further shape the sweeping
+// transaction by adding a set of extra outputs to the sweeping transaction.
+type AuxSweeper interface {
+	// DeriveSweepAddr takes a set of inputs, and the change address we'd
+	// use to sweep them, and maybe results an extra sweep output that we
+	// should add to the sweeping transaction.
+	DeriveSweepAddr(inputs []input.Input,
+		change lnwallet.AddrWithKey) fn.Result[SweepOutput]
+
+	// NotifyBroadcast is used to notify external callers of the broadcast
+	// of a sweep transaction, generated by the passed BumpRequest.
+	NotifyBroadcast(req *BumpRequest, tx *wire.MsgTx,
+		totalFees btcutil.Amount) error
+}

sweep/mock_test.go

@@ -314,3 +314,22 @@ func (m *MockFeeFunction) IncreaseFeeRate(confTarget uint32) (bool, error) {
 
 	return args.Bool(0), args.Error(1)
 }
+
+type MockAuxSweeper struct{}
+
+// DeriveSweepAddr takes a set of inputs, and the change address we'd
+// use to sweep them, and maybe results an extra sweep output that we
+// should add to the sweeping transaction.
+func (*MockAuxSweeper) DeriveSweepAddr(_ []input.Input,
+	_ lnwallet.AddrWithKey) fn.Result[SweepOutput] {
+
+	return fn.Ok(SweepOutput{})
+}
+
+// NotifyBroadcast is used to notify external callers of the broadcast
+// of a sweep transaction, generated by the passed BumpRequest.
+func (*MockAuxSweeper) NotifyBroadcast(_ *BumpRequest, _ *wire.MsgTx,
+	_ btcutil.Amount) error {
+
+	return nil
+}

sweep/sweeper.go

@@ -818,7 +818,7 @@ func (s *UtxoSweeper) sweep(set InputSet) error {
 		Inputs:          set.Inputs(),
 		Budget:          set.Budget(),
 		DeadlineHeight:  set.DeadlineHeight(),
-		DeliveryAddress: sweepAddr.DeliveryAddress,
+		DeliveryAddress: sweepAddr,
 		MaxFeeRate:      s.cfg.MaxFeeRate.FeePerKWeight(),
 		StartingFeeRate: set.StartingFeeRate(),
 		// TODO(yy): pass the strategy here.