xmr: back to flat workflow

src/apps/monero/protocol/signing/step_01_init_transaction.py

@@ -10,12 +10,11 @@
 from apps.monero.xmr import common, crypto, monero
 
 
-async def init_transaction(ctx, address_n, network_type, tsx_data):
+async def init_transaction(state: State, address_n, network_type, tsx_data):
     from apps.monero.xmr.sub.addr import classify_subaddresses
     from apps.monero.protocol import hmac_encryption_keys
 
-    state = State(ctx)
-    state.creds = await misc.monero_get_creds(ctx, address_n, network_type)
+    state.creds = await misc.monero_get_creds(state.ctx, address_n, network_type)
 
     state.tx_priv = crypto.random_scalar()
     state.tx_pub = crypto.scalarmult_base(state.tx_priv)
@@ -111,26 +110,13 @@ async def init_transaction(ctx, address_n, network_type, tsx_data):
 
     rsig_data = MoneroTransactionRsigData(offload_type=state.rsig_offload)
 
-    result_msg = MoneroTransactionInitAck(
+    return MoneroTransactionInitAck(
         in_memory=False,
         many_inputs=True,
         many_outputs=True,
         hmacs=hmacs,
         rsig_data=rsig_data,
     )
-    return await dispatch_and_forward(state, result_msg)
-
-
-async def dispatch_and_forward(state, result_msg):
-    from trezor.messages import MessageType
-    from apps.monero.protocol.signing import step_02_set_input
-
-    await state.ctx.write(result_msg)
-    del result_msg
-
-    received_msg = await state.ctx.read((MessageType.MoneroTransactionSetInputRequest,))
-
-    return await step_02_set_input.set_input(state, received_msg.src_entr)
 
 
 def get_primary_change_address(state: State):

src/apps/monero/protocol/signing/step_02_set_input.py

@@ -106,7 +106,11 @@ async def set_input(state: State, src_entr):
         crypto.encodeint(xi),
     )
 
-    result_msg = MoneroTransactionSetInputAck(
+    # All inputs done?
+    if state.inp_idx + 1 == state.num_inputs():
+        tsx_inputs_done(state)
+
+    return MoneroTransactionSetInputAck(
         vini=vini_bin,
         vini_hmac=hmac_vini,
         pseudo_out=pseudo_out,
@@ -115,44 +119,16 @@ async def set_input(state: State, src_entr):
         spend_enc=spend_enc,
     )
 
-    # All inputs done?
-    if state.inp_idx + 1 != state.num_inputs():
-        return await dispatch_and_next_input(
-            state, result_msg
-        )  # todo: we can merge those two functions
-
-    return await dispatch_and_forward(state, result_msg)
-
-
-async def dispatch_and_next_input(state, result_msg):
-    from trezor.messages import MessageType
-
-    await state.ctx.write(result_msg)
-    del result_msg
-
-    received_msg = await state.ctx.read((MessageType.MoneroTransactionSetInputRequest,))
-    return await set_input(state, received_msg.src_entr)
-
-
-async def dispatch_and_forward(state, result_msg):
-    from trezor.messages import MessageType
-    from apps.monero.protocol.signing import step_03_inputs_permutation
 
+def tsx_inputs_done(state: State):
     """
     All inputs set
     """
-    state.subaddresses = None  # todo why?
-
-    await state.ctx.write(result_msg)
-    del result_msg
+    # self.state.input_done()
+    state.subaddresses = None  # TODO why? remove this?
 
-    received_msg = await state.ctx.read(
-        (MessageType.MoneroTransactionInputsPermutationRequest,)
-    )
-
-    return await step_03_inputs_permutation.tsx_inputs_permutation(
-        state, received_msg.perm
-    )
+    if state.inp_idx + 1 != state.num_inputs():
+        raise ValueError("Input count mismatch")
 
 
 def _gen_commitment(state: State, in_amount):

src/apps/monero/protocol/signing/step_03_inputs_permutation.py

@@ -20,7 +20,7 @@ async def tsx_inputs_permutation(state: State, permutation):
 
     _tsx_inputs_permutation(state, permutation)
 
-    return await dispatch_and_forward(state, MoneroTransactionInputsPermutationAck())
+    return MoneroTransactionInputsPermutationAck()
 
 
 def _tsx_inputs_permutation(state: State, permutation):
@@ -30,24 +30,3 @@ def _tsx_inputs_permutation(state: State, permutation):
     state.source_permutation = permutation
     common.check_permutation(permutation)
     state.inp_idx = -1
-
-
-async def dispatch_and_forward(state, result_msg):
-    from trezor.messages import MessageType
-    from apps.monero.protocol.signing import step_04_input_vini
-
-    await state.ctx.write(result_msg)
-    del result_msg
-
-    received_msg = await state.ctx.read(
-        (MessageType.MoneroTransactionInputViniRequest,)
-    )
-
-    return await step_04_input_vini.input_vini(
-        state,
-        received_msg.src_entr,
-        received_msg.vini,
-        received_msg.vini_hmac,
-        received_msg.pseudo_out,
-        received_msg.pseudo_out_hmac,
-    )

src/apps/monero/protocol/signing/step_04_input_vini.py

@@ -12,12 +12,7 @@
 
 
 async def input_vini(
-    state: State,
-    src_entr,
-    vini_bin,
-    hmac,
-    pseudo_out,
-    pseudo_out_hmac,
+    state: State, src_entr, vini_bin, hmac, pseudo_out, pseudo_out_hmac
 ):
     from trezor.messages.MoneroTransactionInputViniAck import (
         MoneroTransactionInputViniAck
@@ -41,15 +36,12 @@ async def input_vini(
 
     hash_vini_pseudo_out(state, vini_bin, state.inp_idx, pseudo_out, pseudo_out_hmac)
 
-    return await dispatch_and_forward(state, MoneroTransactionInputViniAck())
+    # TODO check input count?
+    return MoneroTransactionInputViniAck()
 
 
 def hash_vini_pseudo_out(
-    state: State,
-    vini_bin,
-    inp_idx,
-    pseudo_out=None,
-    pseudo_out_hmac=None,
+    state: State, vini_bin, inp_idx, pseudo_out=None, pseudo_out_hmac=None
 ):
     """
     Incremental hasing of tx.vin[i] and pseudo output
@@ -68,31 +60,3 @@ def hash_vini_pseudo_out(
         raise ValueError("HMAC invalid for pseudo outs")
 
     state.full_message_hasher.set_pseudo_out(pseudo_out)
-
-
-async def dispatch_and_forward(state, result_msg):
-    from trezor.messages import MessageType
-    from apps.monero.protocol.signing import step_05_all_in_set
-
-    await state.ctx.write(result_msg)
-    del result_msg
-
-    received_msg = await state.ctx.read(
-        (
-            MessageType.MoneroTransactionInputViniRequest,
-            MessageType.MoneroTransactionAllInputsSetRequest,
-        )
-    )
-    # todo check input count
-
-    if received_msg.MESSAGE_WIRE_TYPE == MessageType.MoneroTransactionInputViniRequest:
-        return await input_vini(
-            state,
-            received_msg.src_entr,
-            received_msg.vini,
-            received_msg.vini_hmac,
-            received_msg.pseudo_out,
-            received_msg.pseudo_out_hmac,
-        )
-
-    return await step_05_all_in_set.all_in_set(state, received_msg.rsig_data)

src/apps/monero/protocol/signing/step_05_all_in_set.py

@@ -27,7 +27,7 @@ async def all_in_set(state: State, rsig_data):  # todo: rsig_data not used?
     resp = MoneroTransactionAllInputsSetAck(rsig_data=rsig_data)
 
     if not state.rsig_offload:
-        return await dispatch_and_forward(state, resp)
+        return resp
 
     # Simple offloading - generate random masks that sum to the input mask sum.
     tmp_buff = bytearray(32)
@@ -48,19 +48,5 @@ async def all_in_set(state: State, rsig_data):  # todo: rsig_data not used?
 
     state.assrt(crypto.sc_eq(state.sumout, state.sumpouts_alphas), "Invalid masks sum")
     state.sumout = crypto.sc_init(0)
-    return await dispatch_and_forward(state, resp)
 
-
-async def dispatch_and_forward(state, result_msg):
-    from trezor.messages import MessageType
-    from apps.monero.protocol.signing import step_06_set_out1
-
-    await state.ctx.write(result_msg)
-    del result_msg
-
-    received_msg = await state.ctx.read(
-        (MessageType.MoneroTransactionSetOutputRequest,)
-    )
-    return await step_06_set_out1.set_out1(
-        state, received_msg.dst_entr, received_msg.dst_entr_hmac, received_msg.rsig_data
-    )
+    return resp

src/apps/monero/protocol/signing/step_06_set_out1.py

@@ -14,9 +14,7 @@
 from apps.monero.xmr import common, crypto
 
 
-async def set_out1(
-    state: State, dst_entr, dst_entr_hmac, rsig_data=None
-):
+async def set_out1(state: State, dst_entr, dst_entr_hmac, rsig_data=None):
     state._mem_trace(0, True)
     mods = utils.unimport_begin()
 
@@ -103,14 +101,13 @@ async def set_out1(
     utils.memcpy(out_pk_bin, 0, out_pk.dest, 0, 32)
     utils.memcpy(out_pk_bin, 32, out_pk.mask, 0, 32)
 
-    result = MoneroTransactionSetOutputAck(
+    return MoneroTransactionSetOutputAck(
         tx_out=tx_out_bin,
         vouti_hmac=hmac_vouti,
         rsig_data=_return_rsig_data(state, rsig),
         out_pk=out_pk_bin,
         ecdh_info=ecdh_info_bin,
     )
-    return await dispatch_and_forward(state, result)
 
 
 async def _set_out1_tx_out(state: State, dst_entr, tx_out_key):
@@ -239,9 +236,7 @@ def _return_rsig_data(state, rsig):
         return MoneroTransactionRsigData(rsig=rsig)
 
 
-def _set_out1_ecdh(
-    state: State, dest_pub_key, amount, mask, amount_key
-):
+def _set_out1_ecdh(state: State, dest_pub_key, amount, mask, amount_key):
     from apps.monero.xmr import ring_ct
 
     # Mask sum
@@ -294,9 +289,7 @@ def _set_out1_additional_keys(state: State, dst_entr):
     return additional_txkey_priv
 
 
-def _set_out1_derivation(
-    state: State, dst_entr, additional_txkey_priv
-):
+def _set_out1_derivation(state: State, dst_entr, additional_txkey_priv):
     from apps.monero.xmr.sub.addr import addr_eq
 
     change_addr = state.change_address()
@@ -359,28 +352,3 @@ def _get_out_mask(state: State, idx):
             return crypto.sc_sub(state.sumpouts_alphas, state.sumout)
         else:
             return crypto.random_scalar()
-
-
-async def dispatch_and_forward(state, result_msg):
-    from trezor.messages import MessageType
-    from apps.monero.protocol.signing import step_07_all_out1_set
-
-    await state.ctx.write(result_msg)
-    del result_msg
-
-    received_msg = await state.ctx.read(
-        (
-            MessageType.MoneroTransactionSetOutputRequest,
-            MessageType.MoneroTransactionAllOutSetRequest,
-        )
-    )
-
-    if received_msg.MESSAGE_WIRE_TYPE == MessageType.MoneroTransactionSetOutputRequest:
-        return await set_out1(
-            state,
-            received_msg.dst_entr,
-            received_msg.dst_entr_hmac,
-            received_msg.rsig_data,
-        )
-
-    return await step_07_all_out1_set.all_out1_set(state)

src/apps/monero/protocol/signing/step_07_all_out1_set.py

@@ -85,10 +85,9 @@ async def all_out1_set(state: State):
         rv_type=state.get_rct_type(),
     )
 
-    result = MoneroTransactionAllOutSetAck(
+    return MoneroTransactionAllOutSetAck(
         extra=extra_b, tx_prefix_hash=state.tx_prefix_hash, rv=rv_pb
     )
-    return await dispatch_and_forward(state, result)
 
 
 def _set_tx_extra(state: State):
@@ -114,14 +113,3 @@ def _set_tx_prefix(state: State):
 
     # Hash message to the final_message
     state.full_message_hasher.set_message(state.tx_prefix_hash)
-
-
-async def dispatch_and_forward(state, result_msg):
-    from trezor.messages import MessageType
-    from apps.monero.protocol.signing import step_08_mlsag_done
-
-    await state.ctx.write(result_msg)
-    del result_msg
-
-    await state.ctx.read((MessageType.MoneroTransactionMlsagDoneRequest,))
-    return await step_08_mlsag_done.mlsag_done(state)

src/apps/monero/protocol/signing/step_08_mlsag_done.py

@@ -11,6 +11,7 @@ async def mlsag_done(state: State):
     from trezor.messages.MoneroTransactionMlsagDoneAck import (
         MoneroTransactionMlsagDoneAck
     )
+
     # state.state.set_final_message_done()  todo needed?
     await confirms.transaction_step(state.ctx, state.STEP_MLSAG)
 
@@ -23,8 +24,7 @@ async def mlsag_done(state: State):
     state.full_message = state.full_message_hasher.get_digest()
     state.full_message_hasher = None
 
-    result = MoneroTransactionMlsagDoneAck(full_message_hash=state.full_message)
-    return await dispatch_and_forward(state, result)
+    return MoneroTransactionMlsagDoneAck(full_message_hash=state.full_message)
 
 
 def _ecdh_info(state: State):  # todo why is it here? remove
@@ -43,25 +43,3 @@ def _out_pk(state: State):
 
     for out in state.output_pk:
         state.full_message_hasher.set_out_pk(out)
-
-
-async def dispatch_and_forward(state, result_msg):
-    from trezor.messages import MessageType
-    from apps.monero.protocol.signing import step_09_sign_input
-
-    await state.ctx.write(result_msg)
-    del result_msg
-
-    received_msg = await state.ctx.read(
-        (MessageType.MoneroTransactionSignInputRequest,)
-    )
-    return await step_09_sign_input.sign_input(
-        state,
-        received_msg.src_entr,
-        received_msg.vini,
-        received_msg.vini_hmac,
-        received_msg.pseudo_out,
-        received_msg.pseudo_out_hmac,
-        received_msg.alpha_enc,
-        received_msg.spend_enc,
-    )