Using classical NBGL API for message signing

diff --git a/src/handler/sign_message.c b/src/handler/sign_message.c
index a8f1427..acf8a24 100644
--- a/src/handler/sign_message.c
+++ b/src/handler/sign_message.c
@@ -30,80 +30,12 @@
 
 #include "handlers.h"
 
-#define MAX_DISPLAYBLE_CHUNK_NUMBER \
-    (5 * MESSAGE_CHUNK_PER_DISPLAY)  // If the message is too long we will not display it
-
 extern const char GA_LOADING_MESSAGE[];
 
 static unsigned char const BSM_SIGN_MAGIC[] = {'\x18', 'B', 'i', 't', 'c', 'o', 'i', 'n', ' ',
                                                'S',    'i', 'g', 'n', 'e', 'd', ' ', 'M', 'e',
                                                's',    's', 'a', 'g', 'e', ':', '\n'};
-
-static bool display_message_content_and_confirm(dispatcher_context_t* dc,
-                                                uint8_t* message_merkle_root,
-                                                size_t n_chunks,
-                                                uint8_t* path_str) {
-    reset_streaming_index();
-    while (get_streaming_index() <= (n_chunks - 1) / MESSAGE_CHUNK_PER_DISPLAY) {
-        uint8_t message_chunk[MESSAGE_MAX_DISPLAY_SIZE];
-
-        int total_chunk_len = 0;
-        uint8_t offset = 0;
-
-        if (get_streaming_index() > 0) {
-            message_chunk[offset++] = '.';
-            message_chunk[offset++] = '.';
-            message_chunk[offset++] = '.';
-        }
-
-        total_chunk_len += offset;
-
-        // each UX display will show MESSAGE_CHUNK_PER_DISPLAY chunks
-        size_t group_start_index = get_streaming_index() * MESSAGE_CHUNK_PER_DISPLAY;
-
-        for (int j = 0;
-             j < MESSAGE_CHUNK_PER_DISPLAY &&
-             (group_start_index + j) < (unsigned int) n_chunks;  // make sure not to overflow
-             j++) {
-            offset += j * MESSAGE_CHUNK_SIZE;
-
-            int chunk_len =
-                call_get_merkle_leaf_element(dc,
-                                             message_merkle_root,
-                                             n_chunks,
-                                             get_streaming_index() * MESSAGE_CHUNK_PER_DISPLAY + j,
-                                             message_chunk + offset,
-                                             MESSAGE_CHUNK_SIZE);
-
-            total_chunk_len += chunk_len;
-
-            if (chunk_len < MESSAGE_CHUNK_SIZE) {
-                break;
-            }
-        }
-
-        if ((get_streaming_index() + 1) * MESSAGE_CHUNK_PER_DISPLAY < n_chunks) {
-            message_chunk[total_chunk_len] = '.';
-            message_chunk[total_chunk_len + 1] = '.';
-            message_chunk[total_chunk_len + 2] = '.';
-            message_chunk[total_chunk_len + 3] = '\0';
-        } else {
-            message_chunk[total_chunk_len] = '\0';
-        }
-
-        if (!ui_display_path_and_message_content(dc, (char*) path_str, (char*) message_chunk)) {
-            return false;
-        }
-    }
-
-    if (!ui_display_message_confirm(dc)) {
-        return false;
-    }
-
-    return true;
-}
-
-void handler_sign_message(dispatcher_context_t* dc, uint8_t protocol_version) {
+void handler_sign_message(dispatcher_context_t *dc, uint8_t protocol_version) {
     (void) protocol_version;
 
     uint8_t bip32_path_len;
@@ -141,26 +73,35 @@ void handler_sign_message(dispatcher_context_t* dc, uint8_t protocol_version) {
     crypto_hash_update(&bsm_digest_context.header, BSM_SIGN_MAGIC, sizeof(BSM_SIGN_MAGIC));
     crypto_hash_update_varint(&bsm_digest_context.header, message_length);
 
-    size_t n_chunks = (message_length + MESSAGE_CHUNK_SIZE - 1) / MESSAGE_CHUNK_SIZE;
-
-    if (n_chunks > MAX_DISPLAYBLE_CHUNK_NUMBER) {
+    // Just a flag indicating if we use only one chunk in the case of long message
+    unsigned int not_long_message = 1;
+    if (message_length > MAX_DISPLAYBLE_MESSAGE_LENGTH) {
         printable = false;
+        not_long_message = 0;
     }
 
+    uint8_t message_full[MAX_DISPLAYBLE_MESSAGE_LENGTH + 1];
+    size_t n_chunks = (message_length + MESSAGE_CHUNK_SIZE - 1) / MESSAGE_CHUNK_SIZE;
     for (unsigned int i = 0; i < n_chunks; i++) {
-        uint8_t message_chunk[MESSAGE_CHUNK_SIZE];
+        uint8_t *message_chunk = &message_full[i * MESSAGE_CHUNK_SIZE * not_long_message];
+
         int chunk_len = call_get_merkle_leaf_element(dc,
                                                      message_merkle_root,
                                                      n_chunks,
                                                      i,
                                                      message_chunk,
-                                                     sizeof(message_chunk));
+                                                     MESSAGE_CHUNK_SIZE);
 
         if (chunk_len < 0 || (chunk_len != MESSAGE_CHUNK_SIZE && i != n_chunks - 1)) {
             SEND_SW(dc, SW_BAD_STATE);  // should never happen
             return;
         }
 
+        if (i == n_chunks - 1) {
+            // Last chunk - let's add null terminator at the end
+            message_chunk[chunk_len] = '\0';
+        }
+
         if (printable) {
             for (int j = 0; j < chunk_len; j++) {
                 if (message_chunk[j] < 0x20 || message_chunk[j] > 0x7E) {
@@ -186,17 +127,13 @@ void handler_sign_message(dispatcher_context_t* dc, uint8_t protocol_version) {
     }
 
 #ifndef HAVE_AUTOAPPROVE_FOR_PERF_TESTS
-    ui_pre_processing_message();
     if (printable) {
-        if (!display_message_content_and_confirm(dc,
-                                                 message_merkle_root,
-                                                 n_chunks,
-                                                 (uint8_t*) path_str)) {
+        if (!ui_display_message_and_confirm(dc, path_str, (const char *) message_full, false)) {
             SEND_SW(dc, SW_DENY);
             return;
         }
     } else {
-        if (!ui_display_message_path_hash_and_confirm(dc, path_str, message_hash_str)) {
+        if (!ui_display_message_and_confirm(dc, path_str, message_hash_str, true)) {
             SEND_SW(dc, SW_DENY);
             return;
         }
diff --git a/src/ui/display.c b/src/ui/display.c
index 29e3ece..9e26f68 100644
--- a/src/ui/display.c
+++ b/src/ui/display.c
@@ -117,45 +117,19 @@ bool ui_display_pubkey(dispatcher_context_t *context,
     return io_ui_process(context);
 }
 
-bool ui_display_path_and_message_content(dispatcher_context_t *context,
-                                         const char *path_str,
-                                         const char *message_content) {
+bool ui_display_message_and_confirm(dispatcher_context_t *context,
+                                    const char *path_str,
+                                    const char *message,
+                                    bool is_hash) {
 #ifdef HAVE_AUTOAPPROVE_FOR_PERF_TESTS
     return true;
 #endif
 
     ui_path_and_message_state_t *state = (ui_path_and_message_state_t *) &g_ui_state;
     strncpy(state->bip32_path_str, path_str, sizeof(state->bip32_path_str));
-    strncpy(state->message, message_content, sizeof(state->message));
+    strncpy(state->message, message, sizeof(state->message));
 
-    ui_sign_message_content_flow();
-
-    return io_ui_process(context);
-}
-
-bool ui_display_message_path_hash_and_confirm(dispatcher_context_t *context,
-                                              const char *path_str,
-                                              const char *message_hash) {
-#ifdef HAVE_AUTOAPPROVE_FOR_PERF_TESTS
-    return true;
-#endif
-
-    ui_path_and_message_state_t *state = (ui_path_and_message_state_t *) &g_ui_state;
-    strncpy(state->bip32_path_str, path_str, sizeof(state->bip32_path_str));
-    strncpy(state->message, message_hash, sizeof(state->message));
-
-    ui_sign_message_path_hash_and_confirm_flow();
-
-    return io_ui_process(context);
-}
-
-bool ui_display_message_confirm(dispatcher_context_t *context) {
-#ifdef HAVE_AUTOAPPROVE_FOR_PERF_TESTS
-    return true;
-#endif
-
-    UNUSED(context);
-    ui_sign_message_confirm_flow();
+    ui_sign_message_and_confirm_flow(is_hash);
 
     return io_ui_process(context);
 }
@@ -395,10 +369,6 @@ bool ui_post_processing_confirm_message(dispatcher_context_t *context, bool succ
     return true;
 }
 
-void ui_pre_processing_message(void) {
-    ui_set_display_prompt();
-}
-
 char const *ui_get_processing_screen_text(void) {
     return (G_processing_screen_text != NULL) ? G_processing_screen_text : "Loading";
 }
diff --git a/src/ui/display.h b/src/ui/display.h
index e8b837d..f33a1f0 100644
--- a/src/ui/display.h
+++ b/src/ui/display.h
@@ -15,10 +15,9 @@
 #include "../common/script.h"
 #include "../constants.h"
 
-#define MESSAGE_CHUNK_SIZE        64  // Protocol specific
-#define MESSAGE_CHUNK_PER_DISPLAY 2   // This could be changed depending on screen sizes
-#define MESSAGE_MAX_DISPLAY_SIZE \
-    (MESSAGE_CHUNK_SIZE * MESSAGE_CHUNK_PER_DISPLAY + 2 * sizeof("...") - 1)
+#define MESSAGE_CHUNK_SIZE 64  // Protocol specific
+// Displayed message length - if the message is too long we will not display it
+#define MAX_DISPLAYBLE_MESSAGE_LENGTH (10 * MESSAGE_CHUNK_SIZE)
 
 #if defined(TARGET_STAX) || defined(TARGET_FLEX)
 #define ICON_APP_IMPORTANT IMPORTANT_CIRCLE_ICON
@@ -69,7 +68,7 @@ typedef struct {
 
 typedef struct {
     char bip32_path_str[MAX_SERIALIZED_BIP32_PATH_LENGTH + 1];
-    char message[MESSAGE_MAX_DISPLAY_SIZE];
+    char message[MAX_DISPLAYBLE_MESSAGE_LENGTH + 1];
 } ui_path_and_message_state_t;
 
 typedef struct {
@@ -141,15 +140,10 @@ bool ui_display_pubkey(dispatcher_context_t *context,
                        bool is_path_suspicious,
                        const char *pubkey);
 
-bool ui_display_path_and_message_content(dispatcher_context_t *context,
-                                         const char *path_str,
-                                         const char *message_content);
-
-bool ui_display_message_path_hash_and_confirm(dispatcher_context_t *context,
-                                              const char *path_str,
-                                              const char *message_hash);
-
-bool ui_display_message_confirm(dispatcher_context_t *context);
+bool ui_display_message_and_confirm(dispatcher_context_t *context,
+                                    const char *path_str,
+                                    const char *message,
+                                    bool is_hash);
 
 bool ui_display_address(dispatcher_context_t *dispatcher_context,
                         const char *address,
@@ -211,11 +205,7 @@ void ui_display_pubkey_flow(void);
 
 void ui_display_pubkey_suspicious_flow(void);
 
-void ui_sign_message_path_hash_and_confirm_flow(void);
-
-void ui_sign_message_content_flow(void);
-
-void ui_sign_message_confirm_flow(void);
+void ui_sign_message_and_confirm_flow(bool is_hash);
 
 void ui_display_receive_in_wallet_flow(void);
 
@@ -247,8 +237,6 @@ bool ui_post_processing_confirm_transaction(dispatcher_context_t *context, bool
 
 bool ui_post_processing_confirm_message(dispatcher_context_t *context, bool success);
 
-void ui_pre_processing_message(void);
-
 void ui_display_post_processing_confirm_message(bool success);
 void ui_display_post_processing_confirm_transaction(bool success);
 void ui_set_display_prompt(void);
diff --git a/src/ui/display_nbgl.c b/src/ui/display_nbgl.c
index 0578fa6..d20b8bf 100644
--- a/src/ui/display_nbgl.c
+++ b/src/ui/display_nbgl.c
@@ -11,7 +11,6 @@
 
 static const char *confirmed_status;  // text displayed in confirmation page (after long press)
 static const char *rejected_status;   // text displayed in rejection page (after reject confirmed)
-static bool show_message_start_page;
 
 /* Graphical resources (GA) used by the application and NBGL */
 #ifdef SCREEN_SIZE_WALLET
@@ -467,94 +466,33 @@ void ui_display_pubkey_suspicious_flow(void) {
     nbgl_useCaseGenericReview(&genericContent, "Cancel", status_operation_cancel);
 }
 
-static void message_finish_callback(bool confirm) {
-    if (confirm) {
-        nbgl_useCaseReviewStreamingFinish(GA_SIGN_MESSAGE, start_processing_message_callback);
-    } else {
-        status_message_cancel();
-    }
-}
-
-static void message_display_content_continue(bool confirm) {
-    if (confirm) {
-        increase_streaming_index();
-        ux_flow_response_true();
-    } else {
-        status_message_cancel();
-    }
-}
-
-static void message_display_content(bool confirm) {
-    if (confirm) {
-        pairList.pairs = pairs;
-        pairList.nbPairs = 0;
-
-        if (get_streaming_index() == 0) {
-            pairs[0].item = "Path";
-            pairs[0].value = g_ui_state.path_and_message.bip32_path_str;
-            pairList.nbPairs = 1;
-        }
+void ui_sign_message_and_confirm_flow(bool is_hash) {
+    pairs[0].item = "Path";
+    pairs[0].value = g_ui_state.path_and_message.bip32_path_str;
 
+    if (!is_hash) {
 #ifdef SCREEN_SIZE_WALLET
-        pairs[pairList.nbPairs].item = "Message content";
+        pairs[1].item = "Message content";
 #else
-        pairs[pairList.nbPairs].item = "Message";
+        pairs[1].item = "Message";
 #endif
-        pairs[pairList.nbPairs].value = g_ui_state.path_and_message.message;
-
-        pairList.wrapping = true;
-        pairList.nbPairs++;
-
-        nbgl_useCaseReviewStreamingContinue(&pairList, message_display_content_continue);
     } else {
-        status_message_cancel();
-    }
-}
-
-static void message_display_path(bool confirm) {
-    if (confirm) {
-        pairs[0].item = "Path";
-        pairs[0].value = g_ui_state.path_and_message.bip32_path_str;
-
         pairs[1].item = "Message hash";
-        pairs[1].value = g_ui_state.path_and_message.message;
-
-        pairList.nbPairs = 2;
-        pairList.pairs = pairs;
-
-        nbgl_useCaseReviewStreamingContinue(&pairList, message_finish_callback);
-    } else {
-        status_message_cancel();
-    }
-}
-
-void ui_sign_message_content_flow(void) {
-    if (show_message_start_page == true) {
-        show_message_start_page = false;
-        nbgl_useCaseReviewStreamingStart(TYPE_MESSAGE,
-                                         &ICON_APP_ACTION,
-                                         GA_REVIEW_MESSAGE,
-                                         NULL,
-                                         message_display_content);
-    } else {
-        message_display_content(true);
     }
-}
 
-void ui_sign_message_path_hash_and_confirm_flow(void) {
-    nbgl_useCaseReviewStreamingStart(TYPE_MESSAGE,
-                                     &ICON_APP_ACTION,
-                                     GA_REVIEW_MESSAGE,
-                                     NULL,
-                                     message_display_path);
-}
+    pairs[1].value = g_ui_state.path_and_message.message;
 
-void ui_sign_message_confirm_flow(void) {
-    nbgl_useCaseReviewStreamingFinish(GA_SIGN_MESSAGE, start_processing_message_callback);
-}
+    pairList.wrapping = true;
+    pairList.nbPairs = 2;
+    pairList.pairs = pairs;
 
-void ui_set_display_prompt(void) {
-    show_message_start_page = true;
+    nbgl_useCaseReview(TYPE_MESSAGE,
+                       &pairList,
+                       &ICON_APP_ACTION,
+                       GA_REVIEW_MESSAGE,
+                       NULL,
+                       GA_SIGN_MESSAGE,
+                       start_processing_message_callback);
 }
 
 // Address flow