sweep: Add selectUtxos to CraftSweepAllTx args

Author: Chinwendu20

rpcserver.go

@@ -1363,7 +1363,7 @@ func (r *rpcServer) SendCoins(ctx context.Context,
 		sweepTxPkg, err := sweep.CraftSweepAllTx(
 			feePerKw, maxFeeRate, uint32(bestHeight), nil,
 			targetAddr, wallet, wallet, wallet.WalletController,
-			r.server.cc.Signer, minConfs,
+			r.server.cc.Signer, minConfs, nil,
 		)
 		if err != nil {
 			return nil, err
@@ -1417,7 +1417,7 @@ func (r *rpcServer) SendCoins(ctx context.Context,
 				feePerKw, maxFeeRate, uint32(bestHeight),
 				outputs, targetAddr, wallet, wallet,
 				wallet.WalletController,
-				r.server.cc.Signer, minConfs,
+				r.server.cc.Signer, minConfs, nil,
 			)
 			if err != nil {
 				return nil, err

sweep/walletsweep.go

@@ -10,10 +10,12 @@ import (
 	"github.com/btcsuite/btcd/txscript"
 	"github.com/btcsuite/btcd/wire"
 	"github.com/btcsuite/btcwallet/wtxmgr"
+	"github.com/lightningnetwork/lnd/fn"
 	"github.com/lightningnetwork/lnd/input"
 	"github.com/lightningnetwork/lnd/lnwallet"
 	"github.com/lightningnetwork/lnd/lnwallet/chainfee"
 	"github.com/lightningnetwork/lnd/lnwallet/chanfunding"
+	"golang.org/x/exp/maps"
 )
 
 var (
@@ -24,6 +26,10 @@ var (
 	// ErrFeePreferenceConflict is returned when both a fee rate and a conf
 	// target is set for a fee preference.
 	ErrFeePreferenceConflict = errors.New("fee preference conflict")
+
+	// ErrUnknownUTXO is returned when creating a sweeping tx using an UTXO
+	// that's unknown to the wallet.
+	ErrUnknownUTXO = errors.New("unknown utxo")
 )
 
 // FeePreference defines an interface that allows the caller to specify how the
@@ -181,11 +187,11 @@ type OutputLeaser interface {
 }
 
 // WalletSweepPackage is a package that gives the caller the ability to sweep
-// ALL funds from a wallet in a single transaction. We also package a function
-// closure that allows one to abort the operation.
+// relevant funds from a wallet in a single transaction. We also package a
+// function closure that allows one to abort the operation.
 type WalletSweepPackage struct {
 	// SweepTx is a fully signed, and valid transaction that is broadcast,
-	// will sweep ALL confirmed coins in the wallet with a single
+	// will sweep ALL relevant confirmed coins in the wallet with a single
 	// transaction.
 	SweepTx *wire.MsgTx
 
@@ -208,27 +214,28 @@ type DeliveryAddr struct {
 }
 
 // CraftSweepAllTx attempts to craft a WalletSweepPackage which will allow the
-// caller to sweep ALL outputs within the wallet to a list of outputs. Any
-// leftover amount after these outputs and transaction fee, is sent to a single
-// output, as specified by the change address. The sweep transaction will be
-// crafted with the target fee rate, and will use the utxoSource and
-// outputLeaser as sources for wallet funds.
+// caller to sweep ALL funds in ALL or SELECT outputs within the wallet to a
+// list of outputs. Any leftover amount after these outputs and transaction fee,
+// is sent to a single output, as specified by the change address. The sweep
+// transaction will be crafted with the target fee rate, and will use the
+// utxoSource and outputLeaser as sources for wallet funds.
 func CraftSweepAllTx(feeRate, maxFeeRate chainfee.SatPerKWeight,
 	blockHeight uint32, deliveryAddrs []DeliveryAddr,
 	changeAddr btcutil.Address, coinSelectLocker CoinSelectionLocker,
 	utxoSource UtxoSource, outputLeaser OutputLeaser,
-	signer input.Signer, minConfs int32) (*WalletSweepPackage, error) {
+	signer input.Signer, minConfs int32,
+	selectUtxos fn.Set[wire.OutPoint]) (*WalletSweepPackage, error) {
 
 	// TODO(roasbeef): turn off ATPL as well when available?
 
-	var allOutputs []*lnwallet.Utxo
+	var outputsForSweep []*lnwallet.Utxo
 
 	// We'll make a function closure up front that allows us to unlock all
 	// selected outputs to ensure that they become available again in the
 	// case of an error after the outputs have been locked, but before we
 	// can actually craft a sweeping transaction.
 	unlockOutputs := func() {
-		for _, utxo := range allOutputs {
+		for _, utxo := range outputsForSweep {
 			// Log the error but continue since we're already
 			// handling an error.
 			err := outputLeaser.ReleaseOutput(
@@ -242,9 +249,9 @@ func CraftSweepAllTx(feeRate, maxFeeRate chainfee.SatPerKWeight,
 	}
 
 	// Next, we'll use the coinSelectLocker to ensure that no coin
-	// selection takes place while we fetch and lock all outputs the wallet
-	// knows of.  Otherwise, it may be possible for a new funding flow to
-	// lock an output while we fetch the set of unspent witnesses.
+	// selection takes place while we fetch and lock outputs in the
+	// wallet. Otherwise, it may be possible for a new funding flow to lock
+	// an output while we fetch the set of unspent witnesses.
 	err := coinSelectLocker.WithCoinSelectLock(func() error {
 		log.Trace("[WithCoinSelectLock] entered the lock")
 
@@ -260,6 +267,16 @@ func CraftSweepAllTx(feeRate, maxFeeRate chainfee.SatPerKWeight,
 
 		log.Trace("[WithCoinSelectLock] finished fetching UTXOs")
 
+		// Use select utxos, if provided.
+		if len(selectUtxos) > 0 {
+			utxos, err = fetchUtxosFromOutpoints(
+				utxos, selectUtxos.ToSlice(),
+			)
+			if err != nil {
+				return err
+			}
+		}
+
 		// We'll now lock each UTXO to ensure that other callers don't
 		// attempt to use these UTXOs in transactions while we're
 		// crafting out sweep all transaction.
@@ -278,7 +295,7 @@ func CraftSweepAllTx(feeRate, maxFeeRate chainfee.SatPerKWeight,
 
 		log.Trace("[WithCoinSelectLock] exited the lock")
 
-		allOutputs = append(allOutputs, utxos...)
+		outputsForSweep = append(outputsForSweep, utxos...)
 
 		return nil
 	})
@@ -287,15 +304,15 @@ func CraftSweepAllTx(feeRate, maxFeeRate chainfee.SatPerKWeight,
 		// in case we had any lingering outputs.
 		unlockOutputs()
 
-		return nil, fmt.Errorf("unable to fetch+lock wallet "+
-			"utxos: %v", err)
+		return nil, fmt.Errorf("unable to fetch+lock wallet utxos: %w",
+			err)
 	}
 
 	// Now that we've locked all the potential outputs to sweep, we'll
 	// assemble an input for each of them, so we can hand it off to the
 	// sweeper to generate and sign a transaction for us.
 	var inputsToSweep []input.Input
-	for _, output := range allOutputs {
+	for _, output := range outputsForSweep {
 		// As we'll be signing for outputs under control of the wallet,
 		// we only need to populate the output value and output script.
 		// The rest of the items will be populated internally within
@@ -390,3 +407,24 @@ func CraftSweepAllTx(feeRate, maxFeeRate chainfee.SatPerKWeight,
 		CancelSweepAttempt: unlockOutputs,
 	}, nil
 }
+
+// fetchUtxosFromOutpoints returns UTXOs for given outpoints. Errors if any
+// outpoint is not in the passed slice of utxos.
+func fetchUtxosFromOutpoints(utxos []*lnwallet.Utxo,
+	outpoints []wire.OutPoint) ([]*lnwallet.Utxo, error) {
+
+	lookup := fn.SliceToMap(utxos, func(utxo *lnwallet.Utxo) wire.OutPoint {
+		return utxo.OutPoint
+	}, func(utxo *lnwallet.Utxo) *lnwallet.Utxo {
+		return utxo
+	})
+
+	subMap, err := fn.NewSubMap(lookup, outpoints)
+	if err != nil {
+		return nil, fmt.Errorf("%w: %v", ErrUnknownUTXO, err.Error())
+	}
+
+	fetchedUtxos := maps.Values(subMap)
+
+	return fetchedUtxos, nil
+}

sweep/walletsweep_test.go

@@ -12,6 +12,7 @@ import (
 	"github.com/btcsuite/btcd/txscript"
 	"github.com/btcsuite/btcd/wire"
 	"github.com/btcsuite/btcwallet/wtxmgr"
+	"github.com/lightningnetwork/lnd/fn"
 	"github.com/lightningnetwork/lnd/lntest/mock"
 	"github.com/lightningnetwork/lnd/lnwallet"
 	"github.com/lightningnetwork/lnd/lnwallet/chainfee"
@@ -339,7 +340,7 @@ func TestCraftSweepAllTxCoinSelectFail(t *testing.T) {
 
 	_, err := CraftSweepAllTx(
 		0, 0, 10, nil, nil, coinSelectLocker, utxoSource, utxoLeaser,
-		nil, 0,
+		nil, 0, nil,
 	)
 
 	// Since we instructed the coin select locker to fail above, we should
@@ -365,7 +366,7 @@ func TestCraftSweepAllTxUnknownWitnessType(t *testing.T) {
 
 	_, err := CraftSweepAllTx(
 		0, 0, 10, nil, nil, coinSelectLocker, utxoSource, utxoLeaser,
-		nil, 0,
+		nil, 0, nil,
 	)
 
 	// Since passed in a p2wsh output, which is unknown, we should fail to
@@ -399,7 +400,7 @@ func TestCraftSweepAllTx(t *testing.T) {
 
 	sweepPkg, err := CraftSweepAllTx(
 		0, 0, 10, nil, deliveryAddr, coinSelectLocker, utxoSource,
-		utxoLeaser, signer, 0,
+		utxoLeaser, signer, 0, nil,
 	)
 	require.NoError(t, err, "unable to make sweep tx")
 
@@ -440,3 +441,58 @@ func TestCraftSweepAllTx(t *testing.T) {
 	sweepPkg.CancelSweepAttempt()
 	assertUtxosReleased(t, utxoLeaser, testUtxos[:2])
 }
+
+// TestCraftSweepAllTxWithSelectedUTXO tests that we'll properly lock the
+// selected outputs within the wallet, and craft a single sweep transaction
+// that pays to the target output.
+func TestCraftSweepAllTxWithSelectedUTXO(t *testing.T) {
+	t.Parallel()
+
+	// First, we'll make a mock signer along with a fee estimator, We'll
+	// use zero fees to we can assert a precise output value.
+	signer := &mock.DummySigner{}
+
+	// Grab the first UTXO from the test UTXOs.
+	utxo1 := testUtxos[0]
+	utxoSource := newMockUtxoSource([]*lnwallet.Utxo{utxo1})
+	coinSelectLocker := &mockCoinSelectionLocker{}
+	utxoLeaser := newMockOutputLeaser()
+
+	// Create an unknown utxo.
+	outpointUknown := wire.OutPoint{Index: 4}
+
+	// Sweep using the uknnown utxo and expect an error.
+	sweepPkg, err := CraftSweepAllTx(
+		0, 0, 10, nil, deliveryAddr, coinSelectLocker, utxoSource,
+		utxoLeaser, signer, 0, fn.NewSet(outpointUknown),
+	)
+	require.ErrorIs(t, err, ErrUnknownUTXO)
+	require.Nil(t, sweepPkg)
+
+	// Sweep again using the known utxo and expect no error.
+	sweepPkg, err = CraftSweepAllTx(
+		0, 0, 10, nil, deliveryAddr, coinSelectLocker, utxoSource,
+		utxoLeaser, signer, 0, fn.NewSet(utxo1.OutPoint),
+	)
+	require.NoError(t, err)
+
+	// At this point utxo1 should be locked.
+	assertUtxosLeased(t, utxoLeaser, []*lnwallet.Utxo{utxo1})
+	assertNoUtxosReleased(t, utxoLeaser, []*lnwallet.Utxo{utxo1})
+
+	// Validate the sweeping tx has the expected shape.
+	sweepTx := sweepPkg.SweepTx
+	require.Len(t, sweepTx.TxIn, 1)
+	require.Len(t, sweepTx.TxOut, 1)
+
+	// We should have a single output that pays to our sweep script
+	// generated above.
+	expectedSweepValue := utxo1.Value
+	require.Equal(t, sweepScript, sweepTx.TxOut[0].PkScript)
+	require.EqualValues(t, expectedSweepValue, sweepTx.TxOut[0].Value)
+
+	// If we cancel the sweep attempt, then we should find utxo1 to be
+	// unlocked.
+	sweepPkg.CancelSweepAttempt()
+	assertUtxosReleased(t, utxoLeaser, []*lnwallet.Utxo{utxo1})
+}