From c0eaf8d1066605892502698fbdedf9b9114dc92c Mon Sep 17 00:00:00 2001
From: Alexandre Paillier <alexandre.paillier@ledger.fr>
Date: Thu, 12 May 2022 17:30:26 +0200
Subject: [PATCH] Proper cleanup after EIP712 message

---
 src/main.c                                    |  1 -
 src_features/signMessageEIP712/context.c      | 17 ++++++--
 src_features/signMessageEIP712/context.h      |  5 ++-
 src_features/signMessageEIP712/entrypoint.c   |  8 +++-
 src_features/signMessageEIP712/field_hash.c   |  8 +++-
 src_features/signMessageEIP712/field_hash.h   |  1 +
 src_features/signMessageEIP712/path.c         |  6 ++-
 src_features/signMessageEIP712/path.h         |  1 +
 .../signMessageEIP712/sol_typenames.c         |  2 +-
 .../signMessageEIP712/sol_typenames.h         |  2 +-
 src_features/signMessageEIP712/ui_flow_712.c  |  5 +--
 src_features/signMessageEIP712/ui_logic.c     | 41 +++++++++++++++++++
 src_features/signMessageEIP712/ui_logic.h     |  4 ++
 13 files changed, 88 insertions(+), 13 deletions(-)

diff --git a/src/main.c b/src/main.c
index bd8e67b..09f8eb2 100644
--- a/src/main.c
+++ b/src/main.c
@@ -903,7 +903,6 @@ void coin_main(chain_config_t *coin_config) {
     }
     reset_app_context();
     tmpCtx.transactionContext.currentItemIndex = 0;
-    init_eip712_context();
 
     for (;;) {
         UX_INIT();
diff --git a/src_features/signMessageEIP712/context.c b/src_features/signMessageEIP712/context.c
index 3a08d04..dbea2d6 100644
--- a/src_features/signMessageEIP712/context.c
+++ b/src_features/signMessageEIP712/context.c
@@ -11,17 +11,18 @@
 uint8_t  *typenames_array;
 uint8_t  *structs_array;
 uint8_t  *current_struct_fields_array;
+bool eip712_context_initialized = false;
 
 /**
  *
  * @return a boolean indicating if the initialization was successful or not
  */
-bool    init_eip712_context(void)
+bool    eip712_context_init(void)
 {
     // init global variables
     mem_init();
 
-    if (init_sol_typenames() == false)
+    if (sol_typenames_init() == false)
     {
         return false;
     }
@@ -49,7 +50,17 @@ bool    init_eip712_context(void)
 
     // create len(types)
     *structs_array = 0;
+
+    eip712_context_initialized = true;
+
     return true;
 }
 
-// TODO: Make a deinit function
+void    eip712_context_deinit(void)
+{
+    path_deinit();
+    field_hash_deinit();
+    ui_712_deinit();
+    mem_reset();
+    eip712_context_initialized = false;
+}
diff --git a/src_features/signMessageEIP712/context.h b/src_features/signMessageEIP712/context.h
index 7c4520b..82aff85 100644
--- a/src_features/signMessageEIP712/context.h
+++ b/src_features/signMessageEIP712/context.h
@@ -8,6 +8,9 @@ extern uint8_t  *typenames_array;
 extern uint8_t  *structs_array;
 extern uint8_t  *current_struct_fields_array;
 
-bool    init_eip712_context(void);
+bool    eip712_context_init(void);
+void    eip712_context_deinit(void);
+
+extern bool eip712_context_initialized;
 
 #endif // EIP712_CTX_H_
diff --git a/src_features/signMessageEIP712/entrypoint.c b/src_features/signMessageEIP712/entrypoint.c
index 368deb3..5b67a4f 100644
--- a/src_features/signMessageEIP712/entrypoint.c
+++ b/src_features/signMessageEIP712/entrypoint.c
@@ -367,6 +367,13 @@ bool    handle_eip712_struct_def(const uint8_t *const apdu_buf)
 {
     bool ret = true;
 
+    if (!eip712_context_initialized)
+    {
+        if (!eip712_context_init())
+        {
+            return false;
+        }
+    }
     switch (apdu_buf[OFFSET_P2])
     {
         case P2_NAME:
@@ -391,7 +398,6 @@ bool    handle_eip712_struct_def(const uint8_t *const apdu_buf)
         G_io_apdu_buffer[0] = 0x6A;
         G_io_apdu_buffer[1] = 0x80;
     }
-    //*flags |= IO_ASYNCH_REPLY;
     // Send back the response, do not restart the event loop
     io_exchange(CHANNEL_APDU | IO_RETURN_AFTER_TX, 2);
     return ret;
diff --git a/src_features/signMessageEIP712/field_hash.c b/src_features/signMessageEIP712/field_hash.c
index f8940d5..c8c1fb1 100644
--- a/src_features/signMessageEIP712/field_hash.c
+++ b/src_features/signMessageEIP712/field_hash.c
@@ -10,7 +10,8 @@
 
 static s_field_hashing *fh = NULL;
 
-bool field_hash_init(void)
+
+bool    field_hash_init(void)
 {
     if (fh == NULL)
     {
@@ -23,6 +24,11 @@ bool field_hash_init(void)
     return true;
 }
 
+void    field_hash_deinit(void)
+{
+    fh = NULL;
+}
+
 bool    field_hash(const uint8_t *data,
                    uint8_t data_length,
                    bool partial)
diff --git a/src_features/signMessageEIP712/field_hash.h b/src_features/signMessageEIP712/field_hash.h
index ca08aff..6a3fb4a 100644
--- a/src_features/signMessageEIP712/field_hash.h
+++ b/src_features/signMessageEIP712/field_hash.h
@@ -19,6 +19,7 @@ typedef struct
 }   s_field_hashing;
 
 bool field_hash_init(void);
+void field_hash_deinit(void);
 bool field_hash(const uint8_t *data,
                 uint8_t data_length,
                 bool partial);
diff --git a/src_features/signMessageEIP712/path.c b/src_features/signMessageEIP712/path.c
index fb87573..9cf4116 100644
--- a/src_features/signMessageEIP712/path.c
+++ b/src_features/signMessageEIP712/path.c
@@ -147,7 +147,6 @@ static bool path_depth_list_pop(void)
                 memcpy(tmpCtx.messageSigningContext712.messageHash,
                        shash,
                        KECCAK256_HASH_BYTESIZE);
-                mem_reset();
                 break;
             default:
                 break;
@@ -573,3 +572,8 @@ bool    path_init(void)
     }
     return path_struct != NULL;
 }
+
+void    path_deinit(void)
+{
+    path_struct = NULL;
+}
diff --git a/src_features/signMessageEIP712/path.h b/src_features/signMessageEIP712/path.h
index ffec68f..030e369 100644
--- a/src_features/signMessageEIP712/path.h
+++ b/src_features/signMessageEIP712/path.h
@@ -33,6 +33,7 @@ bool    path_set_root(const char *const struct_name, uint8_t length);
 const void  *path_get_field(void);
 bool    path_advance(void);
 bool    path_init(void);
+void    path_deinit(void);
 bool    path_new_array_depth(uint8_t size);
 
 #endif // PATH_H_
diff --git a/src_features/signMessageEIP712/sol_typenames.c b/src_features/signMessageEIP712/sol_typenames.c
index 11a24ae..f4fc562 100644
--- a/src_features/signMessageEIP712/sol_typenames.c
+++ b/src_features/signMessageEIP712/sol_typenames.c
@@ -39,7 +39,7 @@ static bool find_enum_matches(const uint8_t (*enum_to_idx)[TYPES_COUNT - 1][IDX_
     return (enum_match != NULL);
 }
 
-bool    init_sol_typenames(void)
+bool    sol_typenames_init(void)
 {
     const char *const typenames[] = {
         "int",      // 0
diff --git a/src_features/signMessageEIP712/sol_typenames.h b/src_features/signMessageEIP712/sol_typenames.h
index b1e2a67..bb8e259 100644
--- a/src_features/signMessageEIP712/sol_typenames.h
+++ b/src_features/signMessageEIP712/sol_typenames.h
@@ -4,7 +4,7 @@
 #include <stdbool.h>
 #include <stdint.h>
 
-bool init_sol_typenames(void);
+bool sol_typenames_init(void);
 
 const char *get_struct_field_sol_typename(const uint8_t *ptr,
                                           uint8_t *const length);
diff --git a/src_features/signMessageEIP712/ui_flow_712.c b/src_features/signMessageEIP712/ui_flow_712.c
index 7563407..37d97f1 100644
--- a/src_features/signMessageEIP712/ui_flow_712.c
+++ b/src_features/signMessageEIP712/ui_flow_712.c
@@ -1,7 +1,6 @@
 #include "ui_flow_712.h"
 #include "ui_logic.h"
 #include "shared_context.h" // strings
-#include "common_712.h"
 
 // clang-format off
 UX_STEP_NOCB(
@@ -31,7 +30,7 @@ UX_STEP_INIT(
 UX_STEP_CB(
     ux_712_step_approve,
     pb,
-    ui_712_approve_cb(NULL),
+    ui_712_approve(NULL),
     {
       &C_icon_validate_14,
       "Approve",
@@ -39,7 +38,7 @@ UX_STEP_CB(
 UX_STEP_CB(
     ux_712_step_reject,
     pb,
-    ui_712_reject_cb(NULL),
+    ui_712_reject(NULL),
     {
       &C_icon_crossmark,
       "Reject",
diff --git a/src_features/signMessageEIP712/ui_logic.c b/src_features/signMessageEIP712/ui_logic.c
index 23d4ed9..d52eaf5 100644
--- a/src_features/signMessageEIP712/ui_logic.c
+++ b/src_features/signMessageEIP712/ui_logic.c
@@ -10,9 +10,14 @@
 #include "eip712.h" // get_struct_name
 #include "ethUtils.h" // getEthDisplayableAddress
 #include "utils.h" // uint256_to_decimal
+#include "common_712.h"
+#include "context.h" // eip712_context_deinit
+
 
 static t_ui_context *ui_ctx = NULL;
 
+
+
 /**
  * Called on the intermediate dummy screen between the dynamic step
  * && the approve/reject screen
@@ -43,6 +48,8 @@ void    ui_712_next_field(void)
 
 /**
  * Used to notify of a new struct to review (domain or message)
+ *
+ * @param[in] struct_ptr pointer to the structure
  */
 void    ui_712_new_root_struct(const void *const struct_ptr)
 {
@@ -165,3 +172,37 @@ bool    ui_712_init(void)
     }
     return ui_ctx != NULL;
 }
+
+/**
+ * Deinit function that simply unsets the struct pointer to NULL
+ */
+void    ui_712_deinit(void)
+{
+    ui_ctx = NULL;
+}
+
+/**
+ * Approve button handling, calls the common handler function then
+ * deinitializes the EIP712 context altogether.
+ * @param[in] e unused here, just needed to match the UI function signature
+ * @return unused here, just needed to match the UI function signature
+ */
+unsigned int ui_712_approve(const bagl_element_t *e)
+{
+    ui_712_approve_cb(e);
+    eip712_context_deinit();
+    return 0;
+}
+
+/**
+ * Reject button handling, calls the common handler function then
+ * deinitializes the EIP712 context altogether.
+ * @param[in] e unused here, just needed to match the UI function signature
+ * @return unused here, just needed to match the UI function signature
+ */
+unsigned int ui_712_reject(const bagl_element_t *e)
+{
+    ui_712_reject_cb(e);
+    eip712_context_deinit();
+    return 0;
+}
diff --git a/src_features/signMessageEIP712/ui_logic.h b/src_features/signMessageEIP712/ui_logic.h
index 7e6e82e..c313a9d 100644
--- a/src_features/signMessageEIP712/ui_logic.h
+++ b/src_features/signMessageEIP712/ui_logic.h
@@ -2,6 +2,7 @@
 #define UI_LOGIC_712_H_
 
 #include <stdint.h>
+#include "ux.h"
 
 typedef enum
 {
@@ -17,9 +18,12 @@ typedef struct
 }   t_ui_context;
 
 bool ui_712_init(void);
+void ui_712_deinit(void);
 void ui_712_next_field(void);
 void ui_712_new_root_struct(const void *const struct_ptr);
 void ui_712_new_field(const void *const field_ptr, const uint8_t *const data, uint8_t length);
 void ui_712_end_sign(void);
+unsigned int ui_712_approve(const bagl_element_t *e);
+unsigned int ui_712_reject(const bagl_element_t *e);
 
 #endif // UI_LOGIC_712_H_