core/sign_tx: Refactor BIP-143 signing.

core/src/apps/wallet/sign_tx/bitcoin.py

@@ -28,7 +28,7 @@ from apps.wallet.sign_tx.common import SigningError, ecdsa_sign
 from apps.wallet.sign_tx.matchcheck import MultisigFingerprintChecker, WalletPathChecker
 
 if False:
-    from typing import Set, Union
+    from typing import Set, Tuple, Union
 
 # Default signature hash type in Bitcoin, which signs the entire transaction except scripts.
 _SIGHASH_ALL = const(0x01)
@@ -208,16 +208,19 @@ class Bitcoin:
             raise SigningError(FailureType.DataError, "Wrong input script type")
 
     async def process_segwit_input(self, i: int, txi: TxInputType) -> None:
-        if not txi.amount:
-            raise SigningError(FailureType.DataError, "Segwit input without amount")
-        self.bip143_in += txi.amount
-        self.total_in += txi.amount
+        await self.process_bip143_input(i, txi)
 
     async def process_nonsegwit_input(self, i: int, txi: TxInputType) -> None:
         self.total_in += await self.get_prevtx_output_value(
             txi.prev_hash, txi.prev_index
         )
 
+    async def process_bip143_input(self, i: int, txi: TxInputType) -> None:
+        if not txi.amount:
+            raise SigningError(FailureType.DataError, "Expected input with amount")
+        self.bip143_in += txi.amount
+        self.total_in += txi.amount
+
     async def confirm_output(
         self, i: int, txo: TxOutputType, txo_bin: TxOutputBinType
     ) -> None:
@@ -252,34 +255,45 @@ class Bitcoin:
 
         self.write_tx_input(self.serialized_tx, txi)
 
-    async def sign_segwit_input(self, i: int) -> None:
-        # STAGE_REQUEST_SEGWIT_WITNESS
-        txi = await helpers.request_tx_input(self.tx_req, i, self.coin)
-
+    def sign_bip143_input(self, txi: TxInputType) -> Tuple[bytes, bytes]:
         self.wallet_path.check_input(txi)
         self.multisig_fingerprint.check_input(txi)
 
-        if not input_is_segwit(txi) or txi.amount > self.bip143_in:
+        if txi.amount > self.bip143_in:
             raise SigningError(
                 FailureType.ProcessError, "Transaction has changed during signing"
             )
         self.bip143_in -= txi.amount
 
         node = self.keychain.derive(txi.address_n, self.coin.curve_name)
-        key_sign_pub = node.public_key()
+        public_key = node.public_key()
         hash143_hash = self.hash143.preimage_hash(
             self.coin,
             self.tx,
             txi,
-            addresses.ecdsa_hash_pubkey(key_sign_pub, self.coin),
+            addresses.ecdsa_hash_pubkey(public_key, self.coin),
             self.get_hash_type(),
         )
 
         signature = ecdsa_sign(node, hash143_hash)
+
+        return public_key, signature
+
+    async def sign_segwit_input(self, i: int) -> None:
+        # STAGE_REQUEST_SEGWIT_WITNESS
+        txi = await helpers.request_tx_input(self.tx_req, i, self.coin)
+
+        if not input_is_segwit(txi):
+            raise SigningError(
+                FailureType.ProcessError, "Transaction has changed during signing"
+            )
+
+        public_key, signature = self.sign_bip143_input(txi)
+
         self.set_serialized_signature(i, signature)
         if txi.multisig:
             # find out place of our signature based on the pubkey
-            signature_index = multisig.multisig_pubkey_index(txi.multisig, key_sign_pub)
+            signature_index = multisig.multisig_pubkey_index(txi.multisig, public_key)
             self.serialized_tx.extend(
                 scripts.witness_p2wsh(
                     txi.multisig, signature, signature_index, self.get_hash_type()
@@ -287,7 +301,7 @@ class Bitcoin:
             )
         else:
             self.serialized_tx.extend(
-                scripts.witness_p2wpkh(signature, key_sign_pub, self.get_hash_type())
+                scripts.witness_p2wpkh(signature, public_key, self.get_hash_type())
             )
 
     async def sign_nonsegwit_input(self, i_sign: int) -> None:

core/src/apps/wallet/sign_tx/bitcoinlike.py

@@ -1,14 +1,14 @@
 import gc
 from micropython import const
 
-from trezor.messages import FailureType, InputScriptType
+from trezor.messages import FailureType
 from trezor.messages.SignTx import SignTx
 from trezor.messages.TransactionType import TransactionType
 from trezor.messages.TxInputType import TxInputType
 
-from apps.wallet.sign_tx import addresses, helpers, multisig, writers
-from apps.wallet.sign_tx.bitcoin import Bitcoin
-from apps.wallet.sign_tx.common import SigningError, ecdsa_sign
+from apps.wallet.sign_tx import helpers, multisig, writers
+from apps.wallet.sign_tx.bitcoin import Bitcoin, input_is_nonsegwit
+from apps.wallet.sign_tx.common import SigningError
 
 if False:
     from typing import Union
@@ -28,58 +28,31 @@ class Bitcoinlike(Bitcoin):
         else:
             await super().process_nonsegwit_input(i, txi)
 
-    async def process_bip143_input(self, i: int, txi: TxInputType) -> None:
-        if not txi.amount:
-            raise SigningError(FailureType.DataError, "Expected input with amount")
-        self.bip143_in += txi.amount
-        self.total_in += txi.amount
+    async def sign_nonsegwit_bip143_input(self, i_sign: int) -> None:
+        txi = await helpers.request_tx_input(self.tx_req, i_sign, self.coin)
 
-    async def sign_nonsegwit_input(self, i_sign: int) -> None:
-        if self.coin.force_bip143:
-            await self.sign_bip143_input(i_sign)
-        else:
-            await super().sign_nonsegwit_input(i_sign)
-
-    async def sign_bip143_input(self, i_sign: int) -> None:
-        # STAGE_REQUEST_SEGWIT_INPUT
-        txi_sign = await helpers.request_tx_input(self.tx_req, i_sign, self.coin)
-        self.wallet_path.check_input(txi_sign)
-        self.multisig_fingerprint.check_input(txi_sign)
-
-        is_bip143 = (
-            txi_sign.script_type == InputScriptType.SPENDADDRESS
-            or txi_sign.script_type == InputScriptType.SPENDMULTISIG
-        )
-        if not is_bip143 or txi_sign.amount > self.bip143_in:
+        if not input_is_nonsegwit(txi):
             raise SigningError(
                 FailureType.ProcessError, "Transaction has changed during signing"
             )
-        self.bip143_in -= txi_sign.amount
-
-        key_sign = self.keychain.derive(txi_sign.address_n, self.coin.curve_name)
-        key_sign_pub = key_sign.public_key()
-        hash143_hash = self.hash143.preimage_hash(
-            self.coin,
-            self.tx,
-            txi_sign,
-            addresses.ecdsa_hash_pubkey(key_sign_pub, self.coin),
-            self.get_hash_type(),
-        )
+        public_key, signature = self.sign_bip143_input(txi)
 
         # if multisig, do a sanity check to ensure we are signing with a key that is included in the multisig
-        if txi_sign.multisig:
-            multisig.multisig_pubkey_index(txi_sign.multisig, key_sign_pub)
-
-        signature = ecdsa_sign(key_sign, hash143_hash)
+        if txi.multisig:
+            multisig.multisig_pubkey_index(txi.multisig, public_key)
 
         # serialize input with correct signature
         gc.collect()
-        txi_sign.script_sig = self.input_derive_script(
-            txi_sign, key_sign_pub, signature
-        )
-        writers.write_tx_input(self.serialized_tx, txi_sign)
+        txi.script_sig = self.input_derive_script(txi, public_key, signature)
+        writers.write_tx_input(self.serialized_tx, txi)
         self.set_serialized_signature(i_sign, signature)
 
+    async def sign_nonsegwit_input(self, i_sign: int) -> None:
+        if self.coin.force_bip143:
+            await self.sign_nonsegwit_bip143_input(i_sign)
+        else:
+            await super().sign_nonsegwit_input(i_sign)
+
     def on_negative_fee(self) -> None:
         # some coins require negative fees for reward TX
         if not self.coin.negative_fee:

core/src/apps/wallet/sign_tx/zcash.py

@@ -219,7 +219,7 @@ class Overwintered(Bitcoinlike):
         await self.process_bip143_input(i, txi)
 
     async def sign_nonsegwit_input(self, i_sign: int) -> None:
-        await self.sign_bip143_input(i_sign)
+        await self.sign_nonsegwit_bip143_input(i_sign)
 
     def write_tx_header(
         self, w: Writer, tx: Union[SignTx, TransactionType], has_segwit: bool