diff --git a/core/src/apps/wallet/sign_tx/bitcoin.py b/core/src/apps/wallet/sign_tx/bitcoin.py index 86aea6469..bed2639fd 100644 --- a/core/src/apps/wallet/sign_tx/bitcoin.py +++ b/core/src/apps/wallet/sign_tx/bitcoin.py @@ -128,14 +128,12 @@ class Bitcoin: await self.process_input(txi) async def step2_confirm_outputs(self) -> None: - txo_bin = TxOutputBinType() for i in range(self.tx.outputs_count): # STAGE_REQUEST_3_OUTPUT txo = await helpers.request_tx_output(self.tx_req, i, self.coin) - txo_bin.amount = txo.amount - txo_bin.script_pubkey = self.output_derive_script(txo) - self.weight.add_output(txo_bin.script_pubkey) - await self.confirm_output(txo, txo_bin) + script_pubkey = self.output_derive_script(txo) + self.weight.add_output(script_pubkey) + await self.confirm_output(txo, script_pubkey) async def step3_confirm_tx(self) -> None: fee = self.total_in - self.total_out @@ -218,15 +216,15 @@ class Bitcoin: self.bip143_in += txi.amount self.total_in += txi.amount - async def confirm_output(self, txo: TxOutputType, txo_bin: TxOutputBinType) -> None: + async def confirm_output(self, txo: TxOutputType, script_pubkey: bytes) -> None: if self.change_out == 0 and self.output_is_change(txo): # output is change and does not need confirmation self.change_out = txo.amount elif not await helpers.confirm_output(txo, self.coin): raise SigningError(FailureType.ActionCancelled, "Output cancelled") - self.write_tx_output(self.h_confirmed, txo_bin) - self.hash143_add_output(txo_bin) + self.write_tx_output(self.h_confirmed, txo, script_pubkey) + self.hash143_add_output(txo, script_pubkey) self.total_out += txo.amount def on_negative_fee(self) -> None: @@ -337,14 +335,12 @@ class Bitcoin: writers.write_varint(h_sign, self.tx.outputs_count) - txo_bin = TxOutputBinType() for i in range(self.tx.outputs_count): # STAGE_REQUEST_4_OUTPUT txo = await helpers.request_tx_output(self.tx_req, i, self.coin) - txo_bin.amount = txo.amount - txo_bin.script_pubkey = self.output_derive_script(txo) - self.write_tx_output(h_check, txo_bin) - self.write_tx_output(h_sign, txo_bin) + script_pubkey = self.output_derive_script(txo) + self.write_tx_output(h_check, txo, script_pubkey) + self.write_tx_output(h_sign, txo, script_pubkey) writers.write_uint32(h_sign, self.tx.lock_time) writers.write_uint32(h_sign, self.get_hash_type()) @@ -369,10 +365,8 @@ class Bitcoin: async def serialize_output(self, i: int) -> None: # STAGE_REQUEST_5_OUTPUT txo = await helpers.request_tx_output(self.tx_req, i, self.coin) - txo_bin = TxOutputBinType() - txo_bin.amount = txo.amount - txo_bin.script_pubkey = self.output_derive_script(txo) - self.write_tx_output(self.serialized_tx, txo_bin) + script_pubkey = self.output_derive_script(txo) + self.write_tx_output(self.serialized_tx, txo, script_pubkey) async def get_prevtx_output_value(self, prev_hash: bytes, prev_index: int) -> int: amount_out = 0 # output amount @@ -403,7 +397,7 @@ class Bitcoin: txo_bin = await helpers.request_tx_output( self.tx_req, i, self.coin, prev_hash ) - self.write_tx_output(txh, txo_bin) + self.write_tx_output(txh, txo_bin, txo_bin.script_pubkey) if i == prev_index: amount_out = txo_bin.amount self.check_prevtx_output(txo_bin) @@ -435,8 +429,13 @@ class Bitcoin: ) -> None: writers.write_tx_input(w, txi, script) - def write_tx_output(self, w: writers.Writer, txo_bin: TxOutputBinType) -> None: - writers.write_tx_output(w, txo_bin) + def write_tx_output( + self, + w: writers.Writer, + txo: Union[TxOutputType, TxOutputBinType], + script_pubkey: bytes, + ) -> None: + writers.write_tx_output(w, txo, script_pubkey) def write_tx_header( self, w: writers.Writer, tx: Union[SignTx, TransactionType], has_segwit: bool @@ -519,8 +518,8 @@ class Bitcoin: writers.write_uint32(self.h_prevouts, txi.prev_index) writers.write_uint32(self.h_sequence, txi.sequence) - def hash143_add_output(self, txo_bin: TxOutputBinType) -> None: - writers.write_tx_output(self.h_outputs, txo_bin) + def hash143_add_output(self, txo: TxOutputType, script_pubkey) -> None: + writers.write_tx_output(self.h_outputs, txo, script_pubkey) def hash143_preimage_hash(self, txi: TxInputType, pubkeyhash: bytes) -> bytes: h_preimage = HashWriter(sha256()) diff --git a/core/src/apps/wallet/sign_tx/decred.py b/core/src/apps/wallet/sign_tx/decred.py index dac7971ee..4aa9dfa35 100644 --- a/core/src/apps/wallet/sign_tx/decred.py +++ b/core/src/apps/wallet/sign_tx/decred.py @@ -58,15 +58,14 @@ class Decred(Bitcoin): # Decred serializes inputs early. self.write_tx_input(self.serialized_tx, txi, bytes()) - async def confirm_output(self, txo: TxOutputType, txo_bin: TxOutputBinType) -> None: + async def confirm_output(self, txo: TxOutputType, script_pubkey: bytes) -> None: if txo.decred_script_version != 0: raise SigningError( FailureType.ActionCancelled, "Cannot send to output with script version != 0", ) - txo_bin.decred_script_version = txo.decred_script_version - await super().confirm_output(txo, txo_bin) - self.write_tx_output(self.serialized_tx, txo_bin) + await super().confirm_output(txo, script_pubkey) + self.write_tx_output(self.serialized_tx, txo, script_pubkey) async def step4_serialize_inputs(self) -> None: writers.write_varint(self.serialized_tx, self.tx.inputs_count) @@ -147,16 +146,21 @@ class Decred(Bitcoin): def hash143_add_input(self, txi: TxInputType) -> None: writers.write_tx_input_decred(self.h_prefix, txi) - def hash143_add_output(self, txo_bin: TxOutputBinType) -> None: - writers.write_tx_output_decred(self.h_prefix, txo_bin) + def hash143_add_output(self, txo: TxOutputType, script_pubkey: bytes) -> None: + writers.write_tx_output_decred(self.h_prefix, txo, script_pubkey) def write_tx_input( self, w: writers.Writer, txi: TxInputType, script: bytes ) -> None: writers.write_tx_input_decred(w, txi) - def write_tx_output(self, w: writers.Writer, txo_bin: TxOutputBinType) -> None: - writers.write_tx_output_decred(w, txo_bin) + def write_tx_output( + self, + w: writers.Writer, + txo: Union[TxOutputType, TxOutputBinType], + script_pubkey: bytes, + ) -> None: + writers.write_tx_output_decred(w, txo, script_pubkey) def write_tx_header( self, w: writers.Writer, tx: Union[SignTx, TransactionType], has_segwit: bool diff --git a/core/src/apps/wallet/sign_tx/writers.py b/core/src/apps/wallet/sign_tx/writers.py index a280fc835..16041b068 100644 --- a/core/src/apps/wallet/sign_tx/writers.py +++ b/core/src/apps/wallet/sign_tx/writers.py @@ -3,6 +3,7 @@ from micropython import const from trezor.crypto.hashlib import sha256 from trezor.messages.TxInputType import TxInputType from trezor.messages.TxOutputBinType import TxOutputBinType +from trezor.messages.TxOutputType import TxOutputType from trezor.utils import ensure from apps.common.writers import ( # noqa: F401 @@ -17,6 +18,7 @@ from apps.common.writers import ( # noqa: F401 ) if False: + from typing import Union from apps.common.writers import Writer from trezor.utils import HashWriter @@ -64,15 +66,19 @@ def write_tx_input_decred_witness(w: Writer, i: TxInputType, script_sig: bytes) write_bytes_prefixed(w, script_sig) -def write_tx_output(w: Writer, o: TxOutputBinType) -> None: +def write_tx_output( + w: Writer, o: Union[TxOutputType, TxOutputBinType], script_pubkey: bytes +) -> None: write_uint64(w, o.amount) - write_bytes_prefixed(w, o.script_pubkey) + write_bytes_prefixed(w, script_pubkey) -def write_tx_output_decred(w: Writer, o: TxOutputBinType) -> None: +def write_tx_output_decred( + w: Writer, o: Union[TxOutputType, TxOutputBinType], script_pubkey: bytes +) -> None: write_uint64(w, o.amount) write_uint16(w, o.decred_script_version) - write_bytes_prefixed(w, o.script_pubkey) + write_bytes_prefixed(w, script_pubkey) def write_op_push(w: Writer, n: int) -> None: diff --git a/core/tests/test_apps.wallet.segwit.bip143.native_p2wpkh.py b/core/tests/test_apps.wallet.segwit.bip143.native_p2wpkh.py index 1fec84e32..1223cf7a3 100644 --- a/core/tests/test_apps.wallet.segwit.bip143.native_p2wpkh.py +++ b/core/tests/test_apps.wallet.segwit.bip143.native_p2wpkh.py @@ -69,8 +69,8 @@ class TestSegwitBip143NativeP2WPKH(unittest.TestCase): for txo in [self.out1, self.out2]: txo_bin = TxOutputBinType() txo_bin.amount = txo.amount - txo_bin.script_pubkey = output_derive_script(txo, coin) - bip143.hash143_add_output(txo_bin) + script_pubkey = output_derive_script(txo, coin) + bip143.hash143_add_output(txo_bin, script_pubkey) outputs_hash = get_tx_hash(bip143.h_outputs, double=coin.sign_hash_double) self.assertEqual(hexlify(outputs_hash), b'863ef3e1a92afbfdb97f31ad0fc7683ee943e9abcf2501590ff8f6551f47e5e5') @@ -86,8 +86,8 @@ class TestSegwitBip143NativeP2WPKH(unittest.TestCase): for txo in [self.out1, self.out2]: txo_bin = TxOutputBinType() txo_bin.amount = txo.amount - txo_bin.script_pubkey = output_derive_script(txo, coin) - bip143.hash143_add_output(txo_bin) + script_pubkey = output_derive_script(txo, coin) + bip143.hash143_add_output(txo_bin, script_pubkey) # test data public key hash # only for input 2 - input 1 is not segwit diff --git a/core/tests/test_apps.wallet.segwit.bip143.p2wpkh_in_p2sh.py b/core/tests/test_apps.wallet.segwit.bip143.p2wpkh_in_p2sh.py index cfc0f3b99..cbc8510fb 100644 --- a/core/tests/test_apps.wallet.segwit.bip143.p2wpkh_in_p2sh.py +++ b/core/tests/test_apps.wallet.segwit.bip143.p2wpkh_in_p2sh.py @@ -58,8 +58,8 @@ class TestSegwitBip143(unittest.TestCase): for txo in [self.out1, self.out2]: txo_bin = TxOutputBinType() txo_bin.amount = txo.amount - txo_bin.script_pubkey = output_derive_script(txo, coin) - bip143.hash143_add_output(txo_bin) + script_pubkey = output_derive_script(txo, coin) + bip143.hash143_add_output(txo_bin, script_pubkey) outputs_hash = get_tx_hash(bip143.h_outputs, double=coin.sign_hash_double) self.assertEqual(hexlify(outputs_hash), b'de984f44532e2173ca0d64314fcefe6d30da6f8cf27bafa706da61df8a226c83') @@ -72,8 +72,8 @@ class TestSegwitBip143(unittest.TestCase): for txo in [self.out1, self.out2]: txo_bin = TxOutputBinType() txo_bin.amount = txo.amount - txo_bin.script_pubkey = output_derive_script(txo, coin) - bip143.hash143_add_output(txo_bin) + script_pubkey = output_derive_script(txo, coin) + bip143.hash143_add_output(txo_bin, script_pubkey) # test data public key hash result = bip143.hash143_preimage_hash(self.inp1, unhexlify('79091972186c449eb1ded22b78e40d009bdf0089')) diff --git a/core/tests/test_apps.wallet.zcash.zip143.py b/core/tests/test_apps.wallet.zcash.zip143.py index c42b53708..da84be85e 100644 --- a/core/tests/test_apps.wallet.zcash.zip143.py +++ b/core/tests/test_apps.wallet.zcash.zip143.py @@ -169,7 +169,7 @@ class TestZcashZip143(unittest.TestCase): txo = TxOutputBinType() txo.amount = o["amount"] txo.script_pubkey = unhexlify(o["script_pubkey"]) - zip143.hash143_add_output(txo) + zip143.hash143_add_output(txo, txo.script_pubkey) self.assertEqual(hexlify(get_tx_hash(zip143.h_prevouts)), v["prevouts_hash"]) self.assertEqual(hexlify(get_tx_hash(zip143.h_sequence)), v["sequence_hash"]) diff --git a/core/tests/test_apps.wallet.zcash.zip243.py b/core/tests/test_apps.wallet.zcash.zip243.py index 77ddb5f0f..01aa0a2ef 100644 --- a/core/tests/test_apps.wallet.zcash.zip243.py +++ b/core/tests/test_apps.wallet.zcash.zip243.py @@ -203,7 +203,7 @@ class TestZcashZip243(unittest.TestCase): txo = TxOutputBinType() txo.amount = o["amount"] txo.script_pubkey = unhexlify(o["script_pubkey"]) - zip243.hash143_add_output(txo) + zip243.hash143_add_output(txo, txo.script_pubkey) self.assertEqual(hexlify(get_tx_hash(zip243.h_prevouts)), v["prevouts_hash"]) self.assertEqual(hexlify(get_tx_hash(zip243.h_sequence)), v["sequence_hash"])