From 16edbf30260bf62023f3bce44bb0cd03360feafe Mon Sep 17 00:00:00 2001
From: Olaoluwa Osuntokun <laolu32@gmail.com>
Date: Fri, 27 Sep 2024 17:07:53 +0900
Subject: [PATCH] sweep: ensure we factor in extra change addrs in
 MaxFeeRateAllowed

---
 sweep/fee_bumper.go        | 45 +++++++++++++++++++++++++++++++++-----
 sweep/fee_bumper_test.go   |  8 ++++---
 sweep/tx_input_set_test.go |  2 +-
 sweep/txgenerator.go       |  3 ++-
 4 files changed, 48 insertions(+), 10 deletions(-)

diff --git a/sweep/fee_bumper.go b/sweep/fee_bumper.go
index b1e34c0bf..fd3dfba9e 100644
--- a/sweep/fee_bumper.go
+++ b/sweep/fee_bumper.go
@@ -20,6 +20,7 @@ import (
 	"github.com/lightningnetwork/lnd/lnutils"
 	"github.com/lightningnetwork/lnd/lnwallet"
 	"github.com/lightningnetwork/lnd/lnwallet/chainfee"
+	"github.com/lightningnetwork/lnd/tlv"
 )
 
 var (
@@ -43,6 +44,19 @@ var (
 	ErrThirdPartySpent = errors.New("third party spent the output")
 )
 
+var (
+	// dummyChangePkScript is a dummy tapscript change script that's used
+	// when we don't need a real address, just something that can be used
+	// for fee estimation.
+	dummyChangePkScript = []byte{
+		0x51, 0x20,
+		0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
+		0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
+		0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
+		0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
+	}
+)
+
 // Bumper defines an interface that can be used by other subsystems for fee
 // bumping.
 type Bumper interface {
@@ -129,12 +143,33 @@ type BumpRequest struct {
 // compares it with the specified MaxFeeRate, and returns the smaller of the
 // two.
 func (r *BumpRequest) MaxFeeRateAllowed() (chainfee.SatPerKWeight, error) {
+	// We'll want to know if we have any blobs, as we need to factor this
+	// into the max fee rate for this bump request.
+	hasBlobs := fn.Any(func(i input.Input) bool {
+		return fn.MapOptionZ(
+			i.ResolutionBlob(), func(b tlv.Blob) bool {
+				return len(b) > 0
+			},
+		)
+	}, r.Inputs)
+
+	sweepAddrs := [][]byte{
+		r.DeliveryAddress.DeliveryAddress,
+	}
+
+	// If we have blobs, then we'll add an extra sweep addr for the size
+	// estimate below. We know that these blobs will also always be based on
+	// p2tr addrs.
+	if hasBlobs {
+		// We need to pass in a real address, so we'll use a dummy
+		// tapscript change script that's used elsewhere for tests.
+		sweepAddrs = append(sweepAddrs, dummyChangePkScript)
+	}
+
 	// Get the size of the sweep tx, which will be used to calculate the
 	// budget fee rate.
-	//
-	// TODO(roasbeef): also wants the extra change output?
 	size, err := calcSweepTxWeight(
-		r.Inputs, r.DeliveryAddress.DeliveryAddress,
+		r.Inputs, sweepAddrs,
 	)
 	if err != nil {
 		return 0, err
@@ -163,7 +198,7 @@ func (r *BumpRequest) MaxFeeRateAllowed() (chainfee.SatPerKWeight, error) {
 // calcSweepTxWeight calculates the weight of the sweep tx. It assumes a
 // sweeping tx always has a single output(change).
 func calcSweepTxWeight(inputs []input.Input,
-	outputPkScript []byte) (lntypes.WeightUnit, error) {
+	outputPkScript [][]byte) (lntypes.WeightUnit, error) {
 
 	// Use a const fee rate as we only use the weight estimator to
 	// calculate the size.
@@ -177,7 +212,7 @@ func calcSweepTxWeight(inputs []input.Input,
 	// TODO(yy): we should refactor the weight estimator to not require a
 	// fee rate and max fee rate and make it a pure tx weight calculator.
 	_, estimator, err := getWeightEstimate(
-		inputs, nil, feeRate, 0, [][]byte{outputPkScript},
+		inputs, nil, feeRate, 0, outputPkScript,
 	)
 	if err != nil {
 		return 0, err
diff --git a/sweep/fee_bumper_test.go b/sweep/fee_bumper_test.go
index db9dd4045..53b38607f 100644
--- a/sweep/fee_bumper_test.go
+++ b/sweep/fee_bumper_test.go
@@ -115,13 +115,15 @@ func TestCalcSweepTxWeight(t *testing.T) {
 	inp := createTestInput(100, input.WitnessKeyHash)
 
 	// Use a wrong change script to test the error case.
-	weight, err := calcSweepTxWeight([]input.Input{&inp}, []byte{0})
+	weight, err := calcSweepTxWeight(
+		[]input.Input{&inp}, [][]byte{{0x00}},
+	)
 	require.Error(t, err)
 	require.Zero(t, weight)
 
 	// Use a correct change script to test the success case.
 	weight, err = calcSweepTxWeight(
-		[]input.Input{&inp}, changePkScript.DeliveryAddress,
+		[]input.Input{&inp}, [][]byte{changePkScript.DeliveryAddress},
 	)
 	require.NoError(t, err)
 
@@ -143,7 +145,7 @@ func TestBumpRequestMaxFeeRateAllowed(t *testing.T) {
 
 	// The weight is 487.
 	weight, err := calcSweepTxWeight(
-		[]input.Input{&inp}, changePkScript.DeliveryAddress,
+		[]input.Input{&inp}, [][]byte{changePkScript.DeliveryAddress},
 	)
 	require.NoError(t, err)
 
diff --git a/sweep/tx_input_set_test.go b/sweep/tx_input_set_test.go
index b774eedff..8d0850b20 100644
--- a/sweep/tx_input_set_test.go
+++ b/sweep/tx_input_set_test.go
@@ -92,7 +92,7 @@ func TestNewBudgetInputSet(t *testing.T) {
 	rt.ErrorContains(err, "duplicate inputs")
 	rt.Nil(set)
 
-	// Pass a slice of inputs that only one input has the deadline height,
+	// Pass a slice of inputs that only one input has the deadline height.
 	set, err = NewBudgetInputSet(
 		[]SweeperInput{input0, input3}, testHeight,
 		fn.None[AuxSweeper](),
diff --git a/sweep/txgenerator.go b/sweep/txgenerator.go
index 43fc802ba..993ee9e59 100644
--- a/sweep/txgenerator.go
+++ b/sweep/txgenerator.go
@@ -262,7 +262,8 @@ func getWeightEstimate(inputs []input.Input, outputs []*wire.TxOut,
 
 		default:
 			// Unknown script type.
-			return nil, nil, errors.New("unknown script type")
+			return nil, nil, fmt.Errorf("unknown script "+
+				"type: %x", outputPkScript)
 		}
 	}