test: introduce and use calculate_input_weight helper

Rather than manually estimating an input's weight by adding up all the
involved components (fixed-size skeleton, compact-serialized lengths,
and the actual scriptSig / witness stack items) we can simply take use
of the serialization classes `CTxIn` / `CTxInWitness` instead, to
achieve the same with significantly less code.

The new helper is used in the functional tests rpc_psbt.py and
wallet_send.py, where the previous manual estimation code was
duplicated.

Author: theStack

test/functional/rpc_psbt.py

@@ -16,8 +16,6 @@
     CTxIn,
     CTxOut,
     MAX_BIP125_RBF_SEQUENCE,
-    WITNESS_SCALE_FACTOR,
-    ser_compact_size,
 )
 from test_framework.psbt import (
     PSBT,
@@ -42,6 +40,7 @@
     find_vout_for_address,
 )
 from test_framework.wallet_util import (
+    calculate_input_weight,
     generate_keypair,
     get_generate_key,
 )
@@ -752,17 +751,9 @@ def test_psbt_input_keys(psbt_input, keys):
                 input_idx = i
                 break
         psbt_in = dec["inputs"][input_idx]
-        # Calculate the input weight
-        # (prevout + sequence + length of scriptSig + scriptsig) * WITNESS_SCALE_FACTOR + len of num scriptWitness stack items + (length of stack item + stack item) * N stack items
-        # Note that occasionally this weight estimate may be slightly larger or smaller than the real weight
-        # as sometimes ECDSA signatures are one byte shorter than expected with a probability of 1/128
-        len_scriptsig = len(psbt_in["final_scriptSig"]["hex"]) // 2 if "final_scriptSig" in psbt_in else 0
-        len_scriptsig += len(ser_compact_size(len_scriptsig))
-        len_scriptwitness = (sum([(len(x) // 2) + len(ser_compact_size(len(x) // 2)) for x in psbt_in["final_scriptwitness"]]) + len(ser_compact_size(len(psbt_in["final_scriptwitness"])))) if "final_scriptwitness" in psbt_in else 0
-        len_prevout_txid = 32
-        len_prevout_index = 4
-        len_sequence = 4
-        input_weight = ((len_prevout_txid + len_prevout_index + len_sequence + len_scriptsig) * WITNESS_SCALE_FACTOR) + len_scriptwitness
+        scriptsig_hex = psbt_in["final_scriptSig"]["hex"] if "final_scriptSig" in psbt_in else ""
+        witness_stack_hex = psbt_in["final_scriptwitness"] if "final_scriptwitness" in psbt_in else None
+        input_weight = calculate_input_weight(scriptsig_hex, witness_stack_hex)
         low_input_weight = input_weight // 2
         high_input_weight = input_weight * 2
 

test/functional/test_framework/wallet_util.py

@@ -15,6 +15,11 @@
     script_to_p2wsh,
 )
 from test_framework.key import ECKey
+from test_framework.messages import (
+    CTxIn,
+    CTxInWitness,
+    WITNESS_SCALE_FACTOR,
+)
 from test_framework.script_util import (
     key_to_p2pkh_script,
     key_to_p2wpkh_script,
@@ -123,6 +128,19 @@ def generate_keypair(compressed=True, wif=False):
         privkey = bytes_to_wif(privkey.get_bytes(), compressed)
     return privkey, pubkey
 
+def calculate_input_weight(scriptsig_hex, witness_stack_hex=None):
+    """Given a scriptSig and a list of witness stack items for an input in hex format,
+       calculate the total input weight. If the input has no witness data,
+       `witness_stack_hex` can be set to None."""
+    tx_in = CTxIn(scriptSig=bytes.fromhex(scriptsig_hex))
+    witness_size = 0
+    if witness_stack_hex is not None:
+        tx_inwit = CTxInWitness()
+        for witness_item_hex in witness_stack_hex:
+            tx_inwit.scriptWitness.stack.append(bytes.fromhex(witness_item_hex))
+        witness_size = len(tx_inwit.serialize())
+    return len(tx_in.serialize()) * WITNESS_SCALE_FACTOR + witness_size
+
 class WalletUnlock():
     """
     A context manager for unlocking a wallet with a passphrase and automatically locking it afterward.

test/functional/wallet_send.py

@@ -9,10 +9,6 @@
 
 from test_framework.authproxy import JSONRPCException
 from test_framework.descriptors import descsum_create
-from test_framework.messages import (
-    ser_compact_size,
-    WITNESS_SCALE_FACTOR,
-)
 from test_framework.test_framework import BitcoinTestFramework
 from test_framework.util import (
     assert_equal,
@@ -21,7 +17,10 @@
     assert_raises_rpc_error,
     count_bytes,
 )
-from test_framework.wallet_util import generate_keypair
+from test_framework.wallet_util import (
+    calculate_input_weight,
+    generate_keypair,
+)
 
 
 class WalletSendTest(BitcoinTestFramework):
@@ -543,17 +542,9 @@ def run_test(self):
                 input_idx = i
                 break
         psbt_in = dec["inputs"][input_idx]
-        # Calculate the input weight
-        # (prevout + sequence + length of scriptSig + scriptsig) * WITNESS_SCALE_FACTOR + len of num scriptWitness stack items + (length of stack item + stack item) * N stack items
-        # Note that occasionally this weight estimate may be slightly larger or smaller than the real weight
-        # as sometimes ECDSA signatures are one byte shorter than expected with a probability of 1/128
-        len_scriptsig = len(psbt_in["final_scriptSig"]["hex"]) // 2 if "final_scriptSig" in psbt_in else 0
-        len_scriptsig += len(ser_compact_size(len_scriptsig))
-        len_scriptwitness = (sum([(len(x) // 2) + len(ser_compact_size(len(x) // 2)) for x in psbt_in["final_scriptwitness"]]) + len(ser_compact_size(len(psbt_in["final_scriptwitness"])))) if "final_scriptwitness" in psbt_in else 0
-        len_prevout_txid = 32
-        len_prevout_index = 4
-        len_sequence = 4
-        input_weight = ((len_prevout_txid + len_prevout_index + len_sequence + len_scriptsig) * WITNESS_SCALE_FACTOR) + len_scriptwitness
+        scriptsig_hex = psbt_in["final_scriptSig"]["hex"] if "final_scriptSig" in psbt_in else ""
+        witness_stack_hex = psbt_in["final_scriptwitness"] if "final_scriptwitness" in psbt_in else None
+        input_weight = calculate_input_weight(scriptsig_hex, witness_stack_hex)
 
         # Input weight error conditions
         assert_raises_rpc_error(