From 3303780c70b1013b9adee865e01e0120c3bb486c Mon Sep 17 00:00:00 2001
From: Boris Nagaev <bnagaev@gmail.com>
Date: Thu, 11 Jul 2024 23:19:03 -0300
Subject: [PATCH] swap: fix leaf estimation in AddSuccessToEstimator

According to AddTapscriptInput's godoc, the first argument (leafWitnessSize)
must include not the whole size of witness, but only the size of data consumed
by the revealed script. Previously the value was too high, because it also
included the following extra elements: number_of_witness_elements (1 byte),
witness script and its size, control block and its size.

Tests for greedy_batch_selection were affected by this change.
---
 swap/htlc.go                                | 16 +++--
 sweepbatcher/greedy_batch_selection_test.go | 66 ++++++++++++++++++++-
 2 files changed, 76 insertions(+), 6 deletions(-)

diff --git a/swap/htlc.go b/swap/htlc.go
index dfba66e616..e3899141ee 100644
--- a/swap/htlc.go
+++ b/swap/htlc.go
@@ -266,8 +266,6 @@ func (h *Htlc) GenSuccessWitness(receiverSig []byte,
 
 // AddSuccessToEstimator adds a successful spend to a weight estimator.
 func (h *Htlc) AddSuccessToEstimator(estimator *input.TxWeightEstimator) error {
-	maxSuccessWitnessSize := h.MaxSuccessWitnessSize()
-
 	switch h.OutputType {
 	case HtlcP2TR:
 		// Generate tapscript.
@@ -283,10 +281,20 @@ func (h *Htlc) AddSuccessToEstimator(estimator *input.TxWeightEstimator) error {
 			trHtlc.InternalPubKey, successLeaf, timeoutLeafHash[:],
 		)
 
-		estimator.AddTapscriptInput(maxSuccessWitnessSize, tapscript)
+		// The leaf witness size must be calculated without the byte
+		// that accounts for the number of witness elements, only the
+		// total size of all elements on the stack that are consumed by
+		// the revealed script should be counted. Consumed elements:
+		// - sigLength: 1 byte
+		// - sig: 64 bytes
+		// - preimage_length: 1 byte
+		// - preimage: 32 bytes
+		const leafWitnessSize = 1 + 64 + 1 + 32
+
+		estimator.AddTapscriptInput(leafWitnessSize, tapscript)
 
 	case HtlcP2WSH:
-		estimator.AddWitnessInput(maxSuccessWitnessSize)
+		estimator.AddWitnessInput(h.MaxSuccessWitnessSize())
 	}
 
 	return nil
diff --git a/sweepbatcher/greedy_batch_selection_test.go b/sweepbatcher/greedy_batch_selection_test.go
index a12f0aeac1..a67e7db78f 100644
--- a/sweepbatcher/greedy_batch_selection_test.go
+++ b/sweepbatcher/greedy_batch_selection_test.go
@@ -18,7 +18,7 @@ const (
 	highFeeRate = chainfee.SatPerKWeight(30000)
 
 	coopInputWeight       = lntypes.WeightUnit(230)
-	nonCoopInputWeight    = lntypes.WeightUnit(521)
+	nonCoopInputWeight    = lntypes.WeightUnit(393)
 	nonCoopPenalty        = nonCoopInputWeight - coopInputWeight
 	coopNewBatchWeight    = lntypes.WeightUnit(444)
 	nonCoopNewBatchWeight = coopNewBatchWeight + nonCoopPenalty
@@ -31,7 +31,7 @@ const (
 	coopTwoSweepBatchWeight    = coopNewBatchWeight + coopInputWeight
 	nonCoopTwoSweepBatchWeight = coopTwoSweepBatchWeight +
 		2*nonCoopPenalty
-	v2v3BatchWeight = nonCoopTwoSweepBatchWeight - 153
+	v2v3BatchWeight = nonCoopTwoSweepBatchWeight - 25
 )
 
 // testHtlcV2SuccessEstimator adds weight of non-cooperative input to estimator
@@ -523,6 +523,37 @@ func TestSelectBatches(t *testing.T) {
 				NonCoopWeight: nonCoopNewBatchWeight,
 				NonCoopHint:   true,
 			},
+			wantBestBatchesIds: []int32{2, newBatchSignal, 1},
+		},
+
+		{
+			name: "high fee noncoop sweep, large batches",
+			batches: []feeDetails{
+				{
+					BatchId:       1,
+					FeeRate:       lowFeeRate,
+					CoopWeight:    10000,
+					NonCoopWeight: 15000,
+				},
+				{
+					BatchId:       2,
+					FeeRate:       highFeeRate,
+					CoopWeight:    10000,
+					NonCoopWeight: 15000,
+				},
+			},
+			sweep: feeDetails{
+				FeeRate:       highFeeRate,
+				CoopWeight:    coopInputWeight,
+				NonCoopWeight: nonCoopInputWeight,
+				NonCoopHint:   true,
+			},
+			oneSweepBatch: feeDetails{
+				FeeRate:       highFeeRate,
+				CoopWeight:    coopNewBatchWeight,
+				NonCoopWeight: nonCoopNewBatchWeight,
+				NonCoopHint:   true,
+			},
 			wantBestBatchesIds: []int32{newBatchSignal, 2, 1},
 		},
 
@@ -586,6 +617,37 @@ func TestSelectBatches(t *testing.T) {
 				NonCoopWeight: nonCoopNewBatchWeight,
 				NonCoopHint:   true,
 			},
+			wantBestBatchesIds: []int32{1, newBatchSignal, 2},
+		},
+
+		{
+			name: "low fee noncoop sweep, large batches",
+			batches: []feeDetails{
+				{
+					BatchId:       1,
+					FeeRate:       lowFeeRate,
+					CoopWeight:    10000,
+					NonCoopWeight: 15000,
+				},
+				{
+					BatchId:       2,
+					FeeRate:       highFeeRate,
+					CoopWeight:    10000,
+					NonCoopWeight: 15000,
+				},
+			},
+			sweep: feeDetails{
+				FeeRate:       lowFeeRate,
+				CoopWeight:    coopInputWeight,
+				NonCoopWeight: nonCoopInputWeight,
+				NonCoopHint:   true,
+			},
+			oneSweepBatch: feeDetails{
+				FeeRate:       lowFeeRate,
+				CoopWeight:    coopNewBatchWeight,
+				NonCoopWeight: nonCoopNewBatchWeight,
+				NonCoopHint:   true,
+			},
 			wantBestBatchesIds: []int32{newBatchSignal, 1, 2},
 		},