From dec5919864ba74ddc6ff26cba43fc4c4b6184c2b Mon Sep 17 00:00:00 2001
From: Francois Beutin <francois.beutin@ledger.fr>
Date: Wed, 1 Feb 2023 17:35:12 +0100
Subject: [PATCH] Don't display main menu during swap

---
 src/getPubkey.c           | 4 ++--
 src/signMessage.c         | 9 +++------
 src/signOffchainMessage.c | 4 ++--
 src/utils.c               | 8 +++++---
 src/utils.h               | 2 +-
 5 files changed, 13 insertions(+), 14 deletions(-)

diff --git a/src/getPubkey.c b/src/getPubkey.c
index 5cb6bb2c..e7de168b 100644
--- a/src/getPubkey.c
+++ b/src/getPubkey.c
@@ -28,14 +28,14 @@ UX_STEP_NOCB(ux_display_public_flow_5_step,
              });
 UX_STEP_CB(ux_display_public_flow_6_step,
            pb,
-           sendResponse(set_result_get_pubkey(), true),
+           sendResponse(set_result_get_pubkey(), true, true),
            {
                &C_icon_validate_14,
                "Approve",
            });
 UX_STEP_CB(ux_display_public_flow_7_step,
            pb,
-           sendResponse(0, false),
+           sendResponse(0, false, true),
            {
                &C_icon_crossmark,
                "Reject",
diff --git a/src/signMessage.c b/src/signMessage.c
index 5391adaa..128ead41 100644
--- a/src/signMessage.c
+++ b/src/signMessage.c
@@ -45,22 +45,19 @@ static uint8_t set_result_sign_message() {
     return SIGNATURE_LENGTH;
 }
 
-static void send_result_sign_message(void) {
-    sendResponse(set_result_sign_message(), true);
-}
 
 //////////////////////////////////////////////////////////////////////
 
 UX_STEP_CB(ux_approve_step,
            pb,
-           send_result_sign_message(),
+           sendResponse(set_result_sign_message(), true, true),
            {
                &C_icon_validate_14,
                "Approve",
            });
 UX_STEP_CB(ux_reject_step,
            pb,
-           sendResponse(0, false),
+           sendResponse(0, false, true),
            {
                &C_icon_crossmark,
                "Reject",
@@ -225,7 +222,7 @@ void handle_sign_message_ui(volatile unsigned int *flags) {
             }
             if (check_swap_validity(summary_step_kinds, num_summary_steps)) {
                 PRINTF("Valid swap transaction signed\n");
-                send_result_sign_message();
+                sendResponse(set_result_sign_message(), true, false);
             } else {
                 PRINTF("Refused signing incorrect Swap transaction\n");
                 THROW(ApduReplySolanaSummaryFinalizeFailed);
diff --git a/src/signOffchainMessage.c b/src/signOffchainMessage.c
index 38aaf995..1ed86aa8 100644
--- a/src/signOffchainMessage.c
+++ b/src/signOffchainMessage.c
@@ -117,14 +117,14 @@ UX_STEP_NOCB(ux_sign_msg_text_step,
              });
 UX_STEP_CB(ux_sign_msg_approve_step,
            pb,
-           sendResponse(set_result_sign_message(), true),
+           sendResponse(set_result_sign_message(), true, true),
            {
                &C_icon_validate_14,
                "Approve",
            });
 UX_STEP_CB(ux_sign_msg_reject_step,
            pb,
-           sendResponse(0, false),
+           sendResponse(0, false, true),
            {
                &C_icon_crossmark,
                "Reject",
diff --git a/src/utils.c b/src/utils.c
index 258a1fcc..6a4e0b30 100644
--- a/src/utils.c
+++ b/src/utils.c
@@ -123,13 +123,15 @@ int read_derivation_path(const uint8_t *data_buffer,
     return 0;
 }
 
-void sendResponse(uint8_t tx, bool approve) {
+void sendResponse(uint8_t tx, bool approve, bool display_menu) {
     G_io_apdu_buffer[tx++] = approve ? 0x90 : 0x69;
     G_io_apdu_buffer[tx++] = approve ? 0x00 : 0x85;
     // Send back the response, do not restart the event loop
     io_exchange(CHANNEL_APDU | IO_RETURN_AFTER_TX, tx);
-    // Display back the original UX
-    ui_idle();
+    if (display_menu) {
+        // Display back the original UX
+        ui_idle();
+    }
 }
 
 unsigned int ui_prepro(const bagl_element_t *element) {
diff --git a/src/utils.h b/src/utils.h
index af8f2889..a5c2ce56 100644
--- a/src/utils.h
+++ b/src/utils.h
@@ -53,7 +53,7 @@ int read_derivation_path(const uint8_t *data_buffer,
                          uint32_t *derivation_path,
                          uint32_t *derivation_path_length);
 
-void sendResponse(uint8_t tx, bool approve);
+void sendResponse(uint8_t tx, bool approve, bool display_menu);
 
 // type            userid    x    y   w    h  str rad fill      fg        bg      fid iid  txt
 // touchparams...       ]