From 3bf3e4c38bf98cf264e3014fd31a4a77a41bd5ec Mon Sep 17 00:00:00 2001 From: grdddj Date: Thu, 4 May 2023 16:10:37 +0200 Subject: [PATCH] feat(core/rust): introduce Wordlist component for both bip39 and slip39 --- core/embed/rust/src/trezorhal/bip39.rs | 184 +--------------------- core/embed/rust/src/trezorhal/mod.rs | 1 + core/embed/rust/src/trezorhal/wordlist.rs | 161 +++++++++++++++++++ 3 files changed, 165 insertions(+), 181 deletions(-) create mode 100644 core/embed/rust/src/trezorhal/wordlist.rs diff --git a/core/embed/rust/src/trezorhal/bip39.rs b/core/embed/rust/src/trezorhal/bip39.rs index 1dd0b534a..035ec4e95 100644 --- a/core/embed/rust/src/trezorhal/bip39.rs +++ b/core/embed/rust/src/trezorhal/bip39.rs @@ -1,40 +1,10 @@ -use super::ffi; -use core::cmp::Ordering; -use cstr_core::CStr; - -unsafe fn from_utf8_unchecked<'a>(word: *const cty::c_char) -> &'a str { - // SAFETY: caller must pass a valid 0-terminated UTF-8 string. - // This assumption holds for usage on words of the BIP-39 wordlist. - unsafe { - let word = CStr::from_ptr(word); - core::str::from_utf8_unchecked(word.to_bytes()) - } -} - -/// Compare word from wordlist to a prefix. -/// -/// The comparison returns Less if the word comes lexicographically before all -/// possible words starting with `prefix`, and Greater if it comes after. -/// Equal is returned if the word starts with `prefix`. -unsafe fn prefix_cmp(prefix: &str, word: *const cty::c_char) -> Ordering { - // SAFETY: we assume `word` is a pointer to a 0-terminated string. - for (i, prefix_char) in prefix.as_bytes().iter().enumerate() { - let word_char = unsafe { *(word.add(i)) } as u8; - if word_char == 0 { - // Prefix is longer than word. - return Ordering::Less; - } else if *prefix_char != word_char { - return word_char.cmp(prefix_char); - } - } - Ordering::Equal -} +use super::{ffi, wordlist::Wordlist}; pub fn complete_word(prefix: &str) -> Option<&'static str> { if prefix.is_empty() { None } else { - Wordlist::all().filter_prefix(prefix).iter().next() + Wordlist::bip39().filter_prefix(prefix).iter().next() } } @@ -42,7 +12,7 @@ pub fn options_num(prefix: &str) -> Option { if prefix.is_empty() { None } else { - Some(Wordlist::all().filter_prefix(prefix).len()) + Some(Wordlist::bip39().filter_prefix(prefix).len()) } } @@ -51,151 +21,3 @@ pub fn word_completion_mask(prefix: &str) -> u32 { // passed byte string, making the call safe. unsafe { ffi::mnemonic_word_completion_mask(prefix.as_ptr() as _, prefix.len() as _) } } - -pub struct Wordlist(&'static [*const cty::c_char]); - -impl Wordlist { - pub fn all() -> Self { - Self(unsafe { &ffi::BIP39_WORDLIST_ENGLISH }) - } - - pub const fn empty() -> Self { - Self(&[]) - } - - pub fn filter_prefix(&self, prefix: &str) -> Self { - let mut start = 0usize; - let mut end = self.0.len(); - for (i, word) in self.0.iter().enumerate() { - // SAFETY: We assume our slice is an array of 0-terminated strings. - match unsafe { prefix_cmp(prefix, *word) } { - Ordering::Less => { - start = i + 1; - } - Ordering::Greater => { - end = i; - break; - } - _ => {} - } - } - Self(&self.0[start..end]) - } - - pub fn get(&self, index: usize) -> Option<&'static str> { - // SAFETY: we assume every word in the wordlist is a valid 0-terminated UTF-8 - // string. - self.0 - .get(index) - .map(|word| unsafe { from_utf8_unchecked(*word) }) - } - - pub const fn len(&self) -> usize { - self.0.len() - } - - pub fn iter(&self) -> impl Iterator { - self.0 - .iter() - .map(|word| unsafe { from_utf8_unchecked(*word) }) - } -} - -#[cfg(test)] -mod tests { - use super::*; - use cstr_core::cstr; - - const BIP39_WORD_COUNT: usize = ffi::BIP39_WORD_COUNT as usize; - - #[test] - fn test_prefix_cmp() { - assert_eq!( - unsafe { prefix_cmp("", cstr!("").as_ptr()) }, - Ordering::Equal - ); - - assert_eq!( - unsafe { prefix_cmp("b", cstr!("").as_ptr()) }, - Ordering::Less - ); - assert_eq!( - unsafe { prefix_cmp("b", cstr!("a").as_ptr()) }, - Ordering::Less - ); - assert_eq!( - unsafe { prefix_cmp("b", cstr!("b").as_ptr()) }, - Ordering::Equal - ); - assert_eq!( - unsafe { prefix_cmp("b", cstr!("below").as_ptr()) }, - Ordering::Equal - ); - assert_eq!( - unsafe { prefix_cmp("b", cstr!("c").as_ptr()) }, - Ordering::Greater - ); - - assert_eq!( - unsafe { prefix_cmp("bartender", cstr!("bar").as_ptr()) }, - Ordering::Less - ); - } - - #[test] - fn test_filter_prefix_empty() { - let words = Wordlist::all().filter_prefix(""); - assert_eq!(words.len(), BIP39_WORD_COUNT); - let iter = words.iter(); - assert_eq!(iter.size_hint(), (BIP39_WORD_COUNT, Some(BIP39_WORD_COUNT))); - } - - #[test] - fn test_filter_prefix() { - let expected_result = vec!["strategy", "street", "strike", "strong", "struggle"]; - let result = Wordlist::all() - .filter_prefix("str") - .iter() - .collect::>(); - assert_eq!(result, expected_result); - } - - #[test] - fn test_filter_prefix_refine() { - let expected_result = vec!["strategy", "street", "strike", "strong", "struggle"]; - let words = Wordlist::all().filter_prefix("st"); - let result_a = words.filter_prefix("str").iter().collect::>(); - let result_b = Wordlist::all() - .filter_prefix("str") - .iter() - .collect::>(); - assert_eq!(result_a, expected_result); - assert_eq!(result_b, expected_result); - - let empty = words.filter_prefix("c"); - assert_eq!(empty.len(), 0); - } - - #[test] - fn test_wordlist_get() { - let words = Wordlist::all(); - assert_eq!(words.get(0), Some("abandon")); - assert_eq!(words.get(BIP39_WORD_COUNT - 1), Some("zoo")); - assert_eq!(words.get(BIP39_WORD_COUNT), None); - assert_eq!(words.get(BIP39_WORD_COUNT + 1), None); - - let filtered = words.filter_prefix("str"); - assert_eq!(filtered.get(0), Some("strategy")); - assert_eq!(filtered.get(filtered.len()), None); - } - - #[test] - fn test_filter_prefix_just_one() { - let expected_result = vec!["stick"]; - let result = Wordlist::all() - .filter_prefix("stick") - .iter() - .collect::>(); - assert_eq!(result, expected_result); - } -} diff --git a/core/embed/rust/src/trezorhal/mod.rs b/core/embed/rust/src/trezorhal/mod.rs index a3730ed73..d533ee068 100644 --- a/core/embed/rust/src/trezorhal/mod.rs +++ b/core/embed/rust/src/trezorhal/mod.rs @@ -15,6 +15,7 @@ pub mod slip39; pub mod storage; pub mod usb; pub mod uzlib; +pub mod wordlist; pub mod buffers; #[cfg(not(feature = "micropython"))] diff --git a/core/embed/rust/src/trezorhal/wordlist.rs b/core/embed/rust/src/trezorhal/wordlist.rs new file mode 100644 index 000000000..44d91aea6 --- /dev/null +++ b/core/embed/rust/src/trezorhal/wordlist.rs @@ -0,0 +1,161 @@ +use super::ffi; +use core::cmp::Ordering; +use cstr_core::CStr; + +/// Holds all the possible words with the possibility to interact +/// with the "list" - filtering it further, getting their count, etc. +pub struct Wordlist(&'static [*const cty::c_char]); + +impl Wordlist { + /// Initialize BIP39 wordlist. + pub fn bip39() -> Self { + Self(unsafe { &ffi::BIP39_WORDLIST_ENGLISH }) + } + + /// Initialize SLIP39 wordlist. + pub fn slip39() -> Self { + Self(unsafe { &ffi::SLIP39_WORDLIST }) + } + + /// Only leaves words that have a specified prefix. Throw away others. + pub fn filter_prefix(&self, prefix: &str) -> Self { + let mut start = 0usize; + let mut end = self.0.len(); + for (i, word) in self.0.iter().enumerate() { + // SAFETY: We assume our slice is an array of 0-terminated strings. + match unsafe { prefix_cmp(prefix, *word) } { + Ordering::Less => { + start = i + 1; + } + Ordering::Greater => { + end = i; + break; + } + _ => {} + } + } + Self(&self.0[start..end]) + } + + /// Get a word at the certain position. + pub fn get(&self, index: usize) -> Option<&'static str> { + // SAFETY: we assume every word in the wordlist is a valid 0-terminated UTF-8 + // string. + self.0 + .get(index) + .map(|word| unsafe { from_utf8_unchecked(*word) }) + } + + /// How many words are currently in the list. + pub const fn len(&self) -> usize { + self.0.len() + } + + /// Iterator of all current words. + pub fn iter(&self) -> impl Iterator { + self.0 + .iter() + .map(|word| unsafe { from_utf8_unchecked(*word) }) + } +} + +/// Compare word from wordlist to a prefix. +/// +/// The comparison returns Less if the word comes lexicographically before all +/// possible words starting with `prefix`, and Greater if it comes after. +/// Equal is returned if the word starts with `prefix`. +unsafe fn prefix_cmp(prefix: &str, word: *const cty::c_char) -> Ordering { + // SAFETY: we assume `word` is a pointer to a 0-terminated string. + for (i, prefix_char) in prefix.as_bytes().iter().enumerate() { + let word_char = unsafe { *(word.add(i)) } as u8; + if word_char == 0 { + // Prefix is longer than word. + return Ordering::Less; + } else if *prefix_char != word_char { + return word_char.cmp(prefix_char); + } + } + Ordering::Equal +} + +unsafe fn from_utf8_unchecked<'a>(word: *const cty::c_char) -> &'a str { + // SAFETY: caller must pass a valid 0-terminated UTF-8 string. + // This assumption holds for usage on words of the BIP-39/SLIP-39 wordlists. + unsafe { + let word = CStr::from_ptr(word); + core::str::from_utf8_unchecked(word.to_bytes()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + const BIP39_WORD_COUNT: usize = ffi::BIP39_WORD_COUNT as usize; + const SLIP39_WORD_COUNT: usize = ffi::SLIP39_WORD_COUNT as usize; + + #[test] + fn test_filter_prefix_empty() { + let words = Wordlist::bip39().filter_prefix(""); + assert_eq!(words.len(), BIP39_WORD_COUNT); + let iter = words.iter(); + assert_eq!(iter.size_hint(), (BIP39_WORD_COUNT, Some(BIP39_WORD_COUNT))); + + let words = Wordlist::slip39().filter_prefix(""); + assert_eq!(words.len(), SLIP39_WORD_COUNT); + let iter = words.iter(); + assert_eq!( + iter.size_hint(), + (SLIP39_WORD_COUNT, Some(SLIP39_WORD_COUNT)) + ); + } + + #[test] + fn test_filter_prefix() { + let expected_result = vec!["strategy", "street", "strike", "strong", "struggle"]; + let result = Wordlist::bip39() + .filter_prefix("str") + .iter() + .collect::>(); + assert_eq!(result, expected_result); + } + + #[test] + fn test_filter_prefix_refine() { + let expected_result = vec!["strategy", "street", "strike", "strong", "struggle"]; + let words = Wordlist::bip39().filter_prefix("st"); + let result_a = words.filter_prefix("str").iter().collect::>(); + let result_b = Wordlist::bip39() + .filter_prefix("str") + .iter() + .collect::>(); + assert_eq!(result_a, expected_result); + assert_eq!(result_b, expected_result); + + let empty = words.filter_prefix("c"); + assert_eq!(empty.len(), 0); + } + + #[test] + fn test_wordlist_get() { + let words = Wordlist::bip39(); + assert_eq!(words.get(0), Some("abandon")); + assert_eq!(words.get(BIP39_WORD_COUNT - 1), Some("zoo")); + assert_eq!(words.get(BIP39_WORD_COUNT), None); + assert_eq!(words.get(BIP39_WORD_COUNT + 1), None); + + let filtered = words.filter_prefix("str"); + assert_eq!(filtered.get(0), Some("strategy")); + assert_eq!(filtered.get(filtered.len()), None); + } + + #[test] + fn test_filter_prefix_just_one() { + let expected_result = vec!["stick"]; + let result = Wordlist::bip39() + .filter_prefix("stick") + .iter() + .collect::>(); + assert_eq!(result, expected_result); + } +}