Merge pull request #385 from matheusmoreira/totp-hot-patch

TOTP hotfix: reduce memory usage

2 files changed, 95 insertions, 32 deletions

diff --git a/movement/watch_faces/complication/totp_face.c b/movement/watch_faces/complication/totp_face.c
index 40b083dd..a593e9c9 100644
--- a/movement/watch_faces/complication/totp_face.c
+++ b/movement/watch_faces/complication/totp_face.c
@@ -39,18 +39,22 @@
 #include "TOTP.h"
 #include "base32.h"
 
+#ifndef TOTP_FACE_MAX_KEY_LENGTH
+#define TOTP_FACE_MAX_KEY_LENGTH 128
+#endif
+
 typedef struct {
     unsigned char labels[2];
     hmac_alg algorithm;
     uint32_t period;
-    size_t key_length;
-    unsigned char *key;
+    size_t encoded_key_length;
+    unsigned char *encoded_key;
 } totp_t;
 
 #define CREDENTIAL(label, key_array, algo, timestep) \
     (const totp_t) { \
-        .key = ((unsigned char *) key_array), \
-        .key_length = sizeof(key_array) - 1, \
+        .encoded_key = ((unsigned char *) key_array), \
+        .encoded_key_length = sizeof(key_array) - 1, \
         .period = (timestep), \
         .labels = (#label), \
         .algorithm = (algo), \
@@ -67,15 +71,63 @@ static totp_t credentials[] = {
 // END OF KEY DATA.
 ////////////////////////////////////////////////////////////////////////////////
 
+static inline totp_t *totp_at(size_t i) {
+    return &credentials[i];
+}
+
 static inline totp_t *totp_current(totp_state_t *totp_state) {
-    return &credentials[totp_state->current_index];
+    return totp_at(totp_state->current_index);
 }
 
 static inline size_t totp_total(void) {
     return sizeof(credentials) / sizeof(*credentials);
 }
 
-static void totp_display(totp_state_t *totp_state) {
+static void totp_validate_key_lengths(void) {
+    for (size_t n = totp_total(), i = 0; i < n; ++i) {
+        totp_t *totp = totp_at(i);
+
+        if (UNBASE32_LEN(totp->encoded_key_length) > TOTP_FACE_MAX_KEY_LENGTH) {
+            // Key exceeds static limits, turn it off by zeroing the length
+            totp->encoded_key_length = 0;
+        }
+    }
+}
+
+static void totp_generate(totp_state_t *totp_state) {
+    totp_t *totp = totp_current(totp_state);
+
+    if (totp->encoded_key_length <= 0) {
+        // Key exceeded static limits and was turned off
+        totp_state->current_decoded_key_length = 0;
+        return;
+    }
+
+    totp_state->current_decoded_key_length = base32_decode(totp->encoded_key, totp_state->current_decoded_key);
+
+    if (totp_state->current_decoded_key_length == 0) {
+        // Decoding failed for some reason
+        // Not a base 32 string?
+        return;
+    }
+
+    TOTP(
+        totp_state->current_decoded_key,
+        totp_state->current_decoded_key_length,
+        totp->period,
+        totp->algorithm
+    );
+}
+
+static void totp_display_error(totp_state_t *totp_state) {
+    char buf[10 + 1];
+    totp_t *totp = totp_current(totp_state);
+
+    snprintf(buf, sizeof(buf), "%c%c  ERROR ", totp->labels[0], totp->labels[1]);
+    watch_display_string(buf, 0);
+}
+
+static void totp_display_code(totp_state_t *totp_state) {
     char buf[14];
     div_t result;
     uint8_t valid_for;
@@ -92,46 +144,55 @@ static void totp_display(totp_state_t *totp_state) {
     watch_display_string(buf, 0);
 }
 
-static void totp_face_decode_secrets(void) {
-    for (size_t n = totp_total(), i = 0; i < n; ++i) {
-        totp_t *totp = &credentials[i];
-        unsigned char *key = totp->key;
+static void totp_display(totp_state_t *totp_state) {
+    if (totp_state->current_decoded_key_length > 0) {
+        totp_display_code(totp_state);
+    } else {
+        totp_display_error(totp_state);
+    }
+}
 
-        totp->key = malloc(UNBASE32_LEN(totp->key_length));
-        totp->key_length = base32_decode(key, totp->key);
+static void totp_generate_and_display(totp_state_t *totp_state) {
+    totp_generate(totp_state);
+    totp_display(totp_state);
+}
 
-        if (totp->key_length == 0) {
-            free(totp->key);
-            continue;
-        }
-    }
+static inline uint32_t totp_compute_base_timestamp(movement_settings_t *settings) {
+    return watch_utility_date_time_to_unix_time(watch_rtc_get_date_time(), movement_timezone_offsets[settings->bit.time_zone] * 60);
 }
 
 void totp_face_setup(movement_settings_t *settings, uint8_t watch_face_index, void ** context_ptr) {
     (void) settings;
     (void) watch_face_index;
 
+    totp_validate_key_lengths();
+
     if (*context_ptr == NULL) {
         totp_state_t *totp = malloc(sizeof(totp_state_t));
-        totp_face_decode_secrets();
+        totp->current_decoded_key = malloc(TOTP_FACE_MAX_KEY_LENGTH);
         *context_ptr = totp;
     }
 }
 
 void totp_face_activate(movement_settings_t *settings, void *context) {
     (void) settings;
-    memset(context, 0, sizeof(totp_state_t));
-    totp_state_t *totp_state = (totp_state_t *)context;
-    totp_t *totp = totp_current(totp_state);
-    TOTP(totp->key, totp->key_length, totp->period, totp->algorithm);
-    totp_state->timestamp = watch_utility_date_time_to_unix_time(watch_rtc_get_date_time(), movement_timezone_offsets[settings->bit.time_zone] * 60);
-    totp_state->current_code = getCodeFromTimestamp(totp_state->timestamp);
+
+    totp_state_t *totp = (totp_state_t *) context;
+
+    totp->timestamp = totp_compute_base_timestamp(settings);
+    totp->steps = 0;
+    totp->current_code = 0;
+    totp->current_index = 0;
+    totp->current_decoded_key_length = 0;
+    // totp->current_decoded_key is already initialized in setup
+
+    totp_generate_and_display(totp);
 }
 
 bool totp_face_loop(movement_event_t event, movement_settings_t *settings, void *context) {
     (void) settings;
-    totp_state_t *totp_state = (totp_state_t *)context;
-    totp_t *totp;
+
+    totp_state_t *totp_state = (totp_state_t *) context;
 
     switch (event.event_type) {
         case EVENT_TICK:
@@ -150,9 +211,9 @@ bool totp_face_loop(movement_event_t event, movement_settings_t *settings, void
                 // wrap around to first key
                 totp_state->current_index = 0;
             }
-            totp = totp_current(totp_state);
-            TOTP(totp->key, totp->key_length, totp->period, totp->algorithm);
-            totp_display(totp_state);
+
+            totp_generate_and_display(totp_state);
+
             break;
         case EVENT_LIGHT_BUTTON_UP:
             if (totp_state->current_index == 0) {
@@ -161,9 +222,9 @@ bool totp_face_loop(movement_event_t event, movement_settings_t *settings, void
             } else {
                 totp_state->current_index--;
             }
-            totp = totp_current(totp_state);
-            TOTP(totp->key, totp->key_length, totp->period, totp->algorithm);
-            totp_display(totp_state);
+
+            totp_generate_and_display(totp_state);
+
             break;
         case EVENT_ALARM_BUTTON_DOWN:
         case EVENT_ALARM_LONG_PRESS:
diff --git a/movement/watch_faces/complication/totp_face.h b/movement/watch_faces/complication/totp_face.h
index af269e92..19a2cd45 100644
--- a/movement/watch_faces/complication/totp_face.h
+++ b/movement/watch_faces/complication/totp_face.h
@@ -70,6 +70,8 @@ typedef struct {
     uint8_t steps;
     uint32_t current_code;
     uint8_t current_index;
+    uint8_t *current_decoded_key;
+    size_t current_decoded_key_length;
 } totp_state_t;
 
 void totp_face_setup(movement_settings_t *settings, uint8_t watch_face_index, void ** context_ptr);