Allow baking app to go into locked/screensaver mode on Nano X

src/apdu_pubkey.c

@@ -1,9 +1,11 @@
 #include "apdu_pubkey.h"
 
 #include "apdu.h"
+#include "baking_auth.h"
 #include "cx.h"
 #include "globals.h"
 #include "keys.h"
+#include "key_macros.h"
 #include "protocol.h"
 #include "to_string.h"
 #include "ui.h"
@@ -68,9 +70,23 @@ static void prompt_address(
 #   endif
 }
 
-size_t handle_apdu_get_public_key(uint8_t instruction) {
-    uint8_t *dataBuffer = G_io_apdu_buffer + OFFSET_CDATA;
+#ifdef BAKING_APP
+size_t handle_apdu_authorize_baking(__attribute__((unused)) uint8_t instruction) {
+    if (READ_UNALIGNED_BIG_ENDIAN(uint8_t, &G_io_apdu_buffer[OFFSET_P1]) != 0) THROW(EXC_WRONG_PARAM);
+
+    uint8_t const curve_code = READ_UNALIGNED_BIG_ENDIAN(uint8_t, &G_io_apdu_buffer[OFFSET_CURVE]);
+    G.key.curve = curve_code_to_curve(curve_code);
+
+    size_t const cdata_size = READ_UNALIGNED_BIG_ENDIAN(uint8_t, &G_io_apdu_buffer[OFFSET_LC]);
+    uint8_t const *const cdata = &G_io_apdu_buffer[OFFSET_CDATA];
+    read_bip32_path(&G.key.bip32_path, cdata, cdata_size);
+    if (G.key.bip32_path.length == 0) THROW(EXC_WRONG_LENGTH_FOR_INS);
+
+    prompt_address(true, baking_ok, delay_reject);
+}
+#endif
 
+size_t handle_apdu_get_public_key(uint8_t instruction) {
     if (READ_UNALIGNED_BIG_ENDIAN(uint8_t, &G_io_apdu_buffer[OFFSET_P1]) != 0) THROW(EXC_WRONG_PARAM);
 
     // do not expose pks without prompt over U2F (browser support)
@@ -80,37 +96,15 @@ size_t handle_apdu_get_public_key(uint8_t instruction) {
     G.key.curve = curve_code_to_curve(curve_code);
 
     size_t const cdata_size = READ_UNALIGNED_BIG_ENDIAN(uint8_t, &G_io_apdu_buffer[OFFSET_LC]);
+    uint8_t const *const cdata = &G_io_apdu_buffer[OFFSET_CDATA];
+    read_bip32_path(&G.key.bip32_path, cdata, cdata_size);
 
-#ifdef BAKING_APP
-    if (cdata_size == 0 && instruction == INS_AUTHORIZE_BAKING) {
-        copy_bip32_path_with_curve(&G.key, &N_data.baking_key);
-    } else {
-#endif
-        read_bip32_path(&G.key.bip32_path, dataBuffer, cdata_size);
-#ifdef BAKING_APP
-        if (G.key.bip32_path.length == 0) THROW(EXC_WRONG_LENGTH_FOR_INS);
-    }
-#endif
-    generate_public_key(&G.public_key, G.key.curve, &G.key.bip32_path);
+    generate_public_key_cached(&G.public_key, &G.key);
 
     if (instruction == INS_GET_PUBLIC_KEY) {
         return provide_pubkey(G_io_apdu_buffer, &G.public_key);
     } else {
-        // instruction == INS_PROMPT_PUBLIC_KEY || instruction == INS_AUTHORIZE_BAKING
-        ui_callback_t cb;
-        bool bake;
-#ifdef BAKING_APP
-        if (instruction == INS_AUTHORIZE_BAKING) {
-            cb = baking_ok;
-            bake = true;
-        } else {
-#endif
-            // INS_PROMPT_PUBLIC_KEY
-            cb = pubkey_ok;
-            bake = false;
-#ifdef BAKING_APP
-        }
-#endif
-        prompt_address(bake, cb, delay_reject);
+        // instruction == INS_PROMPT_PUBLIC_KEY
+        prompt_address(false, pubkey_ok, delay_reject);
     }
 }

src/apdu_pubkey.h

@@ -3,3 +3,7 @@
 #include "apdu.h"
 
 size_t handle_apdu_get_public_key(uint8_t instruction);
+
+#ifdef BAKING_APP
+size_t handle_apdu_authorize_baking(uint8_t instruction);
+#endif

src/apdu_setup.c

@@ -31,10 +31,7 @@ static bool ok(void) {
         ram->hwm.test.highest_level = G.hwm.test;
         ram->hwm.test.had_endorsement = false;
     });
-
-    cx_ecfp_public_key_t const *const pubkey = generate_public_key_return_global(
-        G.key.curve, &G.key.bip32_path);
-    delayed_send(provide_pubkey(G_io_apdu_buffer, pubkey));
+    delayed_send(provide_pubkey(G_io_apdu_buffer, get_cached_baking_pubkey()));
     return true;
 }
 

src/apdu_sign.c

@@ -108,7 +108,7 @@ static bool parse_allowed_operations(
         // TODO: Add still other operations
 #   endif
 
-    return parse_operations(out, in, in_size, key->curve, &key->bip32_path, allowed);
+    return parse_operations(out, in, in_size, key, allowed);
 }
 
 #ifdef BAKING_APP // ----------------------------------------------------------
@@ -603,6 +603,7 @@ static int perform_signature(bool const on_hash, bool const send_hash) {
 
     uint8_t const *const data = on_hash ? G.final_hash : G.message_data;
     size_t const data_length = on_hash ? sizeof(G.final_hash) : G.message_data_length;
+
     tx += WITH_KEY_PAIR(G.key, key_pair, size_t, ({
         sign(&G_io_apdu_buffer[tx], MAX_SIGNATURE_SIZE, G.key.curve, key_pair, data, data_length);
     }));

src/globals.c

@@ -1,6 +1,8 @@
 #include "globals.h"
 
 #include "exception.h"
+#include "keys.h"
+#include "key_macros.h"
 #include "to_string.h"
 
 #ifdef TARGET_NANOX
@@ -70,15 +72,12 @@ void calculate_baking_idle_screens_data(void) {
         number_to_string(global.ui.baking_idle_screens.hwm, (level_t const)N_data.hwm.main.highest_level);
 #   endif
 
-    if (N_data.baking_key.bip32_path.length == 0) {
-        STRCPY(global.ui.baking_idle_screens.pkh, "No Key Authorized");
-    } else {
-        cx_ecfp_public_key_t const *const pubkey = generate_public_key_return_global(
-            (cx_curve_t const)N_data.baking_key.curve,
-            (bip32_path_t const *const)&N_data.baking_key.bip32_path);
-        pubkey_to_pkh_string(
+    if (baking_has_authorized_key()) {
+        bip32_path_with_curve_to_pkh_string(
             global.ui.baking_idle_screens.pkh, sizeof(global.ui.baking_idle_screens.pkh),
-            (cx_curve_t const)N_data.baking_key.curve, pubkey);
+            &N_data.baking_key);
+    } else {
+        STRCPY(global.ui.baking_idle_screens.pkh, "No Key Authorized");
     }
 
 #   ifdef TARGET_NANOX
@@ -101,4 +100,19 @@ void update_baking_idle_screens(void) {
     ui_refresh();
 }
 
+void update_baking_keys_cache(void) {
+    if (baking_has_authorized_key()) {
+        generate_public_key(
+            &global.baking_cache.root_public_key,
+            kung_fu_animal_bip32_path_with_curve.curve,
+            &kung_fu_animal_bip32_path_with_curve.bip32_path);
+        generate_key_pair(
+            &global.baking_cache.keys,
+            N_data.baking_key.curve,
+            (const bip32_path_t *)&N_data.baking_key.bip32_path);
+    } else {
+        memset(&global.baking_cache, 0, sizeof(global.baking_cache));
+    }
+}
+
 #endif // #ifdef BAKING_APP

src/globals.h

@@ -62,6 +62,16 @@ typedef struct {
 } apdu_sign_state_t;
 
 typedef struct {
+# ifdef BAKING_APP
+  struct {
+    // cached so we can bake while the device is locked.
+    key_pair_t keys;
+
+    // tezos-client queries the root key to generate kung-fu animal names.
+    cx_ecfp_public_key_t root_public_key;
+  } baking_cache;
+# endif
+
   void *stack_root;
   apdu_handler handlers[INS_MAX + 1];
 
@@ -179,14 +189,37 @@ void calculate_baking_idle_screens_data(void);
 void update_baking_idle_screens(void);
 high_watermark_t volatile *select_hwm_by_chain(chain_id_t const chain_id, nvram_data volatile *const ram);
 
+void update_baking_keys_cache(void);
+
+static inline bool baking_has_authorized_key() {
+    return N_data.baking_key.curve != CX_CURVE_NONE && N_data.baking_key.bip32_path.length > 0;
+}
+
+// Returns NULL if the cache is not valid.
+static inline key_pair_t const *get_cached_baking_key_pair() {
+    return baking_has_authorized_key()
+      ? &global.baking_cache.keys
+      : NULL;
+}
+
+// Returns NULL if the cache is not valid.
+static inline cx_ecfp_public_key_t const *get_cached_baking_pubkey() {
+    key_pair_t const *const key_pair = get_cached_baking_key_pair();
+    return key_pair == NULL
+      ? NULL
+      : &key_pair->public_key;
+}
+
 // Properly updates NVRAM data to prevent any clobbering of data.
 // 'out_param' defines the name of a pointer to the nvram_data struct
 // that 'body' can change to apply updates.
 #define UPDATE_NVRAM(out_name, body) ({ \
     nvram_data *const out_name = &global.baking_auth.new_data; \
     memcpy(&global.baking_auth.new_data, (nvram_data const *const)&N_data, sizeof(global.baking_auth.new_data)); \
     body; \
-    nvm_write((void*)&N_data, &global.baking_auth.new_data, sizeof(N_data)); \
+    bool const baking_keys_changed = memcmp((void const *const)&N_data.baking_key, &global.baking_auth.new_data.baking_key, sizeof(N_data.baking_key)) != 0; \
+    nvm_write((void *const)&N_data, &global.baking_auth.new_data, sizeof(N_data)); \
+    if (baking_keys_changed) { /* requires secure chip so update only when necessary */ update_baking_keys_cache(); } \
     update_baking_idle_screens(); \
 })
 #endif

src/key_macros.h

@@ -3,12 +3,12 @@
 #include "globals.h"
 #include "keys.h"
 
-// Yes you need this oddness if you want to use __LINE__
+// Yes you need this if you want to use __LINE__
 #define CONCAT_(a, b) a##b
 #define CONCAT(a, b) CONCAT_(a, b)
 #define MACROVAR(a, b) CONCAT(____ ## _ ## a ## _ ## b, __LINE__)
 
-#define WITH_KEY_PAIR(bip32_path_with_curve, vname, type, body) ({ \
+#define WITH_KEY_PAIR_NO_CACHE(bip32_path_with_curve, vname, type, body) ({ \
     bip32_path_with_curve_t volatile const *const MACROVAR(vname, key) = &(bip32_path_with_curve); \
     key_pair_t *const MACROVAR(vname, generated_pair) = \
         generate_key_pair_return_global( \
@@ -32,3 +32,44 @@
     END_TRY; \
     MACROVAR(vname, retval); \
 })
+
+// Runs `body` in a scope where `vname` is a `key_pair_t` generated from `bip32_path_with_curve`.
+// The key data will be zero'ed safely when `body` finishes or if an exception occurs.
+// Uses cached key data when available.
+#ifdef BAKING_APP
+#define WITH_KEY_PAIR(bip32_path_with_curve, vname, type, body) ({ \
+    bip32_path_with_curve_t volatile const *const MACROVAR(vname, key) = &(bip32_path_with_curve); \
+    key_pair_t const *const MACROVAR(vname, cached_pair) = get_cached_baking_key_pair(); \
+    volatile type MACROVAR(vname, retval); \
+    if (MACROVAR(vname, cached_pair) != NULL && bip32_path_with_curve_eq(MACROVAR(vname, key), &N_data.baking_key)) { \
+        key_pair_t const *const vname = MACROVAR(vname, cached_pair); \
+        MACROVAR(vname, retval) = body; \
+    } else { \
+        MACROVAR(vname, retval) = WITH_KEY_PAIR_NO_CACHE(*MACROVAR(vname, key), vname, type, body); \
+    } \
+    MACROVAR(vname, retval); \
+})
+#else
+#define WITH_KEY_PAIR(a, b, c, d) WITH_KEY_PAIR_NO_CACHE(a, b, c, d)
+#endif
+
+
+static inline void generate_public_key_cached(
+    cx_ecfp_public_key_t *const out,
+    bip32_path_with_curve_t volatile const *const key
+) {
+    check_null(out);
+    check_null(key);
+#ifdef BAKING_APP
+    cx_ecfp_public_key_t const *const baking_pubkey = get_cached_baking_pubkey();
+    if (baking_pubkey != NULL && bip32_path_with_curve_eq(key, &kung_fu_animal_bip32_path_with_curve)) {
+        memcpy(out, &global.baking_cache.root_public_key, sizeof(*out));
+    } else if (baking_pubkey != NULL && bip32_path_with_curve_eq(key, &N_data.baking_key)) {
+        memcpy(out, baking_pubkey, sizeof(*out));
+    } else {
+#endif
+        generate_public_key(out, key->curve, (bip32_path_t const *const)&key->bip32_path);
+#ifdef BAKING_APP
+    }
+#endif
+}

src/main.c

@@ -20,7 +20,7 @@ void app_main(void) {
     global.handlers[APDU_INS(INS_GIT)] = handle_apdu_git;
     global.handlers[APDU_INS(INS_SIGN_WITH_HASH)] = handle_apdu_sign_with_hash;
 #ifdef BAKING_APP
-    global.handlers[APDU_INS(INS_AUTHORIZE_BAKING)] = handle_apdu_get_public_key;
+    global.handlers[APDU_INS(INS_AUTHORIZE_BAKING)] = handle_apdu_authorize_baking;
     global.handlers[APDU_INS(INS_RESET)] = handle_apdu_reset;
     global.handlers[APDU_INS(INS_QUERY_AUTH_KEY)] = handle_apdu_query_auth_key;
     global.handlers[APDU_INS(INS_QUERY_MAIN_HWM)] = handle_apdu_main_hwm;
@@ -32,5 +32,10 @@ void app_main(void) {
 #else
     global.handlers[APDU_INS(INS_SIGN_UNSAFE)] = handle_apdu_sign;
 #endif
+
+#   ifdef BAKING_APP
+    update_baking_keys_cache();
+#   endif
+
     main_loop(global.handlers, sizeof(global.handlers));
 }

src/operations.c

@@ -2,6 +2,7 @@
 
 #include "apdu.h"
 #include "globals.h"
+#include "key_macros.h"
 #include "memory.h"
 #include "to_string.h"
 #include "ui.h"
@@ -108,18 +109,20 @@ static inline uint64_t parse_z(const void *data, size_t *ix, size_t length, uint
 static inline void compute_pkh(
     cx_ecfp_public_key_t *const compressed_pubkey_out,
     struct parsed_contract *const contract_out,
-    cx_curve_t const curve,
-    bip32_path_t const *const bip32_path
+    bip32_path_with_curve_t const *const key
 ) {
-    check_null(bip32_path);
     check_null(compressed_pubkey_out);
     check_null(contract_out);
-    cx_ecfp_public_key_t const *const pubkey = generate_public_key_return_global(curve, bip32_path);
+    check_null(key);
+
+    cx_ecfp_public_key_t pubkey;
+    generate_public_key_cached(&pubkey, key);
     public_key_hash(
         contract_out->hash, sizeof(contract_out->hash),
         compressed_pubkey_out,
-        curve, pubkey);
-    contract_out->curve_code = curve_to_curve_code(curve);
+        key->curve, &pubkey);
+
+    contract_out->curve_code = curve_to_curve_code(key->curve);
     if (contract_out->curve_code == TEZOS_NO_CURVE) THROW(EXC_MEMORY_ERROR);
     contract_out->originated = 0;
 }
@@ -148,18 +151,17 @@ static void parse_operations_throws_parse_error(
     struct parsed_operation_group *const out,
     void const *const data,
     size_t length,
-    cx_curve_t curve,
-    bip32_path_t const *const bip32_path,
+    bip32_path_with_curve_t const *const key,
     allowed_operation_set ops
 ) {
     check_null(out);
     check_null(data);
-    check_null(bip32_path);
+    check_null(key);
     memset(out, 0, sizeof(*out));
 
     out->operation.tag = OPERATION_TAG_NONE;
 
-    compute_pkh(&out->public_key, &out->signing, curve, bip32_path);
+    compute_pkh(&out->public_key, &out->signing, key);
 
     size_t ix = 0;
 
@@ -328,13 +330,12 @@ bool parse_operations(
     struct parsed_operation_group *const out,
     uint8_t const *const data,
     size_t length,
-    cx_curve_t curve,
-    bip32_path_t const *const bip32_path,
+    bip32_path_with_curve_t const *const key,
     allowed_operation_set ops
 ) {
     BEGIN_TRY {
         TRY {
-            parse_operations_throws_parse_error(out, data, length, curve, bip32_path, ops);
+            parse_operations_throws_parse_error(out, data, length, key, ops);
         }
         CATCH(EXC_PARSE_ERROR) {
             return false;

src/operations.h

@@ -30,7 +30,6 @@ bool parse_operations(
     struct parsed_operation_group *const out,
     uint8_t const *const data,
     size_t length,
-    cx_curve_t curve,
-    bip32_path_t const *const bip32_path,
+    bip32_path_with_curve_t const *const key,
     allowed_operation_set ops
 );