btcwallet: add EstimateInputWeight helper function

This is a helper function that we will need to accurately determine the
weight of inputs specified in a PSBT.
Due to the nature of P2WSH and script-spend P2TR inputs, we can only
accurately estimate their weights if the full witness is already known.
So this helper function rejects inputs that use a script spend path but
don't fully specify the complete witness stack.

lnwallet/btcwallet/psbt.go

@@ -3,6 +3,7 @@ package btcwallet
 import (
 	"bytes"
 	"crypto/sha256"
+	"errors"
 	"fmt"
 
 	"github.com/btcsuite/btcd/btcec/v2"
@@ -30,6 +31,22 @@ var (
 	// the key before signing the input. The value d0 is leet speak for
 	// "do", short for "double".
 	PsbtKeyTypeInputSignatureTweakDouble = []byte{0xd0}
+
+	// ErrInputMissingUTXOInfo is returned if a PSBT input is supplied that
+	// does not specify the witness UTXO info.
+	ErrInputMissingUTXOInfo = errors.New(
+		"input doesn't specify any UTXO info",
+	)
+
+	// ErrScriptSpendFeeEstimationUnsupported is returned if a PSBT input is
+	// of a script spend type.
+	ErrScriptSpendFeeEstimationUnsupported = errors.New(
+		"cannot estimate fee for script spend inputs",
+	)
+
+	// ErrUnsupportedScript is returned if a supplied pk script is not
+	// known or supported.
+	ErrUnsupportedScript = errors.New("unsupported or unknown pk script")
 )
 
 // FundPsbt creates a fully populated PSBT packet that contains enough inputs to
@@ -352,6 +369,62 @@ func validateSigningMethod(in *psbt.PInput) (input.SignMethod, error) {
 	}
 }
 
+// EstimateInputWeight estimates the weight of a PSBT input and adds it to the
+// passed in TxWeightEstimator. It returns an error if the input type is
+// unknown or unsupported. Only inputs that have a known witness size are
+// supported, which is P2WKH, NP2WKH and P2TR (key spend path).
+func EstimateInputWeight(in *psbt.PInput, w *input.TxWeightEstimator) error {
+	if in.WitnessUtxo == nil {
+		return ErrInputMissingUTXOInfo
+	}
+
+	pkScript := in.WitnessUtxo.PkScript
+	switch {
+	case txscript.IsPayToScriptHash(pkScript):
+		w.AddNestedP2WKHInput()
+
+	case txscript.IsPayToWitnessPubKeyHash(pkScript):
+		w.AddP2WKHInput()
+
+	case txscript.IsPayToWitnessScriptHash(pkScript):
+		return fmt.Errorf("P2WSH inputs are not supported, cannot "+
+			"estimate witness size for script spend: %w",
+			ErrScriptSpendFeeEstimationUnsupported)
+
+	case txscript.IsPayToTaproot(pkScript):
+		signMethod, err := validateSigningMethod(in)
+		if err != nil {
+			return fmt.Errorf("error determining p2tr signing "+
+				"method: %w", err)
+		}
+
+		switch signMethod {
+		// For p2tr key spend paths.
+		case input.TaprootKeySpendBIP0086SignMethod,
+			input.TaprootKeySpendSignMethod:
+
+			w.AddTaprootKeySpendInput(in.SighashType)
+
+		// For p2tr script spend path.
+		case input.TaprootScriptSpendSignMethod:
+			return fmt.Errorf("P2TR inputs are not supported, "+
+				"cannot estimate witness size for script "+
+				"spend: %w",
+				ErrScriptSpendFeeEstimationUnsupported)
+
+		default:
+			return fmt.Errorf("unsupported signing method for "+
+				"PSBT signing: %v", signMethod)
+		}
+
+	default:
+		return fmt.Errorf("unknown input type for script %x: %w",
+			pkScript, ErrUnsupportedScript)
+	}
+
+	return nil
+}
+
 // SignSegWitV0 attempts to generate a signature for a SegWit version 0 input
 // and stores it in the PartialSigs (and FinalScriptSig for np2wkh addresses)
 // field.

lnwallet/btcwallet/psbt_test.go

@@ -7,6 +7,7 @@ import (
 	"fmt"
 	"testing"
 
+	"github.com/btcsuite/btcd/blockchain"
 	"github.com/btcsuite/btcd/btcec/v2"
 	"github.com/btcsuite/btcd/btcutil"
 	"github.com/btcsuite/btcd/btcutil/psbt"
@@ -343,3 +344,154 @@ func TestSignPsbt(t *testing.T) {
 		require.NoError(t, vm.Execute())
 	}
 }
+
+// TestEstimateInputWeight tests that we correctly estimate the weight of a
+// PSBT input if it supplies all required information.
+func TestEstimateInputWeight(t *testing.T) {
+	genScript := func(f func([]byte) ([]byte, error)) []byte {
+		pkScript, _ := f([]byte{})
+		return pkScript
+	}
+
+	var (
+		witnessScaleFactor = blockchain.WitnessScaleFactor
+		p2trScript, _      = txscript.PayToTaprootScript(
+			&input.TaprootNUMSKey,
+		)
+		dummyLeaf = txscript.TapLeaf{
+			LeafVersion: txscript.BaseLeafVersion,
+			Script:      []byte("some bitcoin script"),
+		}
+		dummyLeafHash = dummyLeaf.TapHash()
+	)
+
+	testCases := []struct {
+		name string
+		in   *psbt.PInput
+
+		expectedErr       error
+		expectedErrString string
+
+		// expectedWitnessWeight is the expected weight of the content
+		// of the witness of the input (without the base input size that
+		// is constant for all types of inputs).
+		expectedWitnessWeight int
+	}{{
+		name:        "empty input",
+		in:          &psbt.PInput{},
+		expectedErr: ErrInputMissingUTXOInfo,
+	}, {
+		name: "empty pkScript",
+		in: &psbt.PInput{
+			WitnessUtxo: &wire.TxOut{},
+		},
+		expectedErr: ErrUnsupportedScript,
+	}, {
+		name: "nested p2wpkh input",
+		in: &psbt.PInput{
+			WitnessUtxo: &wire.TxOut{
+				PkScript: genScript(input.GenerateP2SH),
+			},
+		},
+		expectedWitnessWeight: input.P2WKHWitnessSize +
+			input.NestedP2WPKHSize*witnessScaleFactor,
+	}, {
+		name: "p2wpkh input",
+		in: &psbt.PInput{
+			WitnessUtxo: &wire.TxOut{
+				PkScript: genScript(input.WitnessPubKeyHash),
+			},
+		},
+		expectedWitnessWeight: input.P2WKHWitnessSize,
+	}, {
+		name: "p2wsh input",
+		in: &psbt.PInput{
+			WitnessUtxo: &wire.TxOut{
+				PkScript: genScript(input.WitnessScriptHash),
+			},
+		},
+		expectedErr: ErrScriptSpendFeeEstimationUnsupported,
+	}, {
+		name: "p2tr with no derivation info",
+		in: &psbt.PInput{
+			WitnessUtxo: &wire.TxOut{
+				PkScript: p2trScript,
+			},
+		},
+		expectedErrString: "cannot sign for taproot input " +
+			"without taproot BIP0032 derivation info",
+	}, {
+		name: "p2tr key spend",
+		in: &psbt.PInput{
+			WitnessUtxo: &wire.TxOut{
+				PkScript: p2trScript,
+			},
+			SighashType: txscript.SigHashSingle,
+			TaprootBip32Derivation: []*psbt.TaprootBip32Derivation{
+				{},
+			},
+		},
+		//nolint:lll
+		expectedWitnessWeight: input.TaprootKeyPathCustomSighashWitnessSize,
+	}, {
+		name: "p2tr script spend",
+		in: &psbt.PInput{
+			WitnessUtxo: &wire.TxOut{
+				PkScript: p2trScript,
+			},
+			TaprootBip32Derivation: []*psbt.TaprootBip32Derivation{
+				{
+					LeafHashes: [][]byte{
+						dummyLeafHash[:],
+					},
+				},
+			},
+			TaprootLeafScript: []*psbt.TaprootTapLeafScript{
+				{
+					LeafVersion: dummyLeaf.LeafVersion,
+					Script:      dummyLeaf.Script,
+				},
+			},
+		},
+		expectedErr: ErrScriptSpendFeeEstimationUnsupported,
+	}}
+
+	// The non-witness weight for a TX with a single input.
+	nonWitnessWeight := input.BaseTxSize + 1 + 1 + input.InputSize
+
+	// The base weight of a witness TX.
+	baseWeight := (nonWitnessWeight * witnessScaleFactor) +
+		input.WitnessHeaderSize
+
+	for _, tc := range testCases {
+		tc := tc
+
+		t.Run(tc.name, func(tt *testing.T) {
+			estimator := input.TxWeightEstimator{}
+			err := EstimateInputWeight(tc.in, &estimator)
+
+			if tc.expectedErr != nil {
+				require.Error(tt, err)
+				require.ErrorIs(tt, err, tc.expectedErr)
+
+				return
+			}
+
+			if tc.expectedErrString != "" {
+				require.Error(tt, err)
+				require.Contains(
+					tt, err.Error(), tc.expectedErrString,
+				)
+
+				return
+			}
+
+			require.NoError(tt, err)
+
+			require.EqualValues(
+				tt, baseWeight+tc.expectedWitnessWeight,
+				estimator.Weight(),
+			)
+		})
+	}
+}