diff --git a/input/size.go b/input/size.go index 6cebc282..c1964716 100644 --- a/input/size.go +++ b/input/size.go @@ -533,6 +533,14 @@ func (twe *TxWeightEstimator) AddNestedP2WSHInput(witnessSize int) *TxWeightEsti return twe } +// AddTxOutput adds a known TxOut to the weight estimator. +func (twe *TxWeightEstimator) AddTxOutput(txOut *wire.TxOut) *TxWeightEstimator { + twe.outputSize += txOut.SerializeSize() + twe.outputCount++ + + return twe +} + // AddP2PKHOutput updates the weight estimate to account for an additional P2PKH // output. func (twe *TxWeightEstimator) AddP2PKHOutput() *TxWeightEstimator { diff --git a/sweep/weight_estimator.go b/sweep/weight_estimator.go index 5a406815..693c3f25 100644 --- a/sweep/weight_estimator.go +++ b/sweep/weight_estimator.go @@ -2,6 +2,7 @@ package sweep import ( "github.com/btcsuite/btcd/chaincfg/chainhash" + "github.com/btcsuite/btcd/wire" "github.com/btcsuite/btcutil" "github.com/lightningnetwork/lnd/input" "github.com/lightningnetwork/lnd/lnwallet/chainfee" @@ -92,6 +93,12 @@ func (w *weightEstimator) addP2WKHOutput() { w.estimator.AddP2WKHOutput() } +// addOutput updates the weight estimate to account for the known +// output given. +func (w *weightEstimator) addOutput(txOut *wire.TxOut) { + w.estimator.AddTxOutput(txOut) +} + // weight gets the estimated weight of the transaction. func (w *weightEstimator) weight() int { return w.estimator.Weight() diff --git a/sweep/weight_estimator_test.go b/sweep/weight_estimator_test.go index e4f1d489..f64b8b89 100644 --- a/sweep/weight_estimator_test.go +++ b/sweep/weight_estimator_test.go @@ -3,7 +3,10 @@ package sweep import ( "testing" + "github.com/btcsuite/btcd/chaincfg" + "github.com/btcsuite/btcd/txscript" "github.com/btcsuite/btcd/wire" + "github.com/btcsuite/btcutil" "github.com/lightningnetwork/lnd/input" "github.com/lightningnetwork/lnd/lnwallet/chainfee" "github.com/stretchr/testify/require" @@ -77,3 +80,32 @@ func TestWeightEstimator(t *testing.T) { require.Equal(t, expectedFee, w.fee()) } + +// TestWeightEstimatorAddOutput tests that adding the raw P2WKH output to the +// estimator yield the same result as an estimated add. +func TestWeightEstimatorAddOutput(t *testing.T) { + testFeeRate := chainfee.SatPerKWeight(20000) + + p2wkhAddr, err := btcutil.NewAddressWitnessPubKeyHash( + make([]byte, 20), &chaincfg.MainNetParams, + ) + require.NoError(t, err) + + p2wkhScript, err := txscript.PayToAddrScript(p2wkhAddr) + require.NoError(t, err) + + // Create two estimators, add the raw P2WKH out to one. + txOut := &wire.TxOut{ + PkScript: p2wkhScript, + Value: 10000, + } + + w1 := newWeightEstimator(testFeeRate) + w1.addOutput(txOut) + + w2 := newWeightEstimator(testFeeRate) + w2.addP2WKHOutput() + + // Estimate hhould be the same. + require.Equal(t, w1.weight(), w2.weight()) +}