use crate::crypto::{Kek, SecretKey};
use crate::wordlist::{BITS_PER_WORD, WORDS};
use sha2::{Digest as _, Sha256};
use zeroize::Zeroize as _;
pub const ENTROPY_LEN: usize = 32;
pub const CHECKSUM_BITS: usize = 4;
#[expect(clippy::integer_division, reason = "编译期常量,静态断言保证整除")]
pub const WORD_COUNT: usize = (ENTROPY_LEN * 8 + CHECKSUM_BITS) / BITS_PER_WORD;
const _: () = assert!(
WORD_COUNT * BITS_PER_WORD == ENTROPY_LEN * 8 + CHECKSUM_BITS,
"词数与熵+校验的总位数必须严格相等,不允许有填充位"
);
const INFO_RECOVERY: &[u8] = b"omy/v1/recovery-kek";
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ParseError {
WordCount {
got: usize,
want: usize,
},
UnknownWord {
index: usize,
word: String,
suggestion: Option<&'static str>,
},
Checksum,
}
impl core::fmt::Display for ParseError {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
match self {
Self::WordCount { got, want } => {
write!(f, "恢复码应有 {want} 个词,实际数到 {got} 个")
}
Self::UnknownWord { index, word, suggestion } => {
let n = index.saturating_add(1);
match suggestion {
Some(s) => write!(f, "第 {n} 个词「{word}」不在词表中,是不是「{s}」?"),
None => write!(f, "第 {n} 个词「{word}」不在词表中"),
}
}
Self::Checksum => write!(
f,
"校验和不匹配:每个词都认识,但组合起来不对。\
常见原因是有两个词写反了顺序,或某个词抄成了另一个词"
),
}
}
}
impl core::error::Error for ParseError {}
#[derive(zeroize::Zeroize, zeroize::ZeroizeOnDrop)]
pub struct RecoveryCode {
entropy: [u8; ENTROPY_LEN],
}
impl core::fmt::Debug for RecoveryCode {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
f.write_str("RecoveryCode(<redacted>)")
}
}
impl RecoveryCode {
#[must_use]
pub fn generate() -> Self {
let mut entropy = [0u8; ENTROPY_LEN];
crate::util::fill_random(&mut entropy);
Self { entropy }
}
#[must_use]
pub const fn from_entropy(entropy: [u8; ENTROPY_LEN]) -> Self {
Self { entropy }
}
#[must_use]
pub const fn entropy(&self) -> &[u8; ENTROPY_LEN] {
&self.entropy
}
#[must_use]
pub fn to_words(&self) -> Vec<&'static str> {
let checksum = compute_checksum(&self.entropy);
let mut out = Vec::with_capacity(WORD_COUNT);
let mut acc: u32 = 0;
let mut bits: usize = 0;
for &b in &self.entropy {
acc = (acc << 8) | u32::from(b);
bits = bits.saturating_add(8);
while bits >= BITS_PER_WORD {
let shift = bits.saturating_sub(BITS_PER_WORD);
let idx = ((acc >> shift) & 0x3FF) as usize;
out.push(word_at(idx));
bits = shift;
}
}
acc = (acc << CHECKSUM_BITS) | u32::from(checksum);
bits = bits.saturating_add(CHECKSUM_BITS);
debug_assert_eq!(bits, BITS_PER_WORD, "尾部必须恰好凑满一个词,不留填充位");
let idx = (acc & 0x3FF) as usize;
out.push(word_at(idx));
debug_assert_eq!(out.len(), WORD_COUNT);
out
}
#[must_use]
pub fn to_phrase(&self) -> String {
self.to_words().join(" ")
}
pub fn from_words(words: &[&str]) -> core::result::Result<Self, ParseError> {
if words.len() != WORD_COUNT {
return Err(ParseError::WordCount { got: words.len(), want: WORD_COUNT });
}
let mut acc: u32 = 0;
let mut bits: usize = 0;
let mut entropy = [0u8; ENTROPY_LEN];
let mut filled = 0usize;
for (i, raw) in words.iter().enumerate() {
let w = raw.trim().to_lowercase();
let idx = index_of(&w).ok_or_else(|| ParseError::UnknownWord {
index: i,
word: w.clone(),
suggestion: nearest(&w),
})?;
let idx = u32::try_from(idx).unwrap_or(0);
acc = (acc << BITS_PER_WORD) | idx;
bits = bits.saturating_add(BITS_PER_WORD);
while bits >= 8 {
let shift = bits.saturating_sub(8);
let byte = ((acc >> shift) & 0xFF) as u8;
if filled < ENTROPY_LEN {
if let Some(slot) = entropy.get_mut(filled) {
*slot = byte;
}
filled = filled.saturating_add(1);
}
bits = shift;
}
}
debug_assert_eq!(bits, CHECKSUM_BITS, "解析后剩余位数必须恰好是校验位宽");
let tail = (acc & 0x0F) as u8;
let expect = compute_checksum(&entropy);
if tail != expect {
entropy.zeroize();
return Err(ParseError::Checksum);
}
Ok(Self { entropy })
}
pub fn from_phrase(phrase: &str) -> core::result::Result<Self, ParseError> {
let parts: Vec<&str> = phrase.split_whitespace().collect();
Self::from_words(&parts)
}
#[must_use]
pub fn to_kek(&self, vault_salt: &[u8; 16]) -> Kek {
let hk = hkdf::Hkdf::<Sha256>::new(Some(vault_salt), &self.entropy);
let mut out = [0u8; 32];
hk.expand(INFO_RECOVERY, &mut out)
.unwrap_or_else(|_| unreachable!("HKDF expand 32 bytes cannot fail"));
Kek::from_key(SecretKey::from_bytes(out))
}
}
fn compute_checksum(entropy: &[u8; ENTROPY_LEN]) -> u8 {
let mut h = Sha256::new();
h.update(b"omy/v1/recovery-checksum");
h.update(entropy);
let digest = h.finalize();
digest.first().map_or(0, |b| b >> 4)
}
fn word_at(idx: usize) -> &'static str {
WORDS.get(idx).copied().unwrap_or("academic")
}
fn index_of(word: &str) -> Option<usize> {
WORDS.binary_search(&word).ok()
}
fn nearest(word: &str) -> Option<&'static str> {
if word.len() >= 4 {
let p: String = word.chars().take(4).collect();
if let Some(hit) = WORDS.iter().find(|w| w.starts_with(&p)) {
return Some(hit);
}
}
WORDS.iter().copied().find(|w| edit_distance_le_1(word, w))
}
fn edit_distance_le_1(a: &str, b: &str) -> bool {
let (ab, bb) = (a.as_bytes(), b.as_bytes());
let (la, lb) = (ab.len(), bb.len());
if la.abs_diff(lb) > 1 {
return false;
}
let mut i = 0usize;
let mut j = 0usize;
let mut diff = 0usize;
while i < la && j < lb {
if ab.get(i) == bb.get(j) {
i = i.saturating_add(1);
j = j.saturating_add(1);
continue;
}
diff = diff.saturating_add(1);
if diff > 1 {
return false;
}
match la.cmp(&lb) {
core::cmp::Ordering::Greater => i = i.saturating_add(1),
core::cmp::Ordering::Less => j = j.saturating_add(1),
core::cmp::Ordering::Equal => {
i = i.saturating_add(1);
j = j.saturating_add(1);
}
}
}
diff.saturating_add(la.saturating_sub(i)).saturating_add(lb.saturating_sub(j)) <= 1
}
#[cfg(test)]
#[expect(
clippy::unwrap_used,
clippy::indexing_slicing,
clippy::panic,
reason = "测试里用断言与 panic 表达失败即可"
)]
mod tests {
use super::*;
#[test]
fn roundtrip_preserves_entropy() {
for _ in 0..50 {
let code = RecoveryCode::generate();
let phrase = code.to_phrase();
let back = RecoveryCode::from_phrase(&phrase).expect("解析自己生成的码");
assert_eq!(code.entropy(), back.entropy(), "往返后熵必须一致");
}
}
#[test]
fn encodes_to_expected_word_count() {
let code = RecoveryCode::from_entropy([0x5Au8; ENTROPY_LEN]);
let words = code.to_words();
assert_eq!(words.len(), 26, "26 词 × 10 bit = 260 = 256 熵 + 4 校验");
for w in &words {
assert!(WORDS.contains(w), "生成的词必须在词表内:{w}");
}
}
#[test]
fn normalizes_case_and_whitespace() {
let code = RecoveryCode::generate();
let words = code.to_words();
let messy = words
.iter()
.enumerate()
.map(|(i, w)| if i % 2 == 0 { w.to_uppercase() } else { format!(" {w} ") })
.collect::<Vec<_>>()
.join(" ");
let back = RecoveryCode::from_phrase(&messy).expect("应容忍大小写与空白");
assert_eq!(code.entropy(), back.entropy());
}
#[test]
fn wrong_word_count_is_reported_precisely() {
let code = RecoveryCode::generate();
let mut words = code.to_words();
words.pop();
let err = RecoveryCode::from_words(&words).unwrap_err();
assert_eq!(err, ParseError::WordCount { got: 25, want: 26 });
assert!(err.to_string().contains("25"), "实际:{err}");
}
#[test]
fn typo_points_at_the_word_and_suggests() {
let code = RecoveryCode::generate();
let mut words: Vec<String> = code.to_words().iter().map(|s| (*s).to_owned()).collect();
let original = words[6].clone();
words[6].pop();
let refs: Vec<&str> = words.iter().map(String::as_str).collect();
match RecoveryCode::from_words(&refs) {
Err(ParseError::UnknownWord { index, suggestion, .. }) => {
assert_eq!(index, 6, "必须指出是第 7 个词(下标 6)");
assert_eq!(
suggestion,
Some(original.as_str()),
"前 4 字母唯一,应当能精确建议回原词"
);
}
Err(ParseError::Checksum) => {}
other => panic!("期望 UnknownWord 或 Checksum,实际 {other:?}"),
}
}
#[test]
fn swapped_words_are_caught_by_checksum() {
let mut caught = 0;
let mut attempts = 0;
for _ in 0..40 {
let code = RecoveryCode::generate();
let mut words = code.to_words();
if words[3] == words[9] {
continue; }
words.swap(3, 9);
attempts += 1;
match RecoveryCode::from_words(&words) {
Err(ParseError::Checksum) => caught += 1,
Ok(_) => {}
other => panic!("交换词序不该报这个错:{other:?}"),
}
}
assert!(
caught * 100 >= attempts * 60,
"校验和应抓到绝大多数词序错误,实际 {caught}/{attempts}"
);
}
#[test]
fn unknown_word_is_rejected() {
let code = RecoveryCode::generate();
let mut words: Vec<String> = code.to_words().iter().map(|s| (*s).to_owned()).collect();
words[0] = String::from("zzzzzz");
let refs: Vec<&str> = words.iter().map(String::as_str).collect();
let err = RecoveryCode::from_words(&refs).unwrap_err();
assert!(matches!(err, ParseError::UnknownWord { index: 0, .. }), "实际 {err:?}");
}
#[test]
fn kek_derivation_is_discriminating() {
let salt_a = [0x11u8; 16];
let salt_b = [0x22u8; 16];
let a = RecoveryCode::from_entropy([1u8; ENTROPY_LEN]);
let b = RecoveryCode::from_entropy([2u8; ENTROPY_LEN]);
let fp = |k: &Kek| *k.derive_slot_key(&[0u8; 16], 0).as_bytes();
assert_ne!(fp(&a.to_kek(&salt_a)), fp(&b.to_kek(&salt_a)), "不同码必须不同 KEK");
assert_ne!(fp(&a.to_kek(&salt_a)), fp(&a.to_kek(&salt_b)), "跨库必须不同 KEK");
assert_eq!(fp(&a.to_kek(&salt_a)), fp(&a.to_kek(&salt_a)), "必须可重现");
}
#[test]
fn recovery_kek_differs_from_password_path() {
let salt = [0x33u8; 16];
let code = RecoveryCode::from_entropy([7u8; ENTROPY_LEN]);
let phrase = code.to_phrase();
let as_password =
Kek::from_password(phrase.as_bytes(), &salt, crate::crypto::Argon2Params::TEST_WEAK)
.expect("argon2");
let fp = |k: &Kek| *k.derive_slot_key(&[0u8; 16], 0).as_bytes();
assert_ne!(fp(&code.to_kek(&salt)), fp(&as_password), "两条派生路径必须域分隔");
}
#[test]
fn debug_is_redacted() {
let code = RecoveryCode::from_entropy([0xABu8; ENTROPY_LEN]);
let s = format!("{code:?}");
assert!(!s.contains("ab"), "Debug 不得打印熵内容:{s}");
assert!(s.contains("redacted"));
}
#[test]
fn edit_distance_helper_is_correct() {
assert!(edit_distance_le_1("academic", "academic"), "相同");
assert!(edit_distance_le_1("acadmic", "academic"), "少一个字母");
assert!(edit_distance_le_1("academics", "academic"), "多一个字母");
assert!(edit_distance_le_1("acedemic", "academic"), "改一个字母");
assert!(!edit_distance_le_1("acdmic", "academic"), "差两个字母");
assert!(!edit_distance_le_1("zebra", "academic"), "完全不同");
}
}