From f334d8ad68afb5fc61cfff70ca72ddeb407b412c Mon Sep 17 00:00:00 2001
From: Dusan Klinec <dusan.klinec@gmail.com>
Date: Wed, 15 Aug 2018 23:18:42 +0200
Subject: [PATCH] xmr: protocol: simplification - require change address to
 equal the main address

---
 src/apps/monero/controller/misc.py            |  5 +++
 src/apps/monero/protocol/tsx_sign_builder.py  | 33 +++++++++++++++++--
 .../monero/protocol/tsx_sign_state_holder.py  |  1 +
 src/apps/monero/xmr/monero.py                 | 19 +++++++++++
 4 files changed, 55 insertions(+), 3 deletions(-)

diff --git a/src/apps/monero/controller/misc.py b/src/apps/monero/controller/misc.py
index 761e17c0c..22ab20543 100644
--- a/src/apps/monero/controller/misc.py
+++ b/src/apps/monero/controller/misc.py
@@ -20,6 +20,11 @@ def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
 
 
+class TrezorChangeAddressError(TrezorError):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+
 class StdObj(object):
     def __init__(self, **kwargs):
         for kw in kwargs:
diff --git a/src/apps/monero/protocol/tsx_sign_builder.py b/src/apps/monero/protocol/tsx_sign_builder.py
index c6fcb7beb..d9f366651 100644
--- a/src/apps/monero/protocol/tsx_sign_builder.py
+++ b/src/apps/monero/protocol/tsx_sign_builder.py
@@ -50,6 +50,7 @@ def __init__(self, trezor=None, creds=None, state=None, **kwargs):
         self.output_change = None
         self.mixin = 0
         self.fee = 0
+        self.account_idx = 0
 
         self.additional_tx_private_keys = []
         self.additional_tx_public_keys = []
@@ -179,9 +180,25 @@ def gen_r(self, use_r=None):
         self.r = crypto.random_scalar() if use_r is None else use_r
         self.r_pub = crypto.scalarmult_base(self.r)
 
+    def get_primary_change_address(self):
+        """
+        Computes primary change address for the current account index
+        :return:
+        """
+        D, C = monero.generate_sub_address_keys(
+            self.creds.view_key_private,
+            self.creds.spend_key_public,
+            self.account_idx,
+            0,
+        )
+        return misc.StdObj(
+            view_public_key=crypto.encodepoint(C),
+            spend_public_key=crypto.encodepoint(D),
+        )
+
     def check_change(self, outputs):
         """
-        Checks if the change address is among tx outputs.
+        Checks if the change address is among tx outputs and it is equal to our address.
         :param outputs:
         :return:
         """
@@ -191,11 +208,20 @@ def check_change(self, outputs):
         if change_addr is None:
             return
 
+        found = False
         for out in outputs:
             if addr_eq(out.addr, change_addr):
-                return True
+                found = True
+                break
+
+        if not found:
+            raise misc.TrezorChangeAddressError("Change address not found in outputs")
+
+        my_addr = self.get_primary_change_address()
+        if not addr_eq(my_addr, change_addr):
+            raise misc.TrezorChangeAddressError("Change address differs from ours")
 
-        raise ValueError("Change address not found in outputs")
+        return True
 
     def in_memory(self):
         """
@@ -453,6 +479,7 @@ async def init_transaction(self, tsx_data, tsx_ctr):
         self.output_change = misc.dst_entry_to_stdobj(tsx_data.change_dts)
         self.mixin = tsx_data.mixin
         self.fee = tsx_data.fee
+        self.account_idx = tsx_data.account
         self.use_simple_rct = self.input_count > 1
         self.use_bulletproof = tsx_data.is_bulletproof
         self.multi_sig = tsx_data.is_multisig
diff --git a/src/apps/monero/protocol/tsx_sign_state_holder.py b/src/apps/monero/protocol/tsx_sign_state_holder.py
index 3d0cd12fa..1d1797d93 100644
--- a/src/apps/monero/protocol/tsx_sign_state_holder.py
+++ b/src/apps/monero/protocol/tsx_sign_state_holder.py
@@ -26,6 +26,7 @@ def __init__(self, **kwargs):
         self.output_change = None
         self.mixin = 0
         self.fee = 0
+        self.account_idx = 0
 
         self.additional_tx_private_keys = []
         self.additional_tx_public_keys = []
diff --git a/src/apps/monero/xmr/monero.py b/src/apps/monero/xmr/monero.py
index 9aca64c73..4070db42d 100644
--- a/src/apps/monero/xmr/monero.py
+++ b/src/apps/monero/xmr/monero.py
@@ -301,3 +301,22 @@ def generate_monero_keys(seed):
     hash = crypto.cn_fast_hash(crypto.encodeint(spend_sec))
     view_sec, view_pub = generate_keys(crypto.decodeint(hash))
     return spend_sec, spend_pub, view_sec, view_pub
+
+
+def generate_sub_address_keys(view_sec, spend_pub, major, minor):
+    """
+    Computes generic public sub-address
+    :param view_sec:
+    :param spend_pub:
+    :param major:
+    :param minor:
+    :return: spend public, view public
+    """
+    if major == 0 and minor == 0:  # special case, Monero-defined
+        return spend_pub, crypto.scalarmult_base(view_sec)
+
+    m = get_subaddress_secret_key(view_sec, major=major, minor=minor)
+    M = crypto.scalarmult_base(m)
+    D = crypto.point_add(spend_pub, M)
+    C = crypto.ge_scalarmult(view_sec, D)
+    return D, C