From 026cf5914372ee9db2a6e6337cc4be996c452947 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?S=C3=A9bastien=20Palmer?= <sebastien.palmer@functori.com>
Date: Wed, 3 Apr 2024 10:19:07 +0200
Subject: [PATCH] UI: do not update all screens data

Update only screen potentially changed
---
 src/apdu_setup.c |  2 +-
 src/apdu_sign.c  |  2 +-
 src/ui.h         | 25 ++++++++++++++++++++++--
 src/ui_bagl.c    | 50 ++++++++++++++++++++++++++++++++++++++++--------
 4 files changed, 67 insertions(+), 12 deletions(-)

diff --git a/src/apdu_setup.c b/src/apdu_setup.c
index b36936c8..a2e6dabd 100644
--- a/src/apdu_setup.c
+++ b/src/apdu_setup.c
@@ -94,7 +94,7 @@ int handle_deauthorize(void) {
     UPDATE_NVRAM(ram, { memset(&ram->baking_key, 0, sizeof(ram->baking_key)); });
 #ifdef HAVE_BAGL
     // Ignore calculation errors
-    calculate_baking_idle_screens_data();
+    calculate_idle_screen_authorized_key();
     refresh_screens();
 #endif  // HAVE_BAGL
 
diff --git a/src/apdu_sign.c b/src/apdu_sign.c
index 6fda5262..42f79a4d 100644
--- a/src/apdu_sign.c
+++ b/src/apdu_sign.c
@@ -211,7 +211,7 @@ static int baking_sign_complete(bool const send_hash) {
             result = perform_signature(send_hash);
 #ifdef HAVE_BAGL
             // Ignore calculation errors
-            calculate_baking_idle_screens_data();
+            calculate_idle_screen_hwm();
 #endif
             break;
 
diff --git a/src/ui.h b/src/ui.h
index 1e05f967..7623b094 100644
--- a/src/ui.h
+++ b/src/ui.h
@@ -42,12 +42,33 @@ void exit_app(void);
 
 #ifdef HAVE_BAGL
 
+/**
+ * @brief Calculates the chain id for the idle screens
+ *
+ * @return tz_exc: exception, SW_OK if none
+ */
+tz_exc calculate_idle_screen_chain_id(void);
+
+/**
+ * @brief Calculates the authorized key for the idle screens
+ *
+ * @return tz_exc: exception, SW_OK if none
+ */
+tz_exc calculate_idle_screen_authorized_key(void);
+
+/**
+ * @brief Calculates the HWM for the idle screens
+ *
+ * @return tz_exc: exception, SW_OK if none
+ */
+tz_exc calculate_idle_screen_hwm(void);
+
 /**
  * @brief Calculates baking values for the idle screens
  *
- * @return bool: whether the values have been calculated successfully or not
+ * @return tz_exc: exception, SW_OK if none
  */
-bool calculate_baking_idle_screens_data(void);
+tz_exc calculate_baking_idle_screens_data(void);
 
 /**
  * @brief Prepare confirmation screens callbacks
diff --git a/src/ui_bagl.c b/src/ui_bagl.c
index 5a183836..70a0ac5c 100644
--- a/src/ui_bagl.c
+++ b/src/ui_bagl.c
@@ -78,16 +78,25 @@ UX_FLOW(ux_idle_flow,
         &ux_idle_quit_step,
         FLOW_LOOP);
 
-bool calculate_baking_idle_screens_data(void) {
+tz_exc calculate_idle_screen_chain_id(void) {
     tz_exc exc = SW_OK;
 
-    memset(&home_context, 0, sizeof(home_context));
+    memset(&home_context.chain_id, 0, sizeof(home_context.chain_id));
 
     TZ_ASSERT(chain_id_to_string_with_aliases(home_context.chain_id,
                                               sizeof(home_context.chain_id),
                                               &N_data.main_chain_id) >= 0,
               EXC_WRONG_LENGTH);
 
+end:
+    return exc;
+}
+
+tz_exc calculate_idle_screen_authorized_key(void) {
+    tz_exc exc = SW_OK;
+
+    memset(&home_context.authorized_key, 0, sizeof(home_context.authorized_key));
+
     if (N_data.baking_key.bip32_path.length == 0u) {
         TZ_ASSERT(copy_string(home_context.authorized_key,
                               sizeof(home_context.authorized_key),
@@ -99,25 +108,50 @@ bool calculate_baking_idle_screens_data(void) {
                                                      &N_data.baking_key));
     }
 
+end:
+    return exc;
+}
+
+tz_exc calculate_idle_screen_hwm(void) {
+    tz_exc exc = SW_OK;
+
+    memset(&home_context.hwm, 0, sizeof(home_context.hwm));
+
     TZ_ASSERT(hwm_to_string(home_context.hwm, sizeof(home_context.hwm), &N_data.hwm.main) >= 0,
               EXC_WRONG_LENGTH);
 
-    return true;
+end:
+    return exc;
+}
+
+tz_exc calculate_baking_idle_screens_data(void) {
+    tz_exc exc = SW_OK;
+
+    TZ_CHECK(calculate_idle_screen_chain_id());
+
+    TZ_CHECK(calculate_idle_screen_authorized_key());
+
+    TZ_CHECK(calculate_idle_screen_hwm());
 
 end:
-    TZ_EXC_PRINT(exc);
-    return false;
+    return exc;
 }
 
 void ui_initial_screen(void) {
+    tz_exc exc = SW_OK;
+
     // reserve a display stack slot if none yet
     if (G_ux.stack_count == 0) {
         ux_stack_push();
     }
 
-    if (calculate_baking_idle_screens_data()) {
-        ux_flow_init(0, ux_idle_flow, NULL);
-    }
+    TZ_CHECK(calculate_baking_idle_screens_data());
+
+    ux_flow_init(0, ux_idle_flow, NULL);
+
+    return;
+end:
+    TZ_EXC_PRINT(exc);
 }
 
 void ux_prepare_confirm_callbacks(ui_callback_t ok_c, ui_callback_t cxl_c) {