From 5e8452cc5d6395d2290860721b94e1ec0dbf69d8 Mon Sep 17 00:00:00 2001
From: yyforyongyu <yy2452@columbia.edu>
Date: Fri, 24 May 2024 23:47:04 +0800
Subject: [PATCH] sweep: make sure the full input is accounted

Fix the case where previously only the witness data is taken into
account when calculating the fees.
---
 sweep/aggregator.go      | 34 ++++++++++++++++++++++++++--------
 sweep/aggregator_test.go | 25 ++++++++++++++++---------
 2 files changed, 42 insertions(+), 17 deletions(-)

diff --git a/sweep/aggregator.go b/sweep/aggregator.go
index 801ac6b83..6028adca3 100644
--- a/sweep/aggregator.go
+++ b/sweep/aggregator.go
@@ -5,6 +5,8 @@ import (
 
 	"github.com/btcsuite/btcd/btcutil"
 	"github.com/btcsuite/btcd/wire"
+	"github.com/lightningnetwork/lnd/input"
+	"github.com/lightningnetwork/lnd/lntypes"
 	"github.com/lightningnetwork/lnd/lnwallet"
 	"github.com/lightningnetwork/lnd/lnwallet/chainfee"
 )
@@ -201,8 +203,8 @@ func (b *BudgetAggregator) filterInputs(inputs InputsMap) InputsMap {
 	for _, pi := range inputs {
 		op := pi.OutPoint()
 
-		// Get the size and skip if there's an error.
-		size, _, err := pi.WitnessType().SizeUpperBound()
+		// Get the size of the witness and skip if there's an error.
+		witnessSize, _, err := pi.WitnessType().SizeUpperBound()
 		if err != nil {
 			log.Warnf("Skipped input=%v: cannot get its size: %v",
 				op, err)
@@ -210,12 +212,27 @@ func (b *BudgetAggregator) filterInputs(inputs InputsMap) InputsMap {
 			continue
 		}
 
+		//nolint:lll
+		// Calculate the size if the input is included in the tx.
+		//
+		// NOTE: When including this input, we need to account the
+		// non-witness data which is expressed in vb.
+		//
+		// TODO(yy): This is not accurate for tapscript input. We need
+		// to unify calculations used in the `TxWeightEstimator` inside
+		// `input/size.go` and `weightEstimator` in
+		// `weight_estimator.go`. And calculate the expected weights
+		// similar to BOLT-3:
+		// https://github.com/lightning/bolts/blob/master/03-transactions.md#appendix-a-expected-weights
+		wu := lntypes.VByte(input.InputSize).ToWU() + witnessSize
+
 		// Skip inputs that has too little budget.
-		minFee := minFeeRate.FeeForWeight(size)
+		minFee := minFeeRate.FeeForWeight(wu)
 		if pi.params.Budget < minFee {
 			log.Warnf("Skipped input=%v: has budget=%v, but the "+
-				"min fee requires %v", op, pi.params.Budget,
-				minFee)
+				"min fee requires %v (feerate=%v), size=%v", op,
+				pi.params.Budget, minFee,
+				minFeeRate.FeePerVByte(), wu.ToVB())
 
 			continue
 		}
@@ -224,11 +241,12 @@ func (b *BudgetAggregator) filterInputs(inputs InputsMap) InputsMap {
 		startingFeeRate := pi.params.StartingFeeRate.UnwrapOr(
 			chainfee.SatPerKWeight(0),
 		)
-		startingFee := startingFeeRate.FeeForWeight(size)
+		startingFee := startingFeeRate.FeeForWeight(wu)
 		if pi.params.Budget < startingFee {
 			log.Errorf("Skipped input=%v: has budget=%v, but the "+
-				"starting fee requires %v", op,
-				pi.params.Budget, minFee)
+				"starting fee requires %v (feerate=%v), "+
+				"size=%v", op, pi.params.Budget, startingFee,
+				startingFeeRate.FeePerVByte(), wu.ToVB())
 
 			continue
 		}
diff --git a/sweep/aggregator_test.go b/sweep/aggregator_test.go
index 0b60bc3a7..a32c849a0 100644
--- a/sweep/aggregator_test.go
+++ b/sweep/aggregator_test.go
@@ -43,8 +43,11 @@ func TestBudgetAggregatorFilterInputs(t *testing.T) {
 	defer wt.AssertExpectations(t)
 
 	// Mock the `SizeUpperBound` method to return the size four times.
-	const wtSize lntypes.WeightUnit = 100
-	wt.On("SizeUpperBound").Return(wtSize, true, nil).Times(4)
+	const wu lntypes.WeightUnit = 100
+	wt.On("SizeUpperBound").Return(wu, true, nil).Times(4)
+
+	// Calculate the input size.
+	inpSize := lntypes.VByte(input.InputSize).ToWU() + wu
 
 	// Create a mock input that will be filtered out due to error.
 	inpErr := &input.MockInput{}
@@ -64,9 +67,9 @@ func TestBudgetAggregatorFilterInputs(t *testing.T) {
 	var (
 		// Define three budget values, one below the min fee rate, one
 		// above and one equal to it.
-		budgetLow   = minFeeRate.FeeForWeight(wtSize) - 1
-		budgetEqual = minFeeRate.FeeForWeight(wtSize)
-		budgetHigh  = minFeeRate.FeeForWeight(wtSize) + 1
+		budgetLow   = minFeeRate.FeeForWeight(inpSize) - 1
+		budgetEqual = minFeeRate.FeeForWeight(inpSize)
+		budgetHigh  = minFeeRate.FeeForWeight(inpSize) + 1
 
 		// Define three outpoints with different budget values.
 		opLow   = wire.OutPoint{Hash: chainhash.Hash{2}}
@@ -398,8 +401,12 @@ func TestBudgetInputSetClusterInputs(t *testing.T) {
 
 	// Mock the `SizeUpperBound` method to return the size 10 times since
 	// we are using ten inputs.
-	const wtSize lntypes.WeightUnit = 100
-	wt.On("SizeUpperBound").Return(wtSize, true, nil).Times(10)
+	const wu lntypes.WeightUnit = 100
+	wt.On("SizeUpperBound").Return(wu, true, nil).Times(10)
+
+	// Calculate the input size.
+	inpSize := lntypes.VByte(input.InputSize).ToWU() + wu
+
 	wt.On("String").Return("mock witness type")
 
 	// Mock the estimator to return a constant fee rate.
@@ -409,8 +416,8 @@ func TestBudgetInputSetClusterInputs(t *testing.T) {
 	var (
 		// Define two budget values, one below the min fee rate and one
 		// above it.
-		budgetLow  = minFeeRate.FeeForWeight(wtSize) - 1
-		budgetHigh = minFeeRate.FeeForWeight(wtSize) + 1
+		budgetLow  = minFeeRate.FeeForWeight(inpSize) - 1
+		budgetHigh = minFeeRate.FeeForWeight(inpSize) + 1
 
 		// Create three deadline heights, which means there are three
 		// groups of inputs to be expected.