From 24bf3525774c8201d6687d27f1b183c5e95583ea Mon Sep 17 00:00:00 2001
From: Andrew Kozlik <andrew.kozlik@gmail.com>
Date: Wed, 8 Jul 2020 18:43:03 +0200
Subject: [PATCH] legacy: Support multiple change-outputs.

---
 legacy/firmware/layout2.c | 10 ++++++++++
 legacy/firmware/layout2.h |  1 +
 legacy/firmware/signing.c | 36 +++++++++++++++++++++++++++++-------
 3 files changed, 40 insertions(+), 7 deletions(-)

diff --git a/legacy/firmware/layout2.c b/legacy/firmware/layout2.c
index 71360f5ad5..9bd88ebd98 100644
--- a/legacy/firmware/layout2.c
+++ b/legacy/firmware/layout2.c
@@ -18,7 +18,9 @@
  */
 
 #include <ctype.h>
+#include <inttypes.h>
 #include <stdint.h>
+#include <stdio.h>
 #include <string.h>
 
 #include "bignum.h"
@@ -440,6 +442,14 @@ void layoutFeeOverThreshold(const CoinInfo *coin, uint64_t fee) {
                     _("Send anyway?"), NULL);
 }
 
+void layoutChangeCountOverThreshold(uint32_t change_count) {
+  char str_change[21] = {0};
+  snprintf(str_change, sizeof(str_change), "There are %" PRIu32, change_count);
+  layoutDialogSwipe(&bmp_icon_question, _("Cancel"), _("Confirm"), NULL,
+                    _("Warning!"), str_change, _("change-outputs."), NULL,
+                    _("Continue?"), NULL);
+}
+
 void layoutSignMessage(const uint8_t *msg, uint32_t len) {
   const char **str = NULL;
   if (!is_valid_ascii(msg, len)) {
diff --git a/legacy/firmware/layout2.h b/legacy/firmware/layout2.h
index 3830596361..b5b2191c74 100644
--- a/legacy/firmware/layout2.h
+++ b/legacy/firmware/layout2.h
@@ -52,6 +52,7 @@ void layoutConfirmOpReturn(const uint8_t *data, uint32_t size);
 void layoutConfirmTx(const CoinInfo *coin, uint64_t amount_out,
                      uint64_t amount_fee);
 void layoutFeeOverThreshold(const CoinInfo *coin, uint64_t fee);
+void layoutChangeCountOverThreshold(uint32_t change_count);
 void layoutSignMessage(const uint8_t *msg, uint32_t len);
 void layoutVerifyAddress(const CoinInfo *coin, const char *address);
 void layoutVerifyMessage(const uint8_t *msg, uint32_t len);
diff --git a/legacy/firmware/signing.c b/legacy/firmware/signing.c
index 7445dff083..4654ac2571 100644
--- a/legacy/firmware/signing.c
+++ b/legacy/firmware/signing.c
@@ -32,6 +32,7 @@
 
 static uint32_t inputs_count;
 static uint32_t outputs_count;
+static uint32_t change_count;
 static const CoinInfo *coin;
 static CONFIDENTIAL HDNode root;
 static CONFIDENTIAL HDNode node;
@@ -103,6 +104,9 @@ static uint32_t tx_weight;
 /* transaction segwit overhead 2 marker */
 #define TXSIZE_SEGWIT_OVERHEAD 2
 
+/* The maximum number of change-outputs allowed without user confirmation. */
+#define MAX_SILENT_CHANGE_COUNT 2
+
 enum {
   SIGHASH_ALL = 1,
   SIGHASH_FORKID = 0x40,
@@ -551,6 +555,7 @@ void signing_init(const SignTx *msg, const CoinInfo *_coin,
   to_spend = 0;
   spending = 0;
   change_spend = 0;
+  change_count = 0;
   memzero(&input, sizeof(TxInputType));
   memzero(&resp, sizeof(TxRequest));
 
@@ -840,11 +845,18 @@ static bool signing_check_output(TxOutputType *txoutput) {
   }
 
   if (is_change) {
-    if (change_spend == 0) {  // not set
-      change_spend = txoutput->amount;
-    } else {
-      /* We only skip confirmation for the first change output */
-      is_change = false;
+    if (change_spend + txoutput->amount < change_spend) {
+      fsm_sendFailure(FailureType_Failure_DataError, _("Value overflow"));
+      signing_abort();
+      return false;
+    }
+    change_spend += txoutput->amount;
+
+    change_count++;
+    if (change_count <= 0) {
+      fsm_sendFailure(FailureType_Failure_DataError, _("Value overflow"));
+      signing_abort();
+      return false;
     }
   }
 
@@ -885,7 +897,7 @@ static bool signing_check_output(TxOutputType *txoutput) {
   return true;
 }
 
-static bool signing_check_fee(void) {
+static bool signing_confirm_tx(void) {
   if (coin->negative_fee) {
     // bypass check for negative fee coins, required for reward TX
   } else {
@@ -912,6 +924,16 @@ static bool signing_check_fee(void) {
   } else {
     fee = 0;
   }
+
+  if (change_count > MAX_SILENT_CHANGE_COUNT) {
+    layoutChangeCountOverThreshold(change_count);
+    if (!protectButton(ButtonRequestType_ButtonRequest_SignTx, false)) {
+      fsm_sendFailure(FailureType_Failure_ActionCancelled, NULL);
+      signing_abort();
+      return false;
+    }
+  }
+
   // last confirmation
   layoutConfirmTx(coin, to_spend - change_spend, fee);
   if (!protectButton(ButtonRequestType_ButtonRequest_SignTx, false)) {
@@ -944,7 +966,7 @@ static void phase1_request_next_output(void) {
     }
 #endif
     hasher_Final(&hasher_outputs, hash_outputs);
-    if (!signing_check_fee()) {
+    if (!signing_confirm_tx()) {
       return;
     }
     // Everything was checked, now phase 2 begins and the transaction is signed.