diff options
Diffstat (limited to 'test/functional/test_framework/wallet.py')
-rw-r--r-- | test/functional/test_framework/wallet.py | 177 |
1 files changed, 113 insertions, 64 deletions
diff --git a/test/functional/test_framework/wallet.py b/test/functional/test_framework/wallet.py index e86f365f11..2164627781 100644 --- a/test/functional/test_framework/wallet.py +++ b/test/functional/test_framework/wallet.py @@ -10,6 +10,7 @@ from enum import Enum from random import choice from typing import ( Any, + List, Optional, ) from test_framework.address import ( @@ -18,9 +19,13 @@ from test_framework.address import ( key_to_p2pkh, key_to_p2sh_p2wpkh, key_to_p2wpkh, + output_key_to_p2tr, ) from test_framework.descriptors import descsum_create -from test_framework.key import ECKey +from test_framework.key import ( + ECKey, + compute_xonly_pubkey, +) from test_framework.messages import ( COIN, COutPoint, @@ -37,6 +42,7 @@ from test_framework.script import ( OP_NOP, OP_TRUE, SIGHASH_ALL, + taproot_construct, ) from test_framework.script_util import ( key_to_p2pk_script, @@ -80,8 +86,7 @@ class MiniWallet: def __init__(self, test_node, *, mode=MiniWalletMode.ADDRESS_OP_TRUE): self._test_node = test_node self._utxos = [] - self._priv_key = None - self._address = None + self._mode = mode assert isinstance(mode, MiniWalletMode) if mode == MiniWalletMode.RAW_OP_TRUE: @@ -96,6 +101,9 @@ class MiniWallet: self._address, self._internal_key = create_deterministic_address_bcrt1_p2tr_op_true() self._scriptPubKey = bytes.fromhex(self._test_node.validateaddress(self._address)['scriptPubKey']) + def _create_utxo(self, *, txid, vout, value, height): + return {"txid": txid, "vout": vout, "value": value, "height": height} + def get_balance(self): return sum(u['value'] for u in self._utxos) @@ -105,17 +113,26 @@ class MiniWallet: res = self._test_node.scantxoutset(action="start", scanobjects=[self.get_descriptor()]) assert_equal(True, res['success']) for utxo in res['unspents']: - self._utxos.append({'txid': utxo['txid'], 'vout': utxo['vout'], 'value': utxo['amount'], 'height': utxo['height']}) + self._utxos.append(self._create_utxo(txid=utxo["txid"], vout=utxo["vout"], value=utxo["amount"], height=utxo["height"])) def scan_tx(self, tx): - """Scan the tx for self._scriptPubKey outputs and add them to self._utxos""" + """Scan the tx and adjust the internal list of owned utxos""" + for spent in tx["vin"]: + # Mark spent. This may happen when the caller has ownership of a + # utxo that remained in this wallet. For example, by passing + # mark_as_spent=False to get_utxo or by using an utxo returned by a + # create_self_transfer* call. + try: + self.get_utxo(txid=spent["txid"], vout=spent["vout"]) + except StopIteration: + pass for out in tx['vout']: if out['scriptPubKey']['hex'] == self._scriptPubKey.hex(): - self._utxos.append({'txid': tx['txid'], 'vout': out['n'], 'value': out['value'], 'height': 0}) + self._utxos.append(self._create_utxo(txid=tx["txid"], vout=out["n"], value=out["value"], height=0)) def sign_tx(self, tx, fixed_length=True): """Sign tx that has been created by MiniWallet in P2PK mode""" - assert self._priv_key is not None + assert_equal(self._mode, MiniWalletMode.RAW_P2PK) (sighash, err) = LegacySignatureHash(CScript(self._scriptPubKey), tx, 0, SIGHASH_ALL) assert err is None # for exact fee calculation, create only signatures with fixed size by default (>49.89% probability): @@ -127,14 +144,19 @@ class MiniWallet: if not fixed_length: break tx.vin[0].scriptSig = CScript([der_sig + bytes(bytearray([SIGHASH_ALL]))]) + tx.rehash() def generate(self, num_blocks, **kwargs): - """Generate blocks with coinbase outputs to the internal address, and append the outputs to the internal list""" + """Generate blocks with coinbase outputs to the internal address, and call rescan_utxos""" blocks = self._test_node.generatetodescriptor(num_blocks, self.get_descriptor(), **kwargs) - for b in blocks: - block_info = self._test_node.getblock(blockhash=b, verbosity=2) - cb_tx = block_info['tx'][0] - self._utxos.append({'txid': cb_tx['txid'], 'vout': 0, 'value': cb_tx['vout'][0]['value'], 'height': block_info['height']}) + # Calling rescan_utxos here makes sure that after a generate the utxo + # set is in a clean state. For example, the wallet will update + # - if the caller consumed utxos, but never used them + # - if the caller sent a transaction that is not mined or got rbf'd + # - after block re-orgs + # - the utxo height for mined mempool txs + # - However, the wallet will not consider remaining mempool txs + self.rescan_utxos() return blocks def get_scriptPubKey(self): @@ -144,9 +166,10 @@ class MiniWallet: return descsum_create(f'raw({self._scriptPubKey.hex()})') def get_address(self): + assert_equal(self._mode, MiniWalletMode.ADDRESS_OP_TRUE) return self._address - def get_utxo(self, *, txid: str = '', vout: Optional[int] = None, mark_as_spent=True): + def get_utxo(self, *, txid: str = '', vout: Optional[int] = None, mark_as_spent=True) -> dict: """ Returns a utxo and marks it as spent (pops it from the internal list) @@ -166,10 +189,17 @@ class MiniWallet: else: return self._utxos[index] - def send_self_transfer(self, **kwargs): - """Create and send a tx with the specified fee_rate. Fee may be exact or at most one satoshi higher than needed.""" + def get_utxos(self, *, mark_as_spent=True): + """Returns the list of all utxos and optionally mark them as spent""" + utxos = deepcopy(self._utxos) + if mark_as_spent: + self._utxos = [] + return utxos + + def send_self_transfer(self, *, from_node, **kwargs): + """Call create_self_transfer and send the transaction.""" tx = self.create_self_transfer(**kwargs) - self.sendrawtransaction(from_node=kwargs['from_node'], tx_hex=tx['hex']) + self.sendrawtransaction(from_node=from_node, tx_hex=tx['hex']) return tx def send_to(self, *, from_node, scriptPubKey, amount, fee=1000): @@ -184,40 +214,45 @@ class MiniWallet: Returns a tuple (txid, n) referring to the created external utxo outpoint. """ - tx = self.create_self_transfer(from_node=from_node, fee_rate=0, mempool_valid=False)['tx'] + tx = self.create_self_transfer(fee_rate=0)["tx"] assert_greater_than_or_equal(tx.vout[0].nValue, amount + fee) tx.vout[0].nValue -= (amount + fee) # change output -> MiniWallet tx.vout.append(CTxOut(amount, scriptPubKey)) # arbitrary output -> to be returned txid = self.sendrawtransaction(from_node=from_node, tx_hex=tx.serialize().hex()) return txid, 1 - def send_self_transfer_multi(self, **kwargs): - """ - Create and send a transaction that spends the given UTXOs and creates a - certain number of outputs with equal amounts. - - Returns a dictionary with - - txid - - serialized transaction in hex format - - transaction as CTransaction instance - - list of newly created UTXOs, ordered by vout index - """ + def send_self_transfer_multi(self, *, from_node, **kwargs): + """Call create_self_transfer_multi and send the transaction.""" tx = self.create_self_transfer_multi(**kwargs) - txid = self.sendrawtransaction(from_node=kwargs['from_node'], tx_hex=tx.serialize().hex()) - return {'new_utxos': [self.get_utxo(txid=txid, vout=vout) for vout in range(len(tx.vout))], - 'txid': txid, 'hex': tx.serialize().hex(), 'tx': tx} + self.sendrawtransaction(from_node=from_node, tx_hex=tx["hex"]) + return tx - def create_self_transfer_multi(self, *, from_node, utxos_to_spend=None, num_outputs=1, fee_per_output=1000): + def create_self_transfer_multi( + self, + *, + utxos_to_spend: Optional[List[dict]] = None, + num_outputs=1, + amount_per_output=0, + sequence=0, + fee_per_output=1000, + ): """ Create and return a transaction that spends the given UTXOs and creates a - certain number of outputs with equal amounts. + certain number of outputs with equal amounts. The output amounts can be + set by amount_per_output or automatically calculated with a fee_per_output. """ utxos_to_spend = utxos_to_spend or [self.get_utxo()] + sequence = [sequence] * len(utxos_to_spend) if type(sequence) is int else sequence + assert_equal(len(utxos_to_spend), len(sequence)) # create simple tx template (1 input, 1 output) - tx = self.create_self_transfer(fee_rate=0, from_node=from_node, utxo_to_spend=utxos_to_spend[0], mempool_valid=False)['tx'] + tx = self.create_self_transfer( + fee_rate=0, + utxo_to_spend=utxos_to_spend[0])["tx"] # duplicate inputs, witnesses and outputs tx.vin = [deepcopy(tx.vin[0]) for _ in range(len(utxos_to_spend))] + for txin, seq in zip(tx.vin, sequence): + txin.nSequence = seq tx.wit.vtxinwit = [deepcopy(tx.wit.vtxinwit[0]) for _ in range(len(utxos_to_spend))] tx.vout = [deepcopy(tx.vout[0]) for _ in range(num_outputs)] @@ -229,54 +264,64 @@ class MiniWallet: inputs_value_total = sum([int(COIN * utxo['value']) for utxo in utxos_to_spend]) outputs_value_total = inputs_value_total - fee_per_output * num_outputs for o in tx.vout: - o.nValue = outputs_value_total // num_outputs - return tx - - def create_self_transfer(self, *, fee_rate=Decimal("0.003"), from_node=None, utxo_to_spend=None, mempool_valid=True, locktime=0, sequence=0): - """Create and return a tx with the specified fee_rate. Fee may be exact or at most one satoshi higher than needed.""" - from_node = from_node or self._test_node + o.nValue = amount_per_output or (outputs_value_total // num_outputs) + txid = tx.rehash() + return { + "new_utxos": [self._create_utxo( + txid=txid, + vout=i, + value=Decimal(tx.vout[i].nValue) / COIN, + height=0, + ) for i in range(len(tx.vout))], + "txid": txid, + "hex": tx.serialize().hex(), + "tx": tx, + } + + def create_self_transfer(self, *, fee_rate=Decimal("0.003"), fee=Decimal("0"), utxo_to_spend=None, locktime=0, sequence=0): + """Create and return a tx with the specified fee. If fee is 0, use fee_rate, where the resulting fee may be exact or at most one satoshi higher than needed.""" utxo_to_spend = utxo_to_spend or self.get_utxo() - if self._priv_key is None: + assert fee_rate >= 0 + assert fee >= 0 + if self._mode in (MiniWalletMode.RAW_OP_TRUE, MiniWalletMode.ADDRESS_OP_TRUE): vsize = Decimal(104) # anyone-can-spend - else: + elif self._mode == MiniWalletMode.RAW_P2PK: vsize = Decimal(168) # P2PK (73 bytes scriptSig + 35 bytes scriptPubKey + 60 bytes other) - send_value = int(COIN * (utxo_to_spend['value'] - fee_rate * (vsize / 1000))) + else: + assert False + send_value = utxo_to_spend["value"] - (fee or (fee_rate * vsize / 1000)) assert send_value > 0 tx = CTransaction() tx.vin = [CTxIn(COutPoint(int(utxo_to_spend['txid'], 16), utxo_to_spend['vout']), nSequence=sequence)] - tx.vout = [CTxOut(send_value, self._scriptPubKey)] + tx.vout = [CTxOut(int(COIN * send_value), bytearray(self._scriptPubKey))] tx.nLockTime = locktime - if not self._address: - # raw script - if self._priv_key is not None: - # P2PK, need to sign - self.sign_tx(tx) - else: - # anyone-can-spend - tx.vin[0].scriptSig = CScript([OP_NOP] * 43) # pad to identical size - else: + if self._mode == MiniWalletMode.RAW_P2PK: + self.sign_tx(tx) + elif self._mode == MiniWalletMode.RAW_OP_TRUE: + tx.vin[0].scriptSig = CScript([OP_NOP] * 43) # pad to identical size + elif self._mode == MiniWalletMode.ADDRESS_OP_TRUE: tx.wit.vtxinwit = [CTxInWitness()] tx.wit.vtxinwit[0].scriptWitness.stack = [CScript([OP_TRUE]), bytes([LEAF_VERSION_TAPSCRIPT]) + self._internal_key] + else: + assert False tx_hex = tx.serialize().hex() - tx_info = from_node.testmempoolaccept([tx_hex])[0] - assert_equal(mempool_valid, tx_info['allowed']) - if mempool_valid: - assert_equal(tx_info['vsize'], vsize) - assert_equal(tx_info['fees']['base'], utxo_to_spend['value'] - Decimal(send_value) / COIN) - return {'txid': tx_info['txid'], 'wtxid': tx_info['wtxid'], 'hex': tx_hex, 'tx': tx} + assert_equal(tx.get_vsize(), vsize) + new_utxo = self._create_utxo(txid=tx.rehash(), vout=0, value=send_value, height=0) + + return {"txid": new_utxo["txid"], "wtxid": tx.getwtxid(), "hex": tx_hex, "tx": tx, "new_utxo": new_utxo} - def sendrawtransaction(self, *, from_node, tx_hex, **kwargs): - txid = from_node.sendrawtransaction(hexstring=tx_hex, **kwargs) + def sendrawtransaction(self, *, from_node, tx_hex, maxfeerate=0, **kwargs): + txid = from_node.sendrawtransaction(hexstring=tx_hex, maxfeerate=maxfeerate, **kwargs) self.scan_tx(from_node.decoderawtransaction(tx_hex)) return txid -def getnewdestination(address_type='bech32'): +def getnewdestination(address_type='bech32m'): """Generate a random destination of the specified type and return the corresponding public key, scriptPubKey and address. Supported types are - 'legacy', 'p2sh-segwit' and 'bech32'. Can be used when a random + 'legacy', 'p2sh-segwit', 'bech32' and 'bech32m'. Can be used when a random destination is needed, but no compiled wallet is available (e.g. as replacement to the getnewaddress/getaddressinfo RPCs).""" key = ECKey() @@ -291,7 +336,11 @@ def getnewdestination(address_type='bech32'): elif address_type == 'bech32': scriptpubkey = key_to_p2wpkh_script(pubkey) address = key_to_p2wpkh(pubkey) - # TODO: also support bech32m (need to generate x-only-pubkey) + elif address_type == 'bech32m': + tap = taproot_construct(compute_xonly_pubkey(key.get_bytes())[0]) + pubkey = tap.output_pubkey + scriptpubkey = tap.scriptPubKey + address = output_key_to_p2tr(pubkey) else: assert False return pubkey, scriptpubkey, address |