sweepbatcher: fix OnChainFeePortion values

There were two mistakes.

In case of a swap with multiple sweeps only the fee of the first sweep
of a swap was accounted.

Rounding diff (the remainder) was attributed to all the sweeps rather than to
the first (primary) sweep of the batch. The sweep to attribute the remainder
was chosen by comparing SignatureScript which is always empty. New approach is
to find the primary sweep and to compare its outpoint directly.

Author: starius

sweepbatcher/sweep_batch.go

@@ -1917,12 +1917,12 @@ func getFeePortionForSweep(spendTx *wire.MsgTx, numSweeps int,
 }
 
 // getFeePortionPaidBySweep returns the fee portion that the sweep should pay
-// for the batch transaction. If the sweep is the first sweep in the batch, it
+// for the batch transaction. If the sweep is the primary sweep in the batch, it
 // pays the rounding difference.
 func getFeePortionPaidBySweep(spendTx *wire.MsgTx, feePortionPerSweep,
-	roundingDiff btcutil.Amount, sweep *sweep) btcutil.Amount {
+	roundingDiff btcutil.Amount, primary bool) btcutil.Amount {
 
-	if bytes.Equal(spendTx.TxIn[0].SignatureScript, sweep.htlc.SigScript) {
+	if primary {
 		return feePortionPerSweep + roundingDiff
 	}
 
@@ -1974,22 +1974,43 @@ func (b *batch) handleSpend(ctx context.Context, spendTx *wire.MsgTx) error {
 		spendTx, len(notifyList), totalSweptAmt,
 	)
 
+	// Calculate fees per swaps. Only the first sweep in a swap has a
+	// notifier, so we calculate total fee per swap and send it to a sweep
+	// having that swap and a notifier.
+	swap2fee := make(map[lntypes.Hash]btcutil.Amount)
+	for _, sweep := range notifyList {
+		primary := sweep.outpoint == b.primarySweepID
+
+		swap2fee[sweep.swapHash] += getFeePortionPaidBySweep(
+			spendTx, feePortionPaidPerSweep,
+			roundingDifference, primary,
+		)
+	}
+
+	// Now send notifications to notifiers.
 	for _, sweep := range notifyList {
 		// If the sweep's notifier is empty then this means that a swap
-		// is not waiting to read an update from it, so we can skip
-		// the notification part.
+		// is not waiting to read an update from it or this is not the
+		// first sweep in a swap, so we can skip the notification part.
 		if sweep.notifier == nil ||
 			*sweep.notifier == (SpendNotifier{}) {
 
 			continue
 		}
 
+		// Make sure there is only one sweep with a notifier per swap
+		// hash, otherwise our fee calculation is incorrect.
+		fee, has := swap2fee[sweep.swapHash]
+		if !has {
+			return fmt.Errorf("no fee for swap %v; maybe "+
+				"multiple sweeps with a notifier per swap?",
+				sweep.swapHash)
+		}
+		delete(swap2fee, sweep.swapHash)
+
 		spendDetail := SpendDetail{
-			Tx: spendTx,
-			OnChainFeePortion: getFeePortionPaidBySweep(
-				spendTx, feePortionPaidPerSweep,
-				roundingDifference, &sweep,
-			),
+			Tx:                spendTx,
+			OnChainFeePortion: fee,
 		}
 
 		// Dispatch the sweep notifier, we don't care about the outcome

sweepbatcher/sweep_batcher.go

@@ -778,7 +778,7 @@ func (b *Batcher) handleSweeps(ctx context.Context, sweeps []*sweep,
 		// and we won't be able to attach this sweep to it.
 		if parentBatch.Confirmed {
 			return b.monitorSpendAndNotify(
-				ctx, sweep, parentBatch.ID, notifier,
+				ctx, sweeps, parentBatch.ID, notifier,
 			)
 		}
 	}
@@ -1061,7 +1061,7 @@ func (b *Batcher) FetchUnconfirmedBatches(ctx context.Context) ([]*batch,
 // the response back to the response channel. It is called if the batch is fully
 // confirmed and we just need to deliver the data back to the caller though
 // SpendNotifier.
-func (b *Batcher) monitorSpendAndNotify(ctx context.Context, sweep *sweep,
+func (b *Batcher) monitorSpendAndNotify(ctx context.Context, sweeps []*sweep,
 	parentBatchID int32, notifier *SpendNotifier) error {
 
 	// If the caller has not provided a notifier, stop.
@@ -1079,6 +1079,17 @@ func (b *Batcher) monitorSpendAndNotify(ctx context.Context, sweep *sweep,
 		return err
 	}
 
+	// Find the primarySweepID.
+	dbSweeps, err := b.store.FetchBatchSweeps(ctx, parentBatchID)
+	if err != nil {
+		cancel()
+
+		return err
+	}
+	primarySweepID := dbSweeps[0].Outpoint
+
+	sweep := sweeps[0]
+
 	spendChan, spendErr, err := b.chainNotifier.RegisterSpendNtfn(
 		spendCtx, &sweep.outpoint, sweep.htlc.PkScript,
 		sweep.initiationHeight,
@@ -1099,6 +1110,7 @@ func (b *Batcher) monitorSpendAndNotify(ctx context.Context, sweep *sweep,
 		select {
 		case spend := <-spendChan:
 			spendTx := spend.SpendingTx
+
 			// Calculate the fee portion that each sweep should pay
 			// for the batch.
 			feePortionPerSweep, roundingDifference :=
@@ -1107,17 +1119,23 @@ func (b *Batcher) monitorSpendAndNotify(ctx context.Context, sweep *sweep,
 					totalSwept,
 				)
 
-			onChainFeePortion := getFeePortionPaidBySweep(
-				spendTx, feePortionPerSweep,
-				roundingDifference, sweep,
-			)
+			// Add onchain fee accross all the sweeps of the swap.
+			var fee btcutil.Amount
+			for _, s := range sweeps {
+				isFirst := s.outpoint == primarySweepID
+
+				fee += getFeePortionPaidBySweep(
+					spendTx, feePortionPerSweep,
+					roundingDifference, isFirst,
+				)
+			}
 
 			// Notify the requester of the spend with the spend
 			// details, including the fee portion for this
 			// particular sweep.
 			spendDetail := &SpendDetail{
 				Tx:                spendTx,
-				OnChainFeePortion: onChainFeePortion,
+				OnChainFeePortion: fee,
 			}
 
 			select {

sweepbatcher/sweep_batcher_presigned_test.go

@@ -1341,7 +1341,12 @@ func testPresigned_purging(t *testing.T, numSwaps, numConfirmedSwaps int,
 
 	require.LessOrEqual(t, numConfirmedSwaps, numSwaps)
 
-	const sweepsPerSwap = 2
+	const (
+		sweepsPerSwap = 2
+		feeRate       = chainfee.SatPerKWeight(10_000)
+		swapAmount    = 3_000_001
+	)
+	sweepAmounts := []btcutil.Amount{1_000_001, 2_000_000}
 
 	lnd := test.NewMockLnd()
 
@@ -1351,7 +1356,7 @@ func testPresigned_purging(t *testing.T, numSwaps, numConfirmedSwaps int,
 	customFeeRate := func(_ context.Context,
 		_ lntypes.Hash) (chainfee.SatPerKWeight, error) {
 
-		return chainfee.SatPerKWeight(10_000), nil
+		return feeRate, nil
 	}
 
 	presignedHelper := newMockPresignedHelper()
@@ -1369,12 +1374,17 @@ func testPresigned_purging(t *testing.T, numSwaps, numConfirmedSwaps int,
 		checkBatcherError(t, err)
 	}()
 
+	swapHashes := make([]lntypes.Hash, numSwaps)
+	groups := make([][]Input, numSwaps)
 	txs := make([]*wire.MsgTx, numSwaps)
 	allOps := make([]wire.OutPoint, 0, numSwaps*sweepsPerSwap)
+	spendChans := make([]<-chan *SpendDetail, numSwaps)
+	confChans := make([]<-chan *chainntnfs.TxConfirmation, numSwaps)
 
 	for i := range numSwaps {
 		// Create a swap of sweepsPerSwap sweeps.
 		swapHash := lntypes.Hash{byte(i + 1)}
+		swapHashes[i] = swapHash
 		ops := make([]wire.OutPoint, sweepsPerSwap)
 		group := make([]Input, sweepsPerSwap)
 		for j := range sweepsPerSwap {
@@ -1386,15 +1396,16 @@ func testPresigned_purging(t *testing.T, numSwaps, numConfirmedSwaps int,
 
 			group[j] = Input{
 				Outpoint: ops[j],
-				Value:    btcutil.Amount(1_000_000 * (j + 1)),
+				Value:    sweepAmounts[j],
 			}
 		}
+		groups[i] = group
 
 		// Create a swap in DB.
 		swap := &loopdb.LoopOutContract{
 			SwapContract: loopdb.SwapContract{
 				CltvExpiry:      111,
-				AmountRequested: 3_000_000,
+				AmountRequested: swapAmount,
 				ProtocolVersion: loopdb.ProtocolVersionMuSig2,
 				HtlcKeys:        htlcKeys,
 
@@ -1421,11 +1432,24 @@ func testPresigned_purging(t *testing.T, numSwaps, numConfirmedSwaps int,
 		)
 		require.NoError(t, err)
 
+		// Create a spending notification channel.
+		spendChan := make(chan *SpendDetail, 1)
+		spendChans[i] = spendChan
+		confChan := make(chan *chainntnfs.TxConfirmation, 1)
+		confChans[i] = confChan
+		notifier := &SpendNotifier{
+			SpendChan:    spendChan,
+			SpendErrChan: make(chan error, 1),
+			ConfChan:     confChan,
+			ConfErrChan:  make(chan error, 1),
+			QuitChan:     make(chan bool, 1),
+		}
+
 		// Add the sweep, triggering the publish attempt.
 		require.NoError(t, batcher.AddSweep(&SweepRequest{
 			SwapHash: swapHash,
 			Inputs:   group,
-			Notifier: &dummyNotifier,
+			Notifier: notifier,
 		}))
 
 		// For the first group it should register for the sweep's spend
@@ -1463,6 +1487,34 @@ func testPresigned_purging(t *testing.T, numSwaps, numConfirmedSwaps int,
 		SpendingHeight:    int32(601 + numSwaps + 1),
 	}
 	lnd.SpendChannel <- spendDetail
+
+	// Calculate the expected on-chain fee of the swap.
+	wantFee := make([]btcutil.Amount, numConfirmedSwaps)
+	for i := range numConfirmedSwaps {
+		batchAmount := swapAmount * btcutil.Amount(numConfirmedSwaps)
+		txFee := batchAmount - btcutil.Amount(tx.TxOut[0].Value)
+		numConfirmedSweeps := numConfirmedSwaps * sweepsPerSwap
+		feePerSweep := txFee / btcutil.Amount(numConfirmedSweeps)
+		roundingDiff := txFee - feePerSweep*btcutil.Amount(
+			numConfirmedSweeps,
+		)
+		swapFee := feePerSweep * 2
+
+		// Add rounding difference to the first swap.
+		if i == 0 {
+			swapFee += roundingDiff
+		}
+
+		wantFee[i] = swapFee
+	}
+
+	// Make sure that notifiers of confirmed sweeps received notifications.
+	for i := range numConfirmedSwaps {
+		spend := <-spendChans[i]
+		require.Equal(t, txHash, spend.Tx.TxHash())
+		require.Equal(t, wantFee[i], spend.OnChainFeePortion)
+	}
+
 	<-lnd.RegisterConfChannel
 	require.NoError(t, lnd.NotifyHeight(
 		int32(601+numSwaps+1+batchConfHeight),
@@ -1474,16 +1526,61 @@ func testPresigned_purging(t *testing.T, numSwaps, numConfirmedSwaps int,
 	// CleanupTransactions is called here.
 	<-presignedHelper.cleanupCalled
 
-	// If all the swaps were confirmed, stop.
-	if numConfirmedSwaps == numSwaps {
-		return
+	// Make sure that notifiers of confirmed sweeps received notifications.
+	for i := range numConfirmedSwaps {
+		conf := <-confChans[i]
+		require.Equal(t, txHash, conf.Tx.TxHash())
 	}
 
 	// Missing sweeps in the confirmed transaction should be re-added to the
 	// batcher as new batch. The groups are added incrementally, so we need
 	// to wait until the batch reaches the expected size.
-	<-lnd.RegisterSpendChannel
-	<-lnd.TxPublishChannel
+	if numConfirmedSwaps != numSwaps {
+		<-lnd.RegisterSpendChannel
+		<-lnd.TxPublishChannel
+	}
+
+	// Now make sure that a correct spenf and conf contification is sent if
+	// AddSweep is called after confirming the sweeps.
+	for i := range numConfirmedSwaps {
+		// Create a spending notification channel.
+		spendChan := make(chan *SpendDetail, 1)
+		confChan := make(chan *chainntnfs.TxConfirmation)
+		notifier := &SpendNotifier{
+			SpendChan:    spendChan,
+			SpendErrChan: make(chan error, 1),
+			ConfChan:     confChan,
+			ConfErrChan:  make(chan error, 1),
+			QuitChan:     make(chan bool, 1),
+		}
+
+		// Add the sweep, triggering the publish attempt.
+		require.NoError(t, batcher.AddSweep(&SweepRequest{
+			SwapHash: swapHashes[i],
+			Inputs:   groups[i],
+			Notifier: notifier,
+		}))
+
+		spendReg := <-lnd.RegisterSpendChannel
+		spendReg.SpendChannel <- spendDetail
+
+		spend := <-spendChan
+		require.Equal(t, txHash, spend.Tx.TxHash())
+		require.Equal(t, wantFee[i], spend.OnChainFeePortion)
+
+		<-lnd.RegisterConfChannel
+		lnd.ConfChannel <- &chainntnfs.TxConfirmation{
+			Tx: tx,
+		}
+
+		conf := <-confChan
+		require.Equal(t, tx.TxHash(), conf.Tx.TxHash())
+	}
+
+	// If all the swaps were confirmed, stop.
+	if numConfirmedSwaps == numSwaps {
+		return
+	}
 
 	// Wait to new batch to appear and to have the expected size.
 	wantSize := (numSwaps - numConfirmedSwaps) * sweepsPerSwap
@@ -1575,5 +1672,6 @@ func TestPresigned(t *testing.T) {
 		testPurging(3, 1)
 		testPurging(3, 2)
 		testPurging(5, 2)
+		testPurging(5, 3)
 	})
 }