xmr: simplify key_image_sync workflow

src/apps/monero/key_image_sync.py

@@ -1,50 +1,108 @@
 import gc
 
-from trezor import log
+from apps.monero.controller import misc
+from apps.monero.layout import confirms
+from apps.monero.xmr import crypto, key_image, monero
+from apps.monero.xmr.enc import chacha_poly
+
+from trezor import log, wire
 from trezor.messages import MessageType
+from trezor.messages.MoneroExportedKeyImage import MoneroExportedKeyImage
+from trezor.messages.MoneroKeyImageExportInitAck import MoneroKeyImageExportInitAck
+from trezor.messages.MoneroKeyImageSyncFinalAck import MoneroKeyImageSyncFinalAck
+from trezor.messages.MoneroKeyImageSyncStepAck import MoneroKeyImageSyncStepAck
 
 
 async def key_image_sync(ctx, msg):
-    state = None
+    state = KeyImageSync()
 
+    res = await _init_step(state, ctx, msg)
     while True:
-        res, state, accept_msgs = await key_image_sync_step(ctx, msg, state)
-        if accept_msgs is None:
+        msg = await ctx.call(
+            res,
+            MessageType.MoneroKeyImageSyncStepRequest,
+            MessageType.MoneroKeyImageSyncFinalRequest,
+        )
+        del res
+        if msg.MESSAGE_WIRE_TYPE == MessageType.MoneroKeyImageSyncStepRequest:
+            res = await _sync_step(state, ctx, msg)
+        else:
+            res = await _final_step(state, ctx)
             break
-        msg = await ctx.call(res, *accept_msgs)
+        gc.collect()
 
     return res
 
 
-async def key_image_sync_step(ctx, msg, state):
-    if __debug__:
-        log.debug(__name__, "f: %s a: %s", gc.mem_free(), gc.mem_alloc())
-        log.debug(__name__, "s: %s", state)
+class KeyImageSync:
+    def __init__(self):
+        self.current_output = -1
+        self.num_outputs = 0
+        self.expected_hash = None
+        self.enc_key = None
+        self.creds = None
+        self.subaddresses = {}
+        self.hasher = crypto.get_keccak()
 
-    from apps.monero.protocol import key_image_sync
 
-    gc.collect()
+async def _init_step(s, ctx, msg):
+    s.creds = await misc.monero_get_creds(ctx, msg.address_n, msg.network_type)
 
-    if msg.MESSAGE_WIRE_TYPE == MessageType.MoneroKeyImageExportInitRequest:
-        state = key_image_sync.KeyImageSync(ctx=ctx)
-        return (
-            await state.init(ctx, msg),
-            state,
-            (MessageType.MoneroKeyImageSyncStepRequest,),
-        )
+    await confirms.require_confirm_keyimage_sync(ctx)
 
-    elif msg.MESSAGE_WIRE_TYPE == MessageType.MoneroKeyImageSyncStepRequest:
-        return (
-            await state.sync(ctx, msg),
-            state,
-            (
-                MessageType.MoneroKeyImageSyncStepRequest,
-                MessageType.MoneroKeyImageSyncFinalRequest,
-            ),
+    s.num_outputs = msg.num
+    s.expected_hash = msg.hash
+    s.enc_key = crypto.random_bytes(32)
+
+    for sub in msg.subs:
+        monero.compute_subaddresses(
+            s.creds, sub.account, sub.minor_indices, s.subaddresses
         )
 
-    elif msg.MESSAGE_WIRE_TYPE == MessageType.MoneroKeyImageSyncFinalRequest:
-        return await state.final(ctx, msg), None, None
+    return MoneroKeyImageExportInitAck()
+
+
+async def _sync_step(s, ctx, tds):
+    if not tds.tdis:
+        raise wire.DataError("Empty")
+
+    kis = []
+    buff = bytearray(32 * 3)
+    buff_mv = memoryview(buff)
+
+    for td in tds.tdis:
+        s.current_output += 1
+        if s.current_output >= s.num_outputs:
+            raise wire.DataError("Too many outputs")
+
+        if __debug__:
+            log.debug(__name__, "ki_sync, step i: %d", s.current_output)
+
+        # Update the control hash
+        s.hasher.update(key_image.compute_hash(td))
+
+        # Compute keyimage + signature
+        ki, sig = key_image.export_key_image(s.creds, s.subaddresses, td)
+
+        # Serialize into buff
+        crypto.encodepoint_into(buff_mv[0:32], ki)
+        crypto.encodeint_into(buff_mv[32:64], sig[0][0])
+        crypto.encodeint_into(buff_mv[64:], sig[0][1])
+
+        # Encrypt with enc_key
+        nonce, ciph, _ = chacha_poly.encrypt(s.enc_key, buff)
+
+        kis.append(MoneroExportedKeyImage(iv=nonce, blob=ciph, tag=b""))
+
+    return MoneroKeyImageSyncStepAck(kis=kis)
+
+
+async def _final_step(s, ctx):
+    if s.current_output + 1 != s.num_outputs:
+        raise wire.DataError("Invalid number of outputs")
+
+    final_hash = s.hasher.digest()
+    if final_hash != s.expected_hash:
+        raise wire.DataError("Invalid number of outputs")
 
-    else:
-        raise ValueError("Unknown error")
+    return MoneroKeyImageSyncFinalAck(enc_key=s.enc_key)

src/apps/monero/layout/confirms.py

@@ -1,12 +1,12 @@
+from apps.common.confirm import require_confirm, require_hold_to_confirm
+
 from trezor import ui
 from trezor.messages import ButtonRequestType
 from trezor.ui.text import Text
 from trezor.utils import chunks
 
 from . import common
 
-from apps.common.confirm import require_confirm, require_hold_to_confirm
-
 
 async def confirm_out(ctx, dst, is_change=False, creds=None, int_payment=None):
     """
@@ -180,23 +180,3 @@ async def transaction_step(ctx, step, sub_step=None, sub_step_total=None):
     text.normal(*info)
 
     await common.simple_text(text, tm=10 * 1000)
-
-
-async def confirm_ki_sync(ctx, init_msg):
-    await require_confirm_keyimage_sync(ctx)
-    return True
-
-
-async def ki_error(ctx, e):
-    # todo
-    pass
-
-
-async def ki_step(ctx, i):
-    # todo
-    pass
-
-
-async def ki_finished(ctx):
-    # todo
-    pass