From 074d3c28851909401f75d37b04eafeec2a5a7a37 Mon Sep 17 00:00:00 2001
From: Salvatore Ingala <6681844+bigspider@users.noreply.github.com>
Date: Fri, 22 Sep 2023 10:41:46 +0000
Subject: [PATCH] Warn if fees are above 10% of total amount

---
 src/handler/sign_psbt.c                       |  9 +++++
 src/ui/display.c                              |  6 ++++
 src/ui/display.h                              |  4 +++
 src/ui/display_bagl.c                         | 23 ++++++++++++
 src/ui/display_nbgl.c                         |  9 +++++
 ...gn_with_default_wallet_accept_highfee.json | 35 +++++++++++++++++++
 tests/test_sign_psbt.py                       | 35 +++++++++++++++++--
 tests/test_sign_psbt_v1.py                    |  8 ++---
 8 files changed, 122 insertions(+), 7 deletions(-)
 create mode 100644 tests/automations/sign_with_default_wallet_accept_highfee.json

diff --git a/src/handler/sign_psbt.c b/src/handler/sign_psbt.c
index e7b79ed0a..646a39176 100644
--- a/src/handler/sign_psbt.c
+++ b/src/handler/sign_psbt.c
@@ -1328,6 +1328,15 @@ confirm_transaction(dispatcher_context_t *dc, sign_psbt_state_t *st) {
             return false;
         }
     } else {
+        // if the value of fees is 10% or more of the amount, and it's more than 10000
+        if (10 * fee >= st->inputs_total_amount && st->inputs_total_amount > 10000) {
+            if (!ui_warn_high_fee(dc)) {
+                SEND_SW(dc, SW_DENY);
+                ui_post_processing_confirm_transaction(dc, false);
+                return false;
+            }
+        }
+
         // Show final user validation UI
         bool is_self_transfer = st->outputs.n_external == 0;
         if (!ui_validate_transaction(dc, COIN_COINID_SHORT, fee, is_self_transfer)) {
diff --git a/src/ui/display.c b/src/ui/display.c
index a4b00178c..8c6ce766e 100644
--- a/src/ui/display.c
+++ b/src/ui/display.c
@@ -199,6 +199,12 @@ bool ui_validate_output(dispatcher_context_t *context,
     return io_ui_process(context, true);
 }
 
+bool ui_warn_high_fee(dispatcher_context_t *context) {
+    ui_warn_high_fee_flow();
+
+    return io_ui_process(context, true);
+}
+
 bool ui_validate_transaction(dispatcher_context_t *context,
                              const char *coin_name,
                              uint64_t fee,
diff --git a/src/ui/display.h b/src/ui/display.h
index d35321396..b8e558693 100644
--- a/src/ui/display.h
+++ b/src/ui/display.h
@@ -132,6 +132,8 @@ bool ui_validate_output(dispatcher_context_t *context,
                         const char *coin_name,
                         uint64_t amount);
 
+bool ui_warn_high_fee(dispatcher_context_t *context);
+
 bool ui_validate_transaction(dispatcher_context_t *context,
                              const char *coin_name,
                              uint64_t fee,
@@ -165,6 +167,8 @@ void ui_display_output_address_amount_flow(int index);
 
 void ui_display_output_address_amount_no_index_flow(int index);
 
+void ui_warn_high_fee_flow();
+
 void ui_accept_transaction_flow(bool is_self_transfer);
 
 void ui_display_transaction_prompt(const int external_outputs_total_count);
diff --git a/src/ui/display_bagl.c b/src/ui/display_bagl.c
index 17edf831c..dd6375666 100644
--- a/src/ui/display_bagl.c
+++ b/src/ui/display_bagl.c
@@ -167,6 +167,15 @@ UX_STEP_NOCB(ux_validate_address_step,
                  .text = g_ui_state.validate_output.address_or_description,
              });
 
+// Step with eye icon and "High fees!" warning, with the percentage
+UX_STEP_NOCB(ux_high_fee_step,
+             pnn,
+             {
+                 &C_icon_eye,
+                 "Fees are above 10%",
+                 "of total amount!",
+             });
+
 UX_STEP_NOCB(ux_confirm_transaction_step, pnn, {&C_icon_eye, "Confirm", "transaction"});
 UX_STEP_NOCB(ux_confirm_selftransfer_step, pnn, {&C_icon_eye, "Confirm", "self-transfer"});
 UX_STEP_NOCB(ux_confirm_transaction_fees_step,
@@ -386,6 +395,16 @@ UX_FLOW(ux_display_output_address_amount_flow,
         &ux_display_approve_step,
         &ux_display_reject_step);
 
+// Finalize see the transaction fees and finally accept signing
+// #1 screen: eye icon + "Confirm transaction"
+// #2 screen: fee amount
+// #3 screen: "Continue", with approve button
+// #4 screen: reject button
+UX_FLOW(ux_warn_high_fee_flow,
+        &ux_high_fee_step,
+        &ux_display_continue_step,
+        &ux_display_reject_step);
+
 // Finalize see the transaction fees and finally accept signing
 // #1 screen: eye icon + "Confirm transaction"
 // #2 screen: fee amount
@@ -466,6 +485,10 @@ void ui_display_output_address_amount_no_index_flow(int index) {
     ui_display_output_address_amount_flow(index);
 }
 
+void ui_warn_high_fee_flow() {
+    ux_flow_init(0, ux_warn_high_fee_flow, NULL);
+}
+
 void ui_accept_transaction_flow(bool is_self_transfer) {
     ux_flow_init(0,
                  is_self_transfer ? ux_accept_selftransfer_flow : ux_accept_transaction_flow,
diff --git a/src/ui/display_nbgl.c b/src/ui/display_nbgl.c
index b921fcefb..3799fa4fe 100644
--- a/src/ui/display_nbgl.c
+++ b/src/ui/display_nbgl.c
@@ -183,6 +183,15 @@ static void transaction_confirm(int token, uint8_t index) {
     }
 }
 
+void ui_warn_high_fee_flow() {
+    nbgl_useCaseChoice(&C_round_warning_64px,
+                       "Warning",
+                       "Fees are above 10%\n of total amount",
+                       "Continue",
+                       "Reject",
+                       ux_flow_response);
+}
+
 void ui_accept_transaction_flow(bool is_self_transfer) {
     if (!is_self_transfer) {
         transactionContext.tagValuePair[0].item = "Fees";
diff --git a/tests/automations/sign_with_default_wallet_accept_highfee.json b/tests/automations/sign_with_default_wallet_accept_highfee.json
new file mode 100644
index 000000000..78cb80505
--- /dev/null
+++ b/tests/automations/sign_with_default_wallet_accept_highfee.json
@@ -0,0 +1,35 @@
+{
+  "version": 1,
+  "rules": [
+    {
+      "regexp": "Approve|Accept|Hold to sign|Continue",
+      "actions": [
+        [ "button", 1, true ],
+        [ "button", 2, true ],
+        [ "button", 2, false ],
+        [ "button", 1, false ],
+        [ "finger", 55, 550, true]
+      ]
+    },
+    {
+      "regexp": "Review|Amount|Address|Confirm|Fees",
+      "actions": [
+        ["button", 2, true],
+        ["button", 2, false]
+      ]
+    },
+    {
+      "regexp": "Tap to continue",
+      "actions": [
+        ["finger", 55, 550, true],
+        ["finger", 55, 550, false]
+      ]
+    },
+    {
+      "regexp": "SIGNED",
+      "actions": [
+        ["finger", 55, 550, false]
+      ]
+    }
+  ]
+}
diff --git a/tests/test_sign_psbt.py b/tests/test_sign_psbt.py
index 4eb8aef61..9fe9487cf 100644
--- a/tests/test_sign_psbt.py
+++ b/tests/test_sign_psbt.py
@@ -34,8 +34,8 @@
 def format_amount(ticker: str, amount: int) -> str:
     """Formats an amounts in sats as shown in the app: divided by 10_000_000, with no trailing zeroes."""
     assert amount >= 0
-
-    return f"{ticker} {str(Decimal(amount) / 100_000_000)}"
+    btc_amount = f"{(amount/100_000_000):.8f}".rstrip('0').rstrip('.')
+    return f"{ticker} {btc_amount}"
 
 
 def should_go_right(event: dict):
@@ -212,6 +212,34 @@ def test_sign_psbt_singlesig_sh_wpkh_1to2(client: Client):
     )]
 
 
+@has_automation("automations/sign_with_default_wallet_accept_highfee.json")
+def test_sign_psbt_highfee(client: Client):
+    # Transactions with fees higher than 10% of total amount
+    # An aditional warning is shown.
+
+    # PSBT for a wrapped segwit 1-input 2-output spend (1 change address)
+    psbt = open_psbt_from_file(f"{tests_root}/psbt/singlesig/sh-wpkh-1to2.psbt")
+
+    # Make sure that the fees are at least 10% of the total amount
+    for out in psbt.tx.vout:
+        out.nValue = int(out.nValue * 0.9)
+
+    # the test is only interesting if the total amount is at least 10000 sats
+    assert sum(input.witness_utxo.nValue for input in psbt.inputs) >= 10000
+
+    wallet = WalletPolicy(
+        "",
+        "sh(wpkh(@0/**))",
+        [
+            "[f5acc2fd/49'/1'/0']tpubDC871vGLAiKPcwAw22EjhKVLk5L98UGXBEcGR8gpcigLQVDDfgcYW24QBEyTHTSFEjgJgbaHU8CdRi9vmG4cPm1kPLmZhJEP17FMBdNheh3"
+        ],
+    )
+
+    result = client.sign_psbt(psbt, wallet, None)
+
+    assert len(result) == 1
+
+
 @has_automation("automations/sign_with_default_wallet_accept.json")
 def test_sign_psbt_singlesig_wpkh_1to2(client: Client):
 
@@ -586,7 +614,8 @@ def test_sign_psbt_singlesig_wpkh_4to3(client: Client, comm: SpeculosClient, is_
     n_outs = 3
 
     in_amounts = [10000 + 10000 * i for i in range(n_ins)]
-    out_amounts = [9999 + 9999 * i for i in range(n_outs)]
+    total_in = sum(in_amounts)
+    out_amounts = [total_in // n_outs - i for i in range(n_outs)]
 
     change_index = 1
 
diff --git a/tests/test_sign_psbt_v1.py b/tests/test_sign_psbt_v1.py
index ced263fe6..d22d91d13 100644
--- a/tests/test_sign_psbt_v1.py
+++ b/tests/test_sign_psbt_v1.py
@@ -36,8 +36,8 @@
 def format_amount(ticker: str, amount: int) -> str:
     """Formats an amounts in sats as shown in the app: divided by 10_000_000, with no trailing zeroes."""
     assert amount >= 0
-
-    return f"{ticker} {str(Decimal(amount) / 100_000_000)}"
+    btc_amount = f"{(amount/100_000_000):.8f}".rstrip('0').rstrip('.')
+    return f"{ticker} {btc_amount}"
 
 
 def should_go_right(event: dict):
@@ -387,7 +387,8 @@ def test_sign_psbt_singlesig_wpkh_4to3_v1(client: Client, comm: SpeculosClient,
     n_outs = 3
 
     in_amounts = [10000 + 10000 * i for i in range(n_ins)]
-    out_amounts = [9999 + 9999 * i for i in range(n_outs)]
+    sum_in = sum(in_amounts)
+    out_amounts = [sum_in // n_outs - i for i in range(n_outs)]
 
     change_index = 1
 
@@ -398,7 +399,6 @@ def test_sign_psbt_singlesig_wpkh_4to3_v1(client: Client, comm: SpeculosClient,
         [i == change_index for i in range(n_outs)]
     )
 
-    sum_in = sum(in_amounts)
     sum_out = sum(out_amounts)
 
     assert sum_out < sum_in