From 8303b42fb7bab913d687d9187780de7496d8114c Mon Sep 17 00:00:00 2001
From: Tomas Susanka <tsusanka@gmail.com>
Date: Thu, 11 Oct 2018 15:35:47 +0200
Subject: [PATCH] xmr: state's use_simple_rct and use_bulletproof modified to
 enums

---
 src/apps/monero/controller/misc.py            | 21 ++++++++++
 src/apps/monero/protocol/signing/rct_type.py  | 13 +++++++
 src/apps/monero/protocol/signing/rsig_type.py | 17 +++++++++
 src/apps/monero/protocol/signing/state.py     | 23 ++++-------
 .../signing/step_01_init_transaction.py       | 38 ++++++++++++++-----
 .../protocol/signing/step_02_set_input.py     |  3 +-
 .../protocol/signing/step_04_input_vini.py    |  6 ++-
 .../signing/step_05_all_inputs_set.py         |  5 ++-
 .../protocol/signing/step_06_set_output.py    |  9 +++--
 .../signing/step_07_all_outputs_set.py        |  9 ++++-
 .../protocol/signing/step_09_sign_input.py    | 11 +++---
 .../monero/xmr/serialize_messages/tx_rsig.py  |  6 ---
 12 files changed, 113 insertions(+), 48 deletions(-)
 create mode 100644 src/apps/monero/protocol/signing/rct_type.py
 create mode 100644 src/apps/monero/protocol/signing/rsig_type.py
 delete mode 100644 src/apps/monero/xmr/serialize_messages/tx_rsig.py

diff --git a/src/apps/monero/controller/misc.py b/src/apps/monero/controller/misc.py
index 14f44a4f9..b82ca0257 100644
--- a/src/apps/monero/controller/misc.py
+++ b/src/apps/monero/controller/misc.py
@@ -110,3 +110,24 @@ def dump_rsig_bp(rsig):
     offset += 32
     memcpy(buff, offset, rsig.t, 0, 32)
     return buff
+
+
+def get_monero_rct_type(rct_type, rsig_type):
+    """
+    This converts our internal representation of RctType and RsigType
+    into what is used in Monero:
+    - Null = 0
+    - Full = 1
+    - Simple = 2
+    - Simple/Full with bulletproof = 3
+    """
+    from apps.monero.protocol.signing.rct_type import RctType
+    from apps.monero.protocol.signing.rsig_type import RsigType
+
+    if rsig_type == RsigType.Bulletproof:
+        return 3  # Bulletproofs
+
+    if rct_type == RctType.Simple:
+        return 2  # Simple
+    else:
+        return 1  # Full
diff --git a/src/apps/monero/protocol/signing/rct_type.py b/src/apps/monero/protocol/signing/rct_type.py
new file mode 100644
index 000000000..659d961d0
--- /dev/null
+++ b/src/apps/monero/protocol/signing/rct_type.py
@@ -0,0 +1,13 @@
+"""
+There are two types of monero Ring Confidential Transactions:
+1. RCTTypeFull = 1 (used if num_inputs == 1)
+2. RCTTypeSimple = 2 (for num_inputs > 1)
+
+There is actually also RCTTypeNull but we ignore that one.
+"""
+# TODO: to be moved somewhere else
+
+
+class RctType:
+    Full = 1
+    Simple = 2
diff --git a/src/apps/monero/protocol/signing/rsig_type.py b/src/apps/monero/protocol/signing/rsig_type.py
new file mode 100644
index 000000000..cd525ab12
--- /dev/null
+++ b/src/apps/monero/protocol/signing/rsig_type.py
@@ -0,0 +1,17 @@
+# TODO: to be moved somewhere else
+"""
+Range signature types
+
+There are four types of range proofs/signatures in official Monero:
+ 1. RangeProofBorromean = 0
+ 2. RangeProofBulletproof = 1
+ 3. RangeProofMultiOutputBulletproof = 2
+ 4. RangeProofPaddedBulletproof = 3
+
+We simplify all the bulletproofs into one.
+"""
+
+
+class RsigType:
+    Borromean = 0
+    Bulletproof = 1
diff --git a/src/apps/monero/protocol/signing/state.py b/src/apps/monero/protocol/signing/state.py
index b7534cd0d..61d83f055 100644
--- a/src/apps/monero/protocol/signing/state.py
+++ b/src/apps/monero/protocol/signing/state.py
@@ -61,9 +61,12 @@ def __init__(self, ctx):
         """
         self.need_additional_txkeys = False
 
-        # TODO to be modified
-        self.use_bulletproof = False
-        self.use_simple_rct = False
+        # Ring Confidential Transaction type
+        # allowed values: RctType.{Full, Simple}
+        self.rct_type = None
+        # Range Signature type (also called range proof)
+        # allowed values: RsigType.{Borromean, Bulletproof}
+        self.rsig_type = None
 
         self.input_count = 0
         self.output_count = 0
@@ -96,7 +99,7 @@ def __init__(self, ctx):
         # the range proofs are calculated in batches, this denotes the grouping
         self.rsig_grouping = []
         # is range proof computing offloaded or not
-        self.rsig_offload = 0
+        self.rsig_offload = False
 
         # sum of all inputs' pseudo out masks
         self.sumpouts_alphas = crypto.sc_0()
@@ -147,15 +150,3 @@ def assrt(self, condition, msg=None):
 
     def change_address(self):
         return self.output_change.addr if self.output_change else None
-
-    def get_rct_type(self):
-        """
-        RCTsig type (simple/full x Borromean/Bulletproof)
-        :return:
-        """
-        from apps.monero.xmr.serialize_messages.tx_rsig import RctType
-
-        if self.use_simple_rct:
-            return RctType.FullBulletproof if self.use_bulletproof else RctType.Simple
-        else:
-            return RctType.Full
diff --git a/src/apps/monero/protocol/signing/step_01_init_transaction.py b/src/apps/monero/protocol/signing/step_01_init_transaction.py
index 610f12f38..41e179c43 100644
--- a/src/apps/monero/protocol/signing/step_01_init_transaction.py
+++ b/src/apps/monero/protocol/signing/step_01_init_transaction.py
@@ -6,6 +6,8 @@
 
 from apps.monero.controller import misc
 from apps.monero.layout import confirms
+from apps.monero.protocol.signing.rct_type import RctType
+from apps.monero.protocol.signing.rsig_type import RsigType
 from apps.monero.protocol.signing.state import State
 from apps.monero.xmr import common, crypto, monero
 
@@ -49,7 +51,7 @@ async def init_transaction(
     if state.output_count < 2:
         raise misc.TrezorNotEnoughOutputs("At least two outputs are required")
 
-    _check_ringct_type(state, tsx_data.rsig_data)
+    _check_rsig_data(state, tsx_data.rsig_data)
     _check_subaddresses(state, tsx_data.outputs)
 
     # Extra processing, payment id
@@ -66,8 +68,10 @@ async def init_transaction(
     state.mem_trace(10, True)
 
     # Final message hasher
-    state.full_message_hasher.init(state.use_simple_rct)  # TODO investigate
-    state.full_message_hasher.set_type_fee(state.get_rct_type(), state.fee)
+    state.full_message_hasher.init(state.rct_type == RctType.Simple)
+    state.full_message_hasher.set_type_fee(
+        misc.get_monero_rct_type(state.rct_type, state.rsig_type), state.fee
+    )
 
     # Sub address precomputation
     if tsx_data.account is not None and tsx_data.minor_indices:
@@ -160,7 +164,7 @@ def _get_primary_change_address(state: State):
     )
 
 
-def _check_ringct_type(state: State, rsig_data: MoneroTransactionRsigData):
+def _check_rsig_data(state: State, rsig_data: MoneroTransactionRsigData):
     """
     There are two types of monero ring confidential transactions:
     1. RCTTypeFull = 1 (used if num_inputs == 1)
@@ -172,11 +176,23 @@ def _check_ringct_type(state: State, rsig_data: MoneroTransactionRsigData):
     3. RangeProofMultiOutputBulletproof = 2
     4. RangeProofPaddedBulletproof = 3
     """
-    # TODO maybe make this more explicit - also in the state
     state.rsig_grouping = rsig_data.grouping
-    state.rsig_offload = rsig_data.rsig_type > 0 and state.output_count > 2
-    state.use_bulletproof = rsig_data.rsig_type > 0
-    state.use_simple_rct = state.input_count > 1 or rsig_data.rsig_type != 0
+
+    if rsig_data.rsig_type == 0:
+        state.rsig_type = RsigType.Borromean
+    elif rsig_data.rsig_type in (1, 2, 3):
+        state.rsig_type = RsigType.Bulletproof
+    else:
+        raise ValueError("Unknown rsig type")
+
+    # unintuitively RctType.Simple is used for more inputs
+    if state.input_count > 1 or state.rsig_type == RsigType.Bulletproof:
+        state.rct_type = RctType.Simple
+    else:
+        state.rct_type = RctType.Full
+
+    if state.rsig_type == RsigType.Bulletproof and state.output_count > 2:
+        state.rsig_offload = True
 
 
 def _check_change(state: State, outputs: list):
@@ -275,8 +291,6 @@ def _process_payment_id(state: State, tsx_data: MoneroTransactionData):
         view_key_pub_enc = _get_key_for_payment_id_encryption(
             tsx_data.outputs, state.change_address()
         )
-        if view_key_pub_enc == crypto.NULL_KEY_ENC:
-            raise ValueError("Invalid key")  # TODO can this be removed?
 
         view_key_pub = crypto.decodepoint(view_key_pub_enc)
         payment_id_encr = _encrypt_payment_id(
@@ -332,6 +346,10 @@ def _get_key_for_payment_id_encryption(destinations: list, change_addr=None):
             )
         addr = dest.addr
         count += 1
+
+    if addr.view_public_key == crypto.NULL_KEY_ENC:
+        raise ValueError("Invalid key")
+
     return addr.view_public_key
 
 
diff --git a/src/apps/monero/protocol/signing/step_02_set_input.py b/src/apps/monero/protocol/signing/step_02_set_input.py
index 8cf188b70..1e71da44b 100644
--- a/src/apps/monero/protocol/signing/step_02_set_input.py
+++ b/src/apps/monero/protocol/signing/step_02_set_input.py
@@ -15,6 +15,7 @@
 
 from apps.monero.controller import misc
 from apps.monero.layout import confirms
+from apps.monero.protocol.signing.rct_type import RctType
 from apps.monero.xmr import crypto, monero
 
 if False:
@@ -99,7 +100,7 @@ async def set_input(state: State, src_entr: MoneroTransactionSourceEntry):
     pseudo_out_hmac = None
     alpha_enc = None
 
-    if state.use_simple_rct:
+    if state.rct_type == RctType.Simple:
         alpha, pseudo_out = _gen_commitment(state, src_entr.amount)
         pseudo_out = crypto.encodepoint(pseudo_out)
 
diff --git a/src/apps/monero/protocol/signing/step_04_input_vini.py b/src/apps/monero/protocol/signing/step_04_input_vini.py
index f279c3fa7..6ed7dca27 100644
--- a/src/apps/monero/protocol/signing/step_04_input_vini.py
+++ b/src/apps/monero/protocol/signing/step_04_input_vini.py
@@ -8,6 +8,8 @@
 
 from apps.monero.layout import confirms
 from apps.monero.protocol import hmac_encryption_keys
+from apps.monero.protocol.signing.rct_type import RctType
+from apps.monero.protocol.signing.rsig_type import RsigType
 from apps.monero.xmr import common, crypto
 
 if False:
@@ -52,8 +54,8 @@ async def input_vini(
     state.tx_prefix_hasher.buffer(vini_bin)
 
     # in monero version >= 8 pseudo outs were moved to a different place
-    # use_bulletproofs implies version >= 8
-    if state.use_simple_rct and not state.use_bulletproof:
+    # bulletproofs imply version >= 8
+    if state.rct_type == RctType.Simple and state.rsig_type != RsigType.Bulletproof:
         _hash_vini_pseudo_out(state, pseudo_out, pseudo_out_hmac)
 
     return MoneroTransactionInputViniAck()
diff --git a/src/apps/monero/protocol/signing/step_05_all_inputs_set.py b/src/apps/monero/protocol/signing/step_05_all_inputs_set.py
index b7b738188..55acb7d1c 100644
--- a/src/apps/monero/protocol/signing/step_05_all_inputs_set.py
+++ b/src/apps/monero/protocol/signing/step_05_all_inputs_set.py
@@ -8,6 +8,7 @@
 from .state import State
 
 from apps.monero.layout import confirms
+from apps.monero.protocol.signing.rct_type import RctType
 from apps.monero.xmr import crypto
 
 
@@ -27,7 +28,7 @@ async def all_inputs_set(state: State):
     for i in range(state.output_count):
         cur_mask = crypto.new_scalar()  # new mask for each output
         is_last = i + 1 == state.output_count
-        if is_last and state.use_simple_rct:
+        if is_last and state.rct_type == RctType.Simple:
             # in SimpleRCT the last mask needs to be calculated as an offset of the sum
             crypto.sc_sub_into(cur_mask, state.sumpouts_alphas, state.sumout)
         else:
@@ -36,7 +37,7 @@ async def all_inputs_set(state: State):
         crypto.sc_add_into(state.sumout, state.sumout, cur_mask)
         state.output_masks.append(cur_mask)
 
-    if state.use_simple_rct:
+    if state.rct_type == RctType.Simple:
         state.assrt(
             crypto.sc_eq(state.sumout, state.sumpouts_alphas), "Invalid masks sum"
         )  # sum check
diff --git a/src/apps/monero/protocol/signing/step_06_set_output.py b/src/apps/monero/protocol/signing/step_06_set_output.py
index 44f5e02be..5ff34a1d9 100644
--- a/src/apps/monero/protocol/signing/step_06_set_output.py
+++ b/src/apps/monero/protocol/signing/step_06_set_output.py
@@ -11,6 +11,7 @@
 from apps.monero.controller import misc
 from apps.monero.layout import confirms
 from apps.monero.protocol import hmac_encryption_keys
+from apps.monero.protocol.signing.rsig_type import RsigType
 from apps.monero.xmr import common, crypto
 
 
@@ -176,7 +177,7 @@ def _range_proof(state, amount, rsig_data):
     C, rsig = None, None
 
     state.mem_trace("pre-rproof" if __debug__ else None, collect=True)
-    if not state.rsig_offload and state.use_bulletproof:
+    if state.rsig_type == RsigType.Bulletproof and not state.rsig_offload:
         """Bulletproof calculation in trezor"""
         rsig = ring_ct.prove_range_bp_batch(state.output_amounts, state.output_masks)
         state.mem_trace("post-bp" if __debug__ else None, collect=True)
@@ -193,7 +194,7 @@ def _range_proof(state, amount, rsig_data):
             "post-bp-ser, size: %s" % len(rsig) if __debug__ else None, collect=True
         )
 
-    elif not state.rsig_offload and not state.use_bulletproof:
+    elif state.rsig_type == RsigType.Borromean and not state.rsig_offload:
         """Borromean calculation in trezor"""
         C, mask, rsig = ring_ct.prove_range_chunked(amount, mask)
         del (ring_ct)
@@ -202,7 +203,7 @@ def _range_proof(state, amount, rsig_data):
         state.full_message_hasher.rsig_val(rsig, False, raw=True)
         _check_out_commitment(state, amount, mask, C)
 
-    elif state.rsig_offload and state.use_bulletproof:
+    elif state.rsig_type == RsigType.Bulletproof and state.rsig_offload:
         """Bulletproof calculated on host, verify in trezor"""
         from apps.monero.xmr.serialize_messages.tx_rsig_bulletproof import Bulletproof
 
@@ -223,7 +224,7 @@ def _range_proof(state, amount, rsig_data):
         state.mem_trace("BP verified" if __debug__ else None, collect=True)
         del (bp_obj, ring_ct)
 
-    elif state.rsig_offload and not state.use_bulletproof:
+    elif state.rsig_type == RsigType.Borromean and state.rsig_offload:
         """Borromean offloading not supported"""
         raise misc.TrezorError(
             "Unsupported rsig state (Borromean offloaded is not supported)"
diff --git a/src/apps/monero/protocol/signing/step_07_all_outputs_set.py b/src/apps/monero/protocol/signing/step_07_all_outputs_set.py
index b3b043ad1..b51d0a46f 100644
--- a/src/apps/monero/protocol/signing/step_07_all_outputs_set.py
+++ b/src/apps/monero/protocol/signing/step_07_all_outputs_set.py
@@ -8,6 +8,7 @@
 
 from .state import State
 
+from apps.monero.controller import misc
 from apps.monero.layout import confirms
 from apps.monero.xmr import crypto
 
@@ -46,7 +47,9 @@ async def all_outputs_set(state: State):
 
     # Initializes RCTsig structure (fee, tx prefix hash, type)
     rv_pb = MoneroRingCtSig(
-        txn_fee=state.fee, message=state.tx_prefix_hash, rv_type=state.get_rct_type()
+        txn_fee=state.fee,
+        message=state.tx_prefix_hash,
+        rv_type=misc.get_monero_rct_type(state.rct_type, state.rsig_type),
     )
 
     return MoneroTransactionAllOutSetAck(
@@ -55,11 +58,13 @@ async def all_outputs_set(state: State):
 
 
 def _validate(state: State):
+    from apps.monero.protocol.signing.rct_type import RctType
+
     if state.current_output_index + 1 != state.output_count:
         raise ValueError("Invalid out num")
 
     # Test if \sum Alpha == \sum A
-    if state.use_simple_rct:
+    if state.rct_type == RctType.Simple:
         state.assrt(crypto.sc_eq(state.sumout, state.sumpouts_alphas))
 
     # Fee test
diff --git a/src/apps/monero/protocol/signing/step_09_sign_input.py b/src/apps/monero/protocol/signing/step_09_sign_input.py
index 637ecf227..f39c2a762 100644
--- a/src/apps/monero/protocol/signing/step_09_sign_input.py
+++ b/src/apps/monero/protocol/signing/step_09_sign_input.py
@@ -8,6 +8,7 @@
 
 from apps.monero.controller import misc
 from apps.monero.layout import confirms
+from apps.monero.protocol.signing.rct_type import RctType
 from apps.monero.xmr import common, crypto
 
 if False:
@@ -47,11 +48,11 @@ async def sign_input(
     state.current_input_index += 1
     if state.current_input_index >= state.input_count:
         raise ValueError("Invalid inputs count")
-    if state.use_simple_rct and pseudo_out is None:
+    if state.rct_type == RctType.Simple and pseudo_out is None:
         raise ValueError("SimpleRCT requires pseudo_out but none provided")
-    if state.use_simple_rct and pseudo_out_alpha_enc is None:
+    if state.rct_type == RctType.Simple and pseudo_out_alpha_enc is None:
         raise ValueError("SimpleRCT requires pseudo_out's mask but none provided")
-    if state.current_input_index >= 1 and not state.use_simple_rct:
+    if state.current_input_index >= 1 and not state.rct_type == RctType.Simple:
         raise ValueError("Two and more inputs must imply SimpleRCT")
 
     input_position = state.source_permutation[state.current_input_index]
@@ -66,7 +67,7 @@ async def sign_input(
     gc.collect()
     state.mem_trace(1)
 
-    if state.use_simple_rct:
+    if state.rct_type == RctType.Simple:
         # both pseudo_out and its mask were offloaded so we need to
         # validate pseudo_out's HMAC and decrypt the alpha
         pseudo_out_hmac_comp = crypto.compute_hmac(
@@ -130,7 +131,7 @@ async def sign_input(
     # RCT signature
     from apps.monero.xmr import mlsag2
 
-    if state.use_simple_rct:
+    if state.rct_type == RctType.Simple:
         # Simple RingCT
         mix_ring = [x.key for x in src_entr.outputs]
         mg, msc = mlsag2.prove_rct_mg_simple(
diff --git a/src/apps/monero/xmr/serialize_messages/tx_rsig.py b/src/apps/monero/xmr/serialize_messages/tx_rsig.py
deleted file mode 100644
index ca6eff911..000000000
--- a/src/apps/monero/xmr/serialize_messages/tx_rsig.py
+++ /dev/null
@@ -1,6 +0,0 @@
-class RctType:
-    Null = 0
-    Full = 1
-    Simple = 2
-    FullBulletproof = 3
-    SimpleBulletproof = 4