diff --git a/core/src/apps/bitcoin/ownership.py b/core/src/apps/bitcoin/ownership.py index e056eb017..793d6023e 100644 --- a/core/src/apps/bitcoin/ownership.py +++ b/core/src/apps/bitcoin/ownership.py @@ -3,12 +3,7 @@ from trezor.crypto import bip32, hashlib, hmac from apps.common.keychain import Keychain from apps.common.readers import read_bitcoin_varint -from apps.common.writers import ( - empty_bytearray, - write_bitcoin_varint, - write_bytes_fixed, - write_uint8, -) +from apps.common.writers import write_bitcoin_varint, write_bytes_fixed, write_uint8 from . import common from .scripts import read_bip322_signature_proof, write_bip322_signature_proof @@ -43,7 +38,7 @@ def generate_proof( if user_confirmed: flags |= _FLAG_USER_CONFIRMED - proof = empty_bytearray(4 + 1 + 1 + len(ownership_ids) * _OWNERSHIP_ID_LEN) + proof = utils.empty_bytearray(4 + 1 + 1 + len(ownership_ids) * _OWNERSHIP_ID_LEN) write_bytes_fixed(proof, _VERSION_MAGIC, 4) write_uint8(proof, flags) diff --git a/core/src/apps/bitcoin/scripts.py b/core/src/apps/bitcoin/scripts.py index b9affa545..4fdeaf946 100644 --- a/core/src/apps/bitcoin/scripts.py +++ b/core/src/apps/bitcoin/scripts.py @@ -5,7 +5,7 @@ from trezor.messages import InputScriptType from apps.common import address_type from apps.common.readers import read_bitcoin_varint -from apps.common.writers import empty_bytearray, write_bitcoin_varint +from apps.common.writers import write_bitcoin_varint from . import common from .multisig import ( @@ -141,7 +141,7 @@ def bip143_derive_script_code( def input_script_p2pkh_or_p2sh( pubkey: bytes, signature: bytes, hash_type: int ) -> bytearray: - w = empty_bytearray(5 + len(signature) + 1 + 5 + len(pubkey)) + w = utils.empty_bytearray(5 + len(signature) + 1 + 5 + len(pubkey)) append_signature(w, signature, hash_type) append_pubkey(w, pubkey) return w @@ -211,7 +211,7 @@ def output_script_native_p2wpkh_or_p2wsh(witprog: bytes) -> bytearray: length = len(witprog) utils.ensure(length == 20 or length == 32) - w = empty_bytearray(3 + length) + w = utils.empty_bytearray(3 + length) w.append(0x00) # witness version byte w.append(length) # pub key hash length is 20 (P2WPKH) or 32 (P2WSH) bytes write_bytes_fixed(w, witprog, length) # pub key hash @@ -231,7 +231,7 @@ def input_script_p2wpkh_in_p2sh(pubkeyhash: bytes) -> bytearray: # Signature is moved to the witness. utils.ensure(len(pubkeyhash) == 20) - w = empty_bytearray(3 + len(pubkeyhash)) + w = utils.empty_bytearray(3 + len(pubkeyhash)) w.append(0x16) # length of the data w.append(0x00) # witness version byte w.append(0x14) # P2WPKH witness program (pub key hash length) @@ -254,7 +254,7 @@ def input_script_p2wsh_in_p2sh(script_hash: bytes) -> bytearray: if len(script_hash) != 32: raise wire.DataError("Redeem script hash should be 32 bytes long") - w = empty_bytearray(3 + len(script_hash)) + w = utils.empty_bytearray(3 + len(script_hash)) w.append(0x22) # length of the data w.append(0x00) # witness version byte w.append(0x20) # P2WSH witness program (redeem script hash length) @@ -267,7 +267,7 @@ def input_script_p2wsh_in_p2sh(script_hash: bytes) -> bytearray: def witness_p2wpkh(signature: bytes, pubkey: bytes, hash_type: int) -> bytearray: - w = empty_bytearray(1 + 5 + len(signature) + 1 + 5 + len(pubkey)) + w = utils.empty_bytearray(1 + 5 + len(signature) + 1 + 5 + len(pubkey)) write_bitcoin_varint(w, 0x02) # num of segwit items, in P2WPKH it's always 2 write_signature_prefixed(w, signature, hash_type) write_bytes_prefixed(w, pubkey) @@ -326,7 +326,7 @@ def witness_multisig( total_length += 1 + len(s) + 1 # length, signature, hash_type total_length += 1 + redeem_script_length # length, script - w = empty_bytearray(total_length) + w = utils.empty_bytearray(total_length) write_bitcoin_varint(w, num_of_witness_items) # Starts with OP_FALSE because of an old OP_CHECKMULTISIG bug, which @@ -399,7 +399,7 @@ def input_script_multisig( total_length += 1 + len(s) + 1 # length, signature, hash_type total_length += 1 + redeem_script_length # length, script - w = empty_bytearray(total_length) + w = utils.empty_bytearray(total_length) # Starts with OP_FALSE because of an old OP_CHECKMULTISIG bug, which # consumes one additional item on the stack: @@ -445,7 +445,7 @@ def parse_input_script_multisig( def output_script_multisig(pubkeys: list[bytes], m: int) -> bytearray: - w = empty_bytearray(output_script_multisig_length(pubkeys, m)) + w = utils.empty_bytearray(output_script_multisig_length(pubkeys, m)) write_output_script_multisig(w, pubkeys, m) return w @@ -507,7 +507,7 @@ def parse_output_script_multisig(script: bytes) -> tuple[list[bytes], int]: def output_script_paytoopreturn(data: bytes) -> bytearray: - w = empty_bytearray(1 + 5 + len(data)) + w = utils.empty_bytearray(1 + 5 + len(data)) w.append(0x6A) # OP_RETURN write_op_push(w, len(data)) w.extend(data) diff --git a/core/src/apps/bitcoin/scripts_decred.py b/core/src/apps/bitcoin/scripts_decred.py index 17cc79f42..2782ea679 100644 --- a/core/src/apps/bitcoin/scripts_decred.py +++ b/core/src/apps/bitcoin/scripts_decred.py @@ -3,7 +3,7 @@ from trezor.crypto import base58 from trezor.crypto.base58 import blake256d_32 from trezor.messages import InputScriptType -from apps.common.writers import empty_bytearray, write_bytes_fixed, write_uint64_le +from apps.common.writers import write_bytes_fixed, write_uint64_le from . import scripts from .multisig import multisig_get_pubkeys, multisig_pubkey_index @@ -65,7 +65,7 @@ def input_script_multisig( total_length += 1 + len(s) + 1 # length, signature, hash_type total_length += 1 + redeem_script_length # length, script - w = empty_bytearray(total_length) + w = utils.empty_bytearray(total_length) for s in signatures: if len(s): @@ -85,7 +85,7 @@ def output_script_sstxsubmissionpkh(addr: str) -> bytearray: except ValueError: raise wire.DataError("Invalid address") - w = empty_bytearray(26) + w = utils.empty_bytearray(26) w.append(0xBA) # OP_SSTX w.append(0x76) # OP_DUP w.append(0xA9) # OP_HASH160 @@ -103,7 +103,7 @@ def output_script_sstxchange(addr: str) -> bytearray: except ValueError: raise wire.DataError("Invalid address") - w = empty_bytearray(26) + w = utils.empty_bytearray(26) w.append(0xBD) # OP_SSTXCHANGE w.append(0x76) # OP_DUP w.append(0xA9) # OP_HASH160 @@ -144,7 +144,7 @@ def output_script_ssgen(pkh: bytes) -> bytearray: # Retrieve pkh bytes from a stake commitment OPRETURN. def sstxcommitment_pkh(pkh: bytes, amount: int) -> bytes: - w = empty_bytearray(30) + w = utils.empty_bytearray(30) write_bytes_fixed(w, pkh, 20) write_uint64_le(w, amount) write_bytes_fixed(w, b"\x00\x58", 2) # standard fee limits diff --git a/core/src/apps/bitcoin/sign_tx/bitcoin.py b/core/src/apps/bitcoin/sign_tx/bitcoin.py index 542ccdfa0..bcf4e4780 100644 --- a/core/src/apps/bitcoin/sign_tx/bitcoin.py +++ b/core/src/apps/bitcoin/sign_tx/bitcoin.py @@ -6,7 +6,7 @@ from trezor.messages import InputScriptType, OutputScriptType from trezor.messages.TxRequest import TxRequest from trezor.messages.TxRequestDetailsType import TxRequestDetailsType from trezor.messages.TxRequestSerializedType import TxRequestSerializedType -from trezor.utils import HashWriter, ensure +from trezor.utils import HashWriter, empty_bytearray, ensure from apps.common.writers import write_bitcoin_varint @@ -93,7 +93,7 @@ class Bitcoin: self.external: set[int] = set() # transaction and signature serialization - self.serialized_tx = writers.empty_bytearray(_MAX_SERIALIZED_CHUNK_SIZE) + self.serialized_tx = empty_bytearray(_MAX_SERIALIZED_CHUNK_SIZE) self.tx_req = TxRequest() self.tx_req.details = TxRequestDetailsType() self.tx_req.serialized = TxRequestSerializedType() diff --git a/core/src/apps/bitcoin/writers.py b/core/src/apps/bitcoin/writers.py index 8eadfceb1..3327b4839 100644 --- a/core/src/apps/bitcoin/writers.py +++ b/core/src/apps/bitcoin/writers.py @@ -4,7 +4,6 @@ from trezor.crypto.hashlib import sha256 from trezor.utils import ensure from apps.common.writers import ( # noqa: F401 - empty_bytearray, write_bitcoin_varint, write_bytes_fixed, write_bytes_reversed, diff --git a/core/src/apps/common/cbor.py b/core/src/apps/common/cbor.py index c04a169f1..3df990e0f 100644 --- a/core/src/apps/common/cbor.py +++ b/core/src/apps/common/cbor.py @@ -7,7 +7,7 @@ from micropython import const from trezor import log, utils -from . import readers, writers +from . import readers if False: from typing import Any, Union, Iterator, Tuple @@ -250,7 +250,7 @@ def encode_chunked(value: Value, max_chunk_size: int) -> Iterator[bytes]: chunks = encode_streamed(value) - chunk_buffer = writers.empty_bytearray(max_chunk_size) + chunk_buffer = utils.empty_bytearray(max_chunk_size) try: current_chunk_view = utils.BufferReader(next(chunks)) while True: diff --git a/core/src/apps/common/writers.py b/core/src/apps/common/writers.py index 5e20bfb78..8169fc746 100644 --- a/core/src/apps/common/writers.py +++ b/core/src/apps/common/writers.py @@ -4,16 +4,6 @@ if False: from trezor.utils import Writer -def empty_bytearray(preallocate: int) -> bytearray: - """ - Returns bytearray that won't allocate for at least `preallocate` bytes. - Useful in case you want to avoid allocating too often. - """ - b = bytearray(preallocate) - b[:] = bytes() - return b - - def write_uint8(w: Writer, n: int) -> int: ensure(0 <= n <= 0xFF) w.append(n) diff --git a/core/src/trezor/utils.py b/core/src/trezor/utils.py index 5009eb5c0..491300863 100644 --- a/core/src/trezor/utils.py +++ b/core/src/trezor/utils.py @@ -268,3 +268,13 @@ def is_empty_iterator(i: Iterator) -> bool: return True else: return False + + +def empty_bytearray(preallocate: int) -> bytearray: + """ + Returns bytearray that won't allocate for at least `preallocate` bytes. + Useful in case you want to avoid allocating too often. + """ + b = bytearray(preallocate) + b[:] = bytes() + return b