crypto: refactor bip39 api

core/embed/extmod/modtrezorcrypto/modtrezorcrypto-bip39.h

@@ -24,54 +24,41 @@
 
 /// package: trezorcrypto.bip39
 
-/// def find_word(prefix: str) -> Optional[str]:
+/// def complete_word(prefix: str) -> Optional[str]:
 ///     """
 ///     Return the first word from the wordlist starting with prefix.
 ///     """
-STATIC mp_obj_t mod_trezorcrypto_bip39_find_word(mp_obj_t prefix) {
-  mp_buffer_info_t pfx;
-  mp_get_buffer_raise(prefix, &pfx, MP_BUFFER_READ);
-  if (pfx.len == 0) {
-    return mp_const_none;
-  }
-  for (const char *const *w = mnemonic_wordlist(); *w != 0; w++) {
-    if (strncmp(*w, pfx.buf, pfx.len) == 0) {
-      return mp_obj_new_str(*w, strlen(*w));
-    }
-  }
-  return mp_const_none;
-}
-STATIC MP_DEFINE_CONST_FUN_OBJ_1(mod_trezorcrypto_bip39_find_word_obj,
-                                 mod_trezorcrypto_bip39_find_word);
-
-/// def complete_word(prefix: str) -> int:
-///     """
-///     Return possible 1-letter suffixes for given word prefix.
-///     Result is a bitmask, with 'a' on the lowest bit, 'b' on the second
-///     lowest, etc.
-///     """
 STATIC mp_obj_t mod_trezorcrypto_bip39_complete_word(mp_obj_t prefix) {
   mp_buffer_info_t pfx;
   mp_get_buffer_raise(prefix, &pfx, MP_BUFFER_READ);
   if (pfx.len == 0) {
-    return mp_obj_new_int(0xFFFFFFFF);  // all letters
+    return mp_const_none;
   }
-  uint32_t res = 0;
-  uint8_t bit;
-  const char *word;
-  const char *const *wlist;
-  for (wlist = mnemonic_wordlist(); *wlist != 0; wlist++) {
-    word = *wlist;
-    if (strncmp(word, pfx.buf, pfx.len) == 0 && strlen(word) > pfx.len) {
-      bit = word[pfx.len] - 'a';
-      res |= 1 << bit;
-    }
+  const char *word = mnemonic_complete_word(pfx.buf, pfx.len);
+  if (word) {
+    return mp_obj_new_str(word, strlen(word));
+  } else {
+    return mp_const_none;
   }
-  return mp_obj_new_int(res);
 }
 STATIC MP_DEFINE_CONST_FUN_OBJ_1(mod_trezorcrypto_bip39_complete_word_obj,
                                  mod_trezorcrypto_bip39_complete_word);
 
+/// def word_completion_mask(prefix: str) -> int:
+///     """
+///     Return possible 1-letter suffixes for given word prefix.
+///     Result is a bitmask, with 'a' on the lowest bit, 'b' on the second
+///     lowest, etc.
+///     """
+STATIC mp_obj_t mod_trezorcrypto_bip39_word_completion_mask(mp_obj_t prefix) {
+  mp_buffer_info_t pfx;
+  mp_get_buffer_raise(prefix, &pfx, MP_BUFFER_READ);
+  return mp_obj_new_int(mnemonic_word_completion_mask(pfx.buf, pfx.len));
+}
+STATIC MP_DEFINE_CONST_FUN_OBJ_1(
+    mod_trezorcrypto_bip39_word_completion_mask_obj,
+    mod_trezorcrypto_bip39_word_completion_mask);
+
 /// def generate(strength: int) -> str:
 ///     """
 ///     Generate a mnemonic of given strength (128, 160, 192, 224 and 256 bits).
@@ -167,10 +154,10 @@ STATIC MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(mod_trezorcrypto_bip39_seed_obj, 2,
 
 STATIC const mp_rom_map_elem_t mod_trezorcrypto_bip39_globals_table[] = {
     {MP_ROM_QSTR(MP_QSTR___name__), MP_ROM_QSTR(MP_QSTR_bip39)},
-    {MP_ROM_QSTR(MP_QSTR_find_word),
-     MP_ROM_PTR(&mod_trezorcrypto_bip39_find_word_obj)},
     {MP_ROM_QSTR(MP_QSTR_complete_word),
      MP_ROM_PTR(&mod_trezorcrypto_bip39_complete_word_obj)},
+    {MP_ROM_QSTR(MP_QSTR_word_completion_mask),
+     MP_ROM_PTR(&mod_trezorcrypto_bip39_word_completion_mask_obj)},
     {MP_ROM_QSTR(MP_QSTR_generate),
      MP_ROM_PTR(&mod_trezorcrypto_bip39_generate_obj)},
     {MP_ROM_QSTR(MP_QSTR_from_data),

core/mocks/generated/trezorcrypto/bip39.pyi

@@ -2,14 +2,14 @@ from typing import *
 
 
 # extmod/modtrezorcrypto/modtrezorcrypto-bip39.h
-def find_word(prefix: str) -> Optional[str]:
+def complete_word(prefix: str) -> Optional[str]:
     """
     Return the first word from the wordlist starting with prefix.
     """
 
 
 # extmod/modtrezorcrypto/modtrezorcrypto-bip39.h
-def complete_word(prefix: str) -> int:
+def word_completion_mask(prefix: str) -> int:
     """
     Return possible 1-letter suffixes for given word prefix.
     Result is a bitmask, with 'a' on the lowest bit, 'b' on the second

core/src/apps/management/recovery_device/keyboard_bip39.py

@@ -172,8 +172,8 @@ class Bip39Keyboard(ui.Layout):
 
         # find the completions
         pending = button is not None
-        word = bip39.find_word(text) or ""
-        mask = bip39.complete_word(text)
+        word = bip39.complete_word(text) or ""
+        mask = bip39.word_completion_mask(text)
 
         # modify the input state
         self.input.edit(text, word, pending)

crypto/bip39.c

@@ -228,4 +228,54 @@ void mnemonic_to_seed(const char *mnemonic, const char *passphrase,
 #endif
 }
 
-const char *const *mnemonic_wordlist(void) { return wordlist; }
+// binary search for finding the word in the wordlist
+int mnemonic_find_word(const char *word) {
+  int lo = 0, hi = BIP39_WORDS - 1;
+  while (lo <= hi) {
+    int mid = lo + (hi - lo) / 2;
+    int cmp = strcmp(word, wordlist[mid]);
+    if (cmp == 0) {
+      return mid;
+    }
+    if (cmp > 0) {
+      lo = mid + 1;
+    } else {
+      hi = mid - 1;
+    }
+  }
+  return -1;
+}
+
+const char *mnemonic_complete_word(const char *prefix, int len) {
+  // we need to perform linear search,
+  // because we want to return the first match
+  for (const char *const *w = wordlist; *w != 0; w++) {
+    if (strncmp(*w, prefix, len) == 0) {
+      return *w;
+    }
+  }
+  return NULL;
+}
+
+const char *mnemonic_get_word(int index) {
+  if (index >= 0 && index < BIP39_WORDS) {
+    return wordlist[index];
+  } else {
+    return NULL;
+  }
+}
+
+uint32_t mnemonic_word_completion_mask(const char *prefix, int len) {
+  if (len <= 0) {
+    return 0x3ffffff;  // all letters (bits 1-26 set)
+  }
+  uint32_t res = 0;
+  for (const char *const *w = wordlist; *w != 0; w++) {
+    const char *word = *w;
+    if (strncmp(word, prefix, len) == 0 && word[len] >= 'a' &&
+        word[len] <= 'z') {
+      res |= 1 << (word[len] - 'a');
+    }
+  }
+  return res;
+}

crypto/bip39.h

@@ -24,8 +24,10 @@
 #ifndef __BIP39_H__
 #define __BIP39_H__
 
+#include <stdbool.h>
 #include <stdint.h>
 
+#define BIP39_WORDS 2048
 #define BIP39_PBKDF2_ROUNDS 2048
 
 const char *mnemonic_generate(int strength);  // strength in bits
@@ -42,6 +44,9 @@ void mnemonic_to_seed(const char *mnemonic, const char *passphrase,
                       void (*progress_callback)(uint32_t current,
                                                 uint32_t total));
 
-const char *const *mnemonic_wordlist(void);
+int mnemonic_find_word(const char *word);
+const char *mnemonic_complete_word(const char *prefix, int len);
+const char *mnemonic_get_word(int index);
+uint32_t mnemonic_word_completion_mask(const char *prefix, int len);
 
 #endif

crypto/tests/test_check.c

@@ -5097,6 +5097,17 @@ START_TEST(test_mnemonic_to_entropy) {
 }
 END_TEST
 
+START_TEST(test_mnemonic_find_word) {
+  ck_assert_int_eq(-1, mnemonic_find_word("aaaa"));
+  ck_assert_int_eq(-1, mnemonic_find_word("zzzz"));
+  for (int i = 0; i < BIP39_WORDS; i++) {
+    const char *word = mnemonic_get_word(i);
+    int index = mnemonic_find_word(word);
+    ck_assert_int_eq(i, index);
+  }
+}
+END_TEST
+
 START_TEST(test_slip39_get_word) {
   static const struct {
     const int index;
@@ -8765,6 +8776,7 @@ Suite *test_suite(void) {
   tcase_add_test(tc, test_mnemonic);
   tcase_add_test(tc, test_mnemonic_check);
   tcase_add_test(tc, test_mnemonic_to_entropy);
+  tcase_add_test(tc, test_mnemonic_find_word);
   suite_add_tcase(s, tc);
 
   tc = tcase_create("slip39");

legacy/firmware/recovery.c

@@ -323,7 +323,6 @@ static void display_choices(bool twoColumn, char choices[9][12], int num) {
  * Generates a new matrix and requests the next pin.
  */
 static void next_matrix(void) {
-  const char *const *wl = mnemonic_wordlist();
   char word_choices[9][12];
   uint32_t idx, num;
   bool last = (word_index % 4) == 3;
@@ -342,7 +341,8 @@ static void next_matrix(void) {
       const uint32_t first = TABLE2(idx);
       num = TABLE2(idx + 1) - first;
       for (uint32_t i = 0; i < num; i++) {
-        strlcpy(word_choices[i], wl[first + i], sizeof(word_choices[i]));
+        strlcpy(word_choices[i], mnemonic_get_word(first + i),
+                sizeof(word_choices[i]));
       }
       break;
 
@@ -354,7 +354,8 @@ static void next_matrix(void) {
       num = TABLE1(word_pincode + 1) - idx;
       for (uint32_t i = 0; i < num; i++) {
         add_choice(word_choices[i], (word_table2[idx + i] >> 12),
-                   wl[TABLE2(idx + i)], wl[TABLE2(idx + i + 1) - 1]);
+                   mnemonic_get_word(TABLE2(idx + i)),
+                   mnemonic_get_word(TABLE2(idx + i + 1) - 1));
       }
       break;
 
@@ -366,8 +367,8 @@ static void next_matrix(void) {
       num = 9;
       for (uint32_t i = 0; i < num; i++) {
         add_choice(word_choices[i], (word_table1[idx + i] >> 12),
-                   wl[TABLE2(TABLE1(idx + i))],
-                   wl[TABLE2(TABLE1(idx + i + 1)) - 1]);
+                   mnemonic_get_word(TABLE2(TABLE1(idx + i))),
+                   mnemonic_get_word(TABLE2(TABLE1(idx + i + 1)) - 1));
       }
       break;
 
@@ -376,8 +377,8 @@ static void next_matrix(void) {
       /* num: the number of choices. */
       num = 9;
       for (uint32_t i = 0; i < num; i++) {
-        add_choice(word_choices[i], 1, wl[TABLE2(TABLE1(9 * i))],
-                   wl[TABLE2(TABLE1(9 * (i + 1))) - 1]);
+        add_choice(word_choices[i], 1, mnemonic_get_word(TABLE2(TABLE1(9 * i))),
+                   mnemonic_get_word(TABLE2(TABLE1(9 * (i + 1))) - 1));
       }
       break;
   }
@@ -427,7 +428,7 @@ static void recovery_digit(const char digit) {
     uint32_t widx = word_index / 4;
 
     word_pincode = 0;
-    strlcpy(words[widx], mnemonic_wordlist()[idx], sizeof(words[widx]));
+    strlcpy(words[widx], mnemonic_get_word(idx), sizeof(words[widx]));
     if (widx + 1 == word_count) {
       recovery_done();
       return;
@@ -449,8 +450,8 @@ void next_word(void) {
   oledDrawStringCenter(OLED_WIDTH / 2, 8, _("Please enter"), FONT_STANDARD);
   word_pos = word_order[word_index];
   if (word_pos == 0) {
-    const char *const *wl = mnemonic_wordlist();
-    strlcpy(fake_word, wl[random_uniform(2048)], sizeof(fake_word));
+    strlcpy(fake_word, mnemonic_get_word(random_uniform(BIP39_WORDS)),
+            sizeof(fake_word));
     oledDrawStringCenter(OLED_WIDTH / 2, 24, fake_word,
                          FONT_FIXED | FONT_DOUBLE);
   } else {
@@ -521,6 +522,10 @@ void recovery_init(uint32_t _word_count, bool passphrase_protection,
 }
 
 static void recovery_scrambledword(const char *word) {
+  int index = -1;
+  if (enforce_wordlist) {  // check if word is valid
+    index = mnemonic_find_word(word);
+  }
   if (word_pos == 0) {  // fake word
     if (strcmp(word, fake_word) != 0) {
       if (!dry_run) {
@@ -531,18 +536,9 @@ static void recovery_scrambledword(const char *word) {
       layoutHome();
       return;
     }
-  } else {                   // real word
-    if (enforce_wordlist) {  // check if word is valid
-      const char *const *wl = mnemonic_wordlist();
-      bool found = false;
-      while (*wl) {
-        if (strcmp(word, *wl) == 0) {
-          found = true;
-          break;
-        }
-        wl++;
-      }
-      if (!found) {
+  } else {  // real word
+    if (enforce_wordlist) {
+      if (index < 0) {  // not found
         if (!dry_run) {
           session_clear(true);
         }