keyfork/keyfork-mnemonic-util/src/lib.rs

251 lines
7.8 KiB
Rust

use std::{collections::HashMap, str::FromStr, sync::Arc};
use sha2::{Digest, Sha256};
use std::{error::Error, fmt::Display};
/// The error type representing a failure to create a [`Mnemonic`]. These errors only occur during
/// [`Mnemonic`] creation.
#[derive(Debug, Clone)]
pub enum MnemonicGenerationError {
/// The amount of bits passed to a mnemonic must be divisible by 32.
InvalidByteCount(usize),
/// The length of a mnemonic in bits must be within the BIP-0039 range, and supported by the
/// library. Currently, only 128, 192 (for testing purposes), and 256 are supported.
InvalidByteLength(usize),
}
impl Display for MnemonicGenerationError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
MnemonicGenerationError::InvalidByteCount(count) => {
write!(f, "Invalid byte count: {count}, must be divisible by 8")
}
MnemonicGenerationError::InvalidByteLength(count) => {
write!(f, "Invalid byte length: {count}, must be 128 or 256")
}
}
}
}
impl Error for MnemonicGenerationError {}
/// A BIP-0039 compatible list of words.
#[derive(Debug, Clone)]
pub struct Wordlist(Vec<String>);
impl Default for Wordlist {
/// Returns the English wordlist in the Bitcoin BIP-0039 specification.
fn default() -> Self {
// TODO: English is the only supported language.
let wordlist_file = include_str!("data/wordlist.txt");
Wordlist(
wordlist_file
.lines()
// skip 1: comment at top of file to point to BIP-0039 source.
.skip(1)
.map(|x| x.trim().to_string())
.collect(),
)
}
}
impl Wordlist {
/// Return an Arced version of the Wordlist
#[allow(clippy::must_use_candidate)]
pub fn arc(self) -> Arc<Self> {
Arc::new(self)
}
/// Given an index, get a word from the wordlist.
fn get_word(&self, word: usize) -> Option<&String> {
self.0.get(word)
}
fn inner(&self) -> &Vec<String> {
&self.0
}
#[cfg(test)]
fn into_inner(self) -> Vec<String> {
self.0
}
}
/// A BIP-0039 mnemonic with reference to a [`Wordlist`].
#[derive(Debug, Clone)]
pub struct Mnemonic {
words: Vec<usize>,
wordlist: Arc<Wordlist>,
}
impl Display for Mnemonic {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
let mut iter = self.words.iter().peekable();
while let Some(word_index) = iter.next() {
let word = self.wordlist.get_word(*word_index).expect("word");
write!(f, "{word}")?;
if iter.peek().is_some() {
write!(f, " ")?;
}
}
Ok(())
}
}
/// The error type representing a failure to parse a [`Mnemonic`]. These errors only occur during
/// [`Mnemonic`] creation.
#[derive(Debug, Clone)]
pub enum MnemonicFromStrError {
/// The amount of words used to parse a mnemonic was not correct.
InvalidWordCount(usize),
/// One of the words used to generate the mnemonic was not found in the default wordlist.
InvalidWord(usize),
}
impl Display for MnemonicFromStrError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
MnemonicFromStrError::InvalidWordCount(count) => {
write!(f, "Incorrect word count: {count}")
}
MnemonicFromStrError::InvalidWord(index) => {
write!(f, "Unknown word at index: {index}")
}
}
}
}
impl Error for MnemonicFromStrError {}
impl FromStr for Mnemonic {
type Err = MnemonicFromStrError;
fn from_str(s: &str) -> Result<Self, Self::Err> {
let wordlist = Wordlist::default().arc();
let hm: HashMap<&str, usize> = wordlist
.inner()
.iter()
.enumerate()
.map(|(a, b)| (b.as_str(), a))
.collect();
let mut words: Vec<usize> = Vec::with_capacity(24);
for (index, word) in s.split_whitespace().enumerate() {
match hm.get(&word) {
Some(id) => words.push(*id),
None => return Err(MnemonicFromStrError::InvalidWord(index)),
}
}
if ![12, 24].contains(&words.len()) {
return Err(MnemonicFromStrError::InvalidWordCount(words.len()));
}
Ok(Mnemonic { words, wordlist })
}
}
fn generate_slice_hash(data: &[u8]) -> Vec<u8> {
let mut hasher = Sha256::new();
hasher.update(data);
hasher.finalize().to_vec()
}
impl Mnemonic {
/// Generate a [`Mnemonic`] from the provided entropy and [`Wordlist`].
///
/// # Errors
/// An error may be returned if the entropy is not within the acceptable lengths.
pub fn from_entropy(
bytes: &[u8],
wordlist: Arc<Wordlist>,
) -> Result<Mnemonic, MnemonicGenerationError> {
let bit_count = bytes.len() * 8;
let hash = generate_slice_hash(bytes);
if bit_count % 32 != 0 {
return Err(MnemonicGenerationError::InvalidByteCount(bit_count));
}
// 192 supported for test suite
if ![128, 192, 256].contains(&bit_count) {
return Err(MnemonicGenerationError::InvalidByteLength(bit_count));
}
let mut bits = vec![false; bit_count + bit_count / 32];
for byte_index in 0..bit_count / 8 {
for bit_index in 0..8 {
bits[byte_index * 8 + bit_index] = (bytes[byte_index] & (1 << (7 - bit_index))) > 0;
}
}
for check_bit in 0..bit_count / 32 {
bits[bit_count + check_bit] = (hash[check_bit / 8] & (1 << (7 - (check_bit % 8)))) > 0;
}
let words = bits
// NOTE: Tested with all approved variants. Always divisible by 11.
.chunks_exact(11)
.map(|chunk| {
// NOTE: usize to use for indexing wordlist later
let mut num = 0usize;
for i in 0..11 {
num += usize::from(chunk[10 - i]) << i;
}
num
})
.collect::<Vec<_>>();
Ok(Mnemonic { words, wordlist })
}
#[must_use]
pub fn into_inner(self) -> (Vec<usize>, Arc<Wordlist>) {
(self.words, self.wordlist)
}
}
#[cfg(test)]
mod tests {
use std::{fs::File, io::Read};
use super::*;
#[test]
fn wordlist_word_count_correct() {
let wordlist = Wordlist::default().into_inner();
assert_eq!(
wordlist.len(),
2usize.pow(11),
"Wordlist did not include correct word count"
);
}
#[test]
fn conforms_to_trezor_tests() {
let content = include_str!("data/vectors.json");
let jsonobj: serde_json::Value = serde_json::from_str(content).unwrap();
let wordlist = Wordlist::default().arc();
for test in jsonobj["english"].as_array().unwrap() {
let [ref hex_, ref seed, ..] = test.as_array().unwrap()[..] else {
panic!("bad test: {test}");
};
let hex = hex::decode(hex_.as_str().unwrap()).unwrap();
let mnemonic = Mnemonic::from_entropy(&hex, wordlist.clone()).unwrap();
assert_eq!(mnemonic.to_string(), seed.as_str().unwrap());
}
}
#[test]
fn matches_bip39_crate() {
let mut random_handle = File::open("/dev/random").unwrap();
let entropy = &mut [0u8; 256 / 8];
random_handle.read_exact(&mut entropy[..]).unwrap();
let wordlist = Wordlist::default().arc();
let my_mnemonic = super::Mnemonic::from_entropy(&entropy[..256 / 8], wordlist).unwrap();
let their_mnemonic = bip39::Mnemonic::from_entropy(&entropy[..256 / 8]).unwrap();
assert_eq!(my_mnemonic.to_string(), their_mnemonic.to_string());
}
}