sweepbatcher: consider change in presigning and batch tx

Presigning sweeps takes change outputs into account.
Each sweep belonging to the same sweep group points
to the same change output, if existent.
sweepbatcher.presign scans all passed sweeps for
change outputs and passes them to constructUnsignedTx.
Optional change of a swap is encoded in its sweeps
as a pointer to the same change output. This change
is taken into account when constructing the unsigned
batch transaction when it comes to tx weight and
outputs.

Author: hieblmi

sweepbatcher/presigned.go

@@ -51,6 +51,7 @@ func ensurePresigned(ctx context.Context, newSweeps []*sweep,
 			outpoint:  s.outpoint,
 			value:     s.value,
 			presigned: s.presigned,
+			change:    s.change,
 		}
 	}
 
@@ -66,14 +67,20 @@ func ensurePresigned(ctx context.Context, newSweeps []*sweep,
 		return fmt.Errorf("failed to find destination address: %w", err)
 	}
 
+	// Get the change outputs for each sweep group.
+	changeOutputs, err := getChangeOutputs(sweeps, chainParams)
+	if err != nil {
+		return fmt.Errorf("failed to get change outputs: %w", err)
+	}
+
 	// Set LockTime to 0. It is not critical.
 	const currentHeight = 0
 
 	// Check if we can sign with minimum fee rate.
 	const feeRate = chainfee.FeePerKwFloor
 
 	tx, _, _, _, err := constructUnsignedTx(
-		sweeps, destAddr, currentHeight, feeRate,
+		sweeps, destAddr, changeOutputs, currentHeight, feeRate,
 	)
 	if err != nil {
 		return fmt.Errorf("failed to construct unsigned tx "+
@@ -257,7 +264,7 @@ func (b *batch) presign(ctx context.Context, newSweeps []*sweep) error {
 
 		err = presign(
 			ctx, b.cfg.presignedHelper, destAddr, primarySweepID,
-			sweeps, nextBlockFeeRate,
+			sweeps, nextBlockFeeRate, b.cfg.chainParams,
 		)
 		if err != nil {
 			return fmt.Errorf("failed to presign a transaction "+
@@ -299,7 +306,8 @@ type presigner interface {
 // 10x of the current next block feerate.
 func presign(ctx context.Context, presigner presigner, destAddr btcutil.Address,
 	primarySweepID wire.OutPoint, sweeps []sweep,
-	nextBlockFeeRate chainfee.SatPerKWeight) error {
+	nextBlockFeeRate chainfee.SatPerKWeight,
+	chainParams *chaincfg.Params) error {
 
 	if presigner == nil {
 		return fmt.Errorf("presigner is not installed")
@@ -328,6 +336,12 @@ func presign(ctx context.Context, presigner presigner, destAddr btcutil.Address,
 		return fmt.Errorf("timeout is invalid: %d", timeout)
 	}
 
+	// Get the change outputs of each sweep group.
+	changeOutputs, err := getChangeOutputs(sweeps, chainParams)
+	if err != nil {
+		return fmt.Errorf("failed to get change outputs: %w", err)
+	}
+
 	// Go from the floor (1.01 sat/vbyte) to 2k sat/vbyte with step of 1.2x.
 	const (
 		start            = chainfee.FeePerKwFloor
@@ -353,7 +367,7 @@ func presign(ctx context.Context, presigner presigner, destAddr btcutil.Address,
 	for fr := start; fr <= stop; fr = (fr * factorPPM) / 1_000_000 {
 		// Construct an unsigned transaction for this fee rate.
 		tx, _, feeForWeight, fee, err := constructUnsignedTx(
-			sweeps, destAddr, currentHeight, fr,
+			sweeps, destAddr, changeOutputs, currentHeight, fr,
 		)
 		if err != nil {
 			return fmt.Errorf("failed to construct unsigned tx "+
@@ -438,9 +452,15 @@ func (b *batch) publishPresigned(ctx context.Context) (btcutil.Amount, error,
 			err), false
 	}
 
+	changeOutputs, err := getChangeOutputs(sweeps, b.cfg.chainParams)
+	if err != nil {
+		return 0, fmt.Errorf("failed to get change outputs: %w", err),
+			false
+	}
+
 	// Construct unsigned batch transaction.
 	tx, weight, _, fee, err := constructUnsignedTx(
-		sweeps, address, currentHeight, feeRate,
+		sweeps, address, changeOutputs, currentHeight, feeRate,
 	)
 	if err != nil {
 		return 0, fmt.Errorf("failed to construct tx: %w", err),
@@ -493,10 +513,12 @@ func (b *batch) publishPresigned(ctx context.Context) (btcutil.Amount, error,
 	signedFeeRate := chainfee.NewSatPerKWeight(fee, realWeight)
 
 	numSweeps := len(tx.TxIn)
+	numChange := len(tx.TxOut) - 1
 	b.Infof("attempting to publish custom signed tx=%v, desiredFeerate=%v,"+
-		" signedFeeRate=%v, weight=%v, fee=%v, sweeps=%d, destAddr=%s",
+		" signedFeeRate=%v, weight=%v, fee=%v, sweeps=%d, "+
+		"changeOutputs=%d, destAddr=%s",
 		txHash, feeRate, signedFeeRate, realWeight, fee, numSweeps,
-		address)
+		numChange, address)
 	b.debugLogTx("serialized batch", tx)
 
 	// Publish the transaction.
@@ -557,6 +579,46 @@ func getPresignedSweepsDestAddr(ctx context.Context, helper destPkScripter,
 	return address, nil
 }
 
+// getChangeOutputs retrieves the change output references of each sweep and
+// de-duplicates them. The function must be used in presigned mode only.
+func getChangeOutputs(sweeps []sweep, chainParams *chaincfg.Params) (
+	map[*wire.TxOut]btcutil.Address, error) {
+
+	changeOutputs := make(map[*wire.TxOut]btcutil.Address)
+	for _, sweep := range sweeps {
+		// If the sweep has a change output, add it to the changeOutputs
+		// map to avoid duplicates.
+		if sweep.change == nil {
+			continue
+		}
+
+		// If the change output is already in the map, skip it.
+		if _, exists := changeOutputs[sweep.change]; exists {
+			continue
+		}
+
+		// Convert the change output's pkScript to an address.
+		changePkScript, err := txscript.ParsePkScript(
+			sweep.change.PkScript,
+		)
+		if err != nil {
+			return nil, fmt.Errorf("failed to parse change "+
+				"output pkScript: %w", err)
+		}
+
+		address, err := changePkScript.Address(chainParams)
+		if err != nil {
+			return nil, fmt.Errorf("pkScript.Address failed for "+
+				"pkScript %x returned for change output: %w",
+				sweep.change.PkScript, err)
+		}
+
+		changeOutputs[sweep.change] = address
+	}
+
+	return changeOutputs, nil
+}
+
 // CheckSignedTx makes sure that signedTx matches the unsignedTx. It checks
 // according to criteria specified in the description of PresignedHelper.SignTx.
 func CheckSignedTx(unsignedTx, signedTx *wire.MsgTx, inputAmt btcutil.Amount,
@@ -593,23 +655,31 @@ func CheckSignedTx(unsignedTx, signedTx *wire.MsgTx, inputAmt btcutil.Amount,
 	}
 
 	// Compare outputs.
-	if len(unsignedTx.TxOut) != 1 {
-		return fmt.Errorf("unsigned tx has %d outputs, want 1",
-			len(unsignedTx.TxOut))
-	}
-	if len(signedTx.TxOut) != 1 {
-		return fmt.Errorf("the signed tx has %d outputs, want 1",
+	if len(unsignedTx.TxOut) != len(signedTx.TxOut) {
+		return fmt.Errorf("unsigned tx has %d outputs, signed tx has "+
+			"%d outputs, should be equal", len(unsignedTx.TxOut),
 			len(signedTx.TxOut))
 	}
-	unsignedOut := unsignedTx.TxOut[0]
-	signedOut := signedTx.TxOut[0]
-	if !bytes.Equal(unsignedOut.PkScript, signedOut.PkScript) {
-		return fmt.Errorf("mismatch of output pkScript: %x, %x",
-			unsignedOut.PkScript, signedOut.PkScript)
+	for i, o := range unsignedTx.TxOut {
+		if !bytes.Equal(o.PkScript, signedTx.TxOut[i].PkScript) {
+			return fmt.Errorf("mismatch of output pkScript: %x, %x",
+				o.PkScript, signedTx.TxOut[i].PkScript)
+		}
+	}
+
+	// The first output is always the batch output.
+	// TODO(hieblmi): ensure this.
+	batchOutput := signedTx.TxOut[0]
+
+	// Calculate the total value of change outputs to help determine the
+	// transaction fee.
+	totalChangeValue := btcutil.Amount(0)
+	for i := 1; i < len(signedTx.TxOut); i++ {
+		totalChangeValue += btcutil.Amount(signedTx.TxOut[i].Value)
 	}
 
 	// Find the feerate of signedTx.
-	fee := inputAmt - btcutil.Amount(signedOut.Value)
+	fee := inputAmt - btcutil.Amount(batchOutput.Value) - totalChangeValue
 	weight := lntypes.WeightUnit(
 		blockchain.GetTransactionWeight(btcutil.NewTx(signedTx)),
 	)

sweepbatcher/presigned_test.go

@@ -1011,6 +1011,7 @@ func TestPresign(t *testing.T) {
 				ctx, tc.presigner, tc.destAddr,
 				tc.primarySweepID, tc.sweeps,
 				tc.nextBlockFeeRate,
+				&chaincfg.RegressionNetParams,
 			)
 			if tc.wantErr != "" {
 				require.Error(t, err)
@@ -1460,7 +1461,8 @@ func TestCheckSignedTx(t *testing.T) {
 			},
 			inputAmt:    3_000_000,
 			minRelayFee: 253,
-			wantErr:     "unsigned tx has 2 outputs, want 1",
+			wantErr: "unsigned tx has 2 outputs, signed tx " +
+				"has 1 outputs, should be equal",
 		},
 
 		{
@@ -1517,7 +1519,8 @@ func TestCheckSignedTx(t *testing.T) {
 			},
 			inputAmt:    3_000_000,
 			minRelayFee: 253,
-			wantErr:     "the signed tx has 2 outputs, want 1",
+			wantErr: "unsigned tx has 1 outputs, signed tx " +
+				"has 2 outputs, should be equal",
 		},
 
 		{
@@ -1642,3 +1645,124 @@ func TestCheckSignedTx(t *testing.T) {
 		})
 	}
 }
+
+// TestGetChangeOutputs tests that the change aggregation across sweeps works as
+// intended. Each sweep of a sweep group should have a pointer to the same
+// change output which is aggregated in getChangeOutput.
+func TestGetChangeOutputs(t *testing.T) {
+	// Prepare the necessary data for test cases.
+	op1 := wire.OutPoint{
+		Hash:  chainhash.Hash{1, 1, 1},
+		Index: 1,
+	}
+	op2 := wire.OutPoint{
+		Hash:  chainhash.Hash{2, 2, 2},
+		Index: 2,
+	}
+	op3 := wire.OutPoint{
+		Hash:  chainhash.Hash{3, 3, 3},
+		Index: 3,
+	}
+
+	batchPkScript, err := txscript.PayToAddrScript(destAddr)
+	require.NoError(t, err)
+
+	changeOutput1 := &wire.TxOut{
+		Value:    100_000,
+		PkScript: batchPkScript,
+	}
+	changeOutput2 := &wire.TxOut{
+		Value:    200_000,
+		PkScript: batchPkScript,
+	}
+
+	cases := []struct {
+		name        string
+		sweeps      []sweep
+		wantOutputs map[*wire.TxOut]btcutil.Address
+		wantErr     string
+	}{
+		{
+			name: "no change",
+			sweeps: []sweep{
+				{
+					outpoint: op1,
+					value:    1_000_000,
+					change:   nil,
+				},
+			},
+			wantOutputs: map[*wire.TxOut]btcutil.Address{},
+		},
+		{
+			name: "single sweep, single change",
+			sweeps: []sweep{
+				{
+					outpoint: op1,
+					value:    1_000_000,
+					change:   changeOutput1,
+				},
+			},
+			wantOutputs: map[*wire.TxOut]btcutil.Address{
+				changeOutput1: destAddr,
+			},
+		},
+		{
+			name: "double sweep, single change",
+			sweeps: []sweep{
+				{
+					outpoint: op1,
+					value:    1_000_000,
+					change:   changeOutput1,
+				},
+				{
+					outpoint: op2,
+					value:    1_000_000,
+					change:   changeOutput1,
+				},
+			},
+			wantOutputs: map[*wire.TxOut]btcutil.Address{
+				changeOutput1: destAddr,
+			},
+		},
+		{
+			name: "double sweep, double change",
+			sweeps: []sweep{
+				{
+					outpoint: op1,
+					value:    1_000_000,
+					change:   changeOutput1,
+				},
+				{
+					outpoint: op2,
+					value:    1_000_000,
+					change:   changeOutput1,
+				},
+				{
+					outpoint: op3,
+					value:    1_000_000,
+					change:   changeOutput2,
+				},
+			},
+			wantOutputs: map[*wire.TxOut]btcutil.Address{
+				changeOutput1: destAddr,
+				changeOutput2: destAddr,
+			},
+		},
+	}
+
+	for _, tc := range cases {
+		t.Run(tc.name, func(t *testing.T) {
+			changeOutputs, err := getChangeOutputs(
+				tc.sweeps, &chaincfg.RegressionNetParams,
+			)
+			if tc.wantErr != "" {
+				require.Error(t, err)
+				require.ErrorContains(t, err, tc.wantErr)
+			} else {
+				require.NoError(t, err)
+			}
+
+			require.Equal(t, tc.wantOutputs, changeOutputs)
+		})
+	}
+}