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.

sweepbatcher/sweep_batch.go

@@ -1927,12 +1927,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 {
+func getFeePortionPaidBySweep(feePortionPerSweep, roundingDiff btcutil.Amount,
+	primary bool) btcutil.Amount {
 
-	if bytes.Equal(spendTx.TxIn[0].SignatureScript, sweep.htlc.SigScript) {
+	if primary {
 		return feePortionPerSweep + roundingDiff
 	}
 
@@ -1983,22 +1983,42 @@ 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(
+			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
@@ -2193,6 +2213,18 @@ func (b *batch) handleConf(ctx context.Context,
 		spendTx, len(b.sweeps), 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 b.sweeps {
+		primary := sweep.outpoint == b.primarySweepID
+
+		swap2fee[sweep.swapHash] += getFeePortionPaidBySweep(
+			feePortionPaidPerSweep, roundingDifference, primary,
+		)
+	}
+
 	// Send the confirmation to all the notifiers.
 	for _, s := range b.sweeps {
 		// If the sweep's notifier is empty then this means that
@@ -2202,12 +2234,19 @@ func (b *batch) handleConf(ctx context.Context,
 			continue
 		}
 
+		// Make sure there is only one sweep with a notifier per swap
+		// hash, otherwise our fee calculation is incorrect.
+		fee, has := swap2fee[s.swapHash]
+		if !has {
+			return fmt.Errorf("no fee for swap %v; maybe "+
+				"multiple sweeps with a notifier per swap?",
+				s.swapHash)
+		}
+		delete(swap2fee, s.swapHash)
+
 		confDetail := &ConfDetail{
-			TxConfirmation: conf,
-			OnChainFeePortion: getFeePortionPaidBySweep(
-				spendTx, feePortionPaidPerSweep,
-				roundingDifference, &s,
-			),
+			TxConfirmation:    conf,
+			OnChainFeePortion: fee,
 		}
 
 		// Notify the caller in a goroutine to avoid possible dead-lock.

sweepbatcher/sweep_batcher.go

@@ -886,7 +886,7 @@ func (b *Batcher) handleSweeps(ctx context.Context, sweeps []*sweep,
 	// Instead we directly detect and return the spend here.
 	if completed && parentBatch.Confirmed {
 		return b.monitorSpendAndNotify(
-			ctx, sweep, parentBatch.ID, notifier,
+			ctx, sweeps, parentBatch.ID, notifier,
 		)
 	}
 
@@ -1186,7 +1186,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.
@@ -1204,6 +1204,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,
@@ -1224,6 +1235,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 :=
@@ -1232,17 +1244,23 @@ func (b *Batcher) monitorSpendAndNotify(ctx context.Context, sweep *sweep,
 					totalSwept,
 				)
 
-			onChainFeePortion := getFeePortionPaidBySweep(
-				spendTx, feePortionPerSweep,
-				roundingDifference, sweep,
-			)
+			// Sum onchain fee across all the sweeps of the swap.
+			var fee btcutil.Amount
+			for _, s := range sweeps {
+				isFirst := s.outpoint == primarySweepID
+
+				fee += getFeePortionPaidBySweep(
+					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 {
@@ -1250,7 +1268,7 @@ func (b *Batcher) monitorSpendAndNotify(ctx context.Context, sweep *sweep,
 			case notifier.SpendChan <- spendDetail:
 				err := b.monitorConfAndNotify(
 					ctx, sweep, notifier, spendTx,
-					onChainFeePortion,
+					fee,
 				)
 				if err != nil {
 					b.writeToErrChan(

sweepbatcher/sweep_batcher_presigned_test.go

@@ -1568,10 +1568,31 @@ func testPresigned_purging(t *testing.T, numSwaps, numConfirmedSwaps int,
 	}
 	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
@@ -1594,6 +1615,7 @@ func testPresigned_purging(t *testing.T, numSwaps, numConfirmedSwaps int,
 	for i := range numConfirmedSwaps {
 		conf := <-confChans[i]
 		require.Equal(t, txHash, conf.Tx.TxHash())
+		require.Equal(t, wantFee[i], conf.OnChainFeePortion)
 	}
 
 	if !online && numConfirmedSwaps != numSwaps {
@@ -1631,6 +1653,7 @@ func testPresigned_purging(t *testing.T, numSwaps, numConfirmedSwaps int,
 
 		spend := <-spendChan
 		require.Equal(t, txHash, spend.Tx.TxHash())
+		require.Equal(t, wantFee[i], spend.OnChainFeePortion)
 
 		<-lnd.RegisterConfChannel
 		lnd.ConfChannel <- &chainntnfs.TxConfirmation{
@@ -1639,6 +1662,7 @@ func testPresigned_purging(t *testing.T, numSwaps, numConfirmedSwaps int,
 
 		conf := <-confChan
 		require.Equal(t, tx.TxHash(), conf.Tx.TxHash())
+		require.Equal(t, wantFee[i], conf.OnChainFeePortion)
 	}
 
 	// If all the swaps were confirmed, stop.