1use crate::crypto::{Kek, SecretKey};
41use crate::wordlist::{BITS_PER_WORD, WORDS};
42use sha2::{Digest as _, Sha256};
43use zeroize::Zeroize as _;
44
45pub const ENTROPY_LEN: usize = 32;
47pub const CHECKSUM_BITS: usize = 4;
49#[expect(clippy::integer_division, reason = "编译期常量,静态断言保证整除")]
55pub const WORD_COUNT: usize = (ENTROPY_LEN * 8 + CHECKSUM_BITS) / BITS_PER_WORD;
56
57const _: () = assert!(
60 WORD_COUNT * BITS_PER_WORD == ENTROPY_LEN * 8 + CHECKSUM_BITS,
61 "词数与熵+校验的总位数必须严格相等,不允许有填充位"
62);
63
64const INFO_RECOVERY: &[u8] = b"omy/v1/recovery-kek";
69
70#[derive(Debug, Clone, PartialEq, Eq)]
76pub enum ParseError {
77 WordCount {
79 got: usize,
81 want: usize,
83 },
84 UnknownWord {
90 index: usize,
92 word: String,
94 suggestion: Option<&'static str>,
96 },
97 Checksum,
102}
103
104impl core::fmt::Display for ParseError {
105 fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
106 match self {
107 Self::WordCount { got, want } => {
108 write!(f, "恢复码应有 {want} 个词,实际数到 {got} 个")
109 }
110 Self::UnknownWord { index, word, suggestion } => {
111 let n = index.saturating_add(1);
112 match suggestion {
113 Some(s) => write!(f, "第 {n} 个词「{word}」不在词表中,是不是「{s}」?"),
114 None => write!(f, "第 {n} 个词「{word}」不在词表中"),
115 }
116 }
117 Self::Checksum => write!(
118 f,
119 "校验和不匹配:每个词都认识,但组合起来不对。\
120 常见原因是有两个词写反了顺序,或某个词抄成了另一个词"
121 ),
122 }
123 }
124}
125
126impl core::error::Error for ParseError {}
127
128#[derive(zeroize::Zeroize, zeroize::ZeroizeOnDrop)]
133pub struct RecoveryCode {
134 entropy: [u8; ENTROPY_LEN],
135}
136
137impl core::fmt::Debug for RecoveryCode {
138 fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
140 f.write_str("RecoveryCode(<redacted>)")
141 }
142}
143
144impl RecoveryCode {
145 #[must_use]
147 pub fn generate() -> Self {
148 let mut entropy = [0u8; ENTROPY_LEN];
149 crate::util::fill_random(&mut entropy);
150 Self { entropy }
151 }
152
153 #[must_use]
155 pub const fn from_entropy(entropy: [u8; ENTROPY_LEN]) -> Self {
156 Self { entropy }
157 }
158
159 #[must_use]
161 pub const fn entropy(&self) -> &[u8; ENTROPY_LEN] {
162 &self.entropy
163 }
164
165 #[must_use]
167 pub fn to_words(&self) -> Vec<&'static str> {
168 let checksum = compute_checksum(&self.entropy);
169 let mut out = Vec::with_capacity(WORD_COUNT);
173 let mut acc: u32 = 0;
174 let mut bits: usize = 0;
175 for &b in &self.entropy {
176 acc = (acc << 8) | u32::from(b);
177 bits = bits.saturating_add(8);
178 while bits >= BITS_PER_WORD {
179 let shift = bits.saturating_sub(BITS_PER_WORD);
180 let idx = ((acc >> shift) & 0x3FF) as usize;
181 out.push(word_at(idx));
182 bits = shift;
183 }
184 }
185 acc = (acc << CHECKSUM_BITS) | u32::from(checksum);
187 bits = bits.saturating_add(CHECKSUM_BITS);
188 debug_assert_eq!(bits, BITS_PER_WORD, "尾部必须恰好凑满一个词,不留填充位");
189 let idx = (acc & 0x3FF) as usize;
190 out.push(word_at(idx));
191
192 debug_assert_eq!(out.len(), WORD_COUNT);
193 out
194 }
195
196 #[must_use]
198 pub fn to_phrase(&self) -> String {
199 self.to_words().join(" ")
200 }
201
202 pub fn from_words(words: &[&str]) -> core::result::Result<Self, ParseError> {
211 if words.len() != WORD_COUNT {
212 return Err(ParseError::WordCount { got: words.len(), want: WORD_COUNT });
213 }
214
215 let mut acc: u32 = 0;
216 let mut bits: usize = 0;
217 let mut entropy = [0u8; ENTROPY_LEN];
218 let mut filled = 0usize;
219
220 for (i, raw) in words.iter().enumerate() {
221 let w = raw.trim().to_lowercase();
222 let idx = index_of(&w).ok_or_else(|| ParseError::UnknownWord {
223 index: i,
224 word: w.clone(),
225 suggestion: nearest(&w),
226 })?;
227 let idx = u32::try_from(idx).unwrap_or(0);
231 acc = (acc << BITS_PER_WORD) | idx;
232 bits = bits.saturating_add(BITS_PER_WORD);
233 while bits >= 8 {
234 let shift = bits.saturating_sub(8);
235 let byte = ((acc >> shift) & 0xFF) as u8;
236 if filled < ENTROPY_LEN {
237 if let Some(slot) = entropy.get_mut(filled) {
239 *slot = byte;
240 }
241 filled = filled.saturating_add(1);
242 }
243 bits = shift;
244 }
245 }
246 debug_assert_eq!(bits, CHECKSUM_BITS, "解析后剩余位数必须恰好是校验位宽");
248 let tail = (acc & 0x0F) as u8;
249
250 let expect = compute_checksum(&entropy);
251 if tail != expect {
252 entropy.zeroize();
255 return Err(ParseError::Checksum);
256 }
257
258 Ok(Self { entropy })
259 }
260
261 pub fn from_phrase(phrase: &str) -> core::result::Result<Self, ParseError> {
267 let parts: Vec<&str> = phrase.split_whitespace().collect();
268 Self::from_words(&parts)
269 }
270
271 #[must_use]
276 pub fn to_kek(&self, vault_salt: &[u8; 16]) -> Kek {
277 let hk = hkdf::Hkdf::<Sha256>::new(Some(vault_salt), &self.entropy);
278 let mut out = [0u8; 32];
279 hk.expand(INFO_RECOVERY, &mut out)
281 .unwrap_or_else(|_| unreachable!("HKDF expand 32 bytes cannot fail"));
282 Kek::from_key(SecretKey::from_bytes(out))
283 }
284}
285
286fn compute_checksum(entropy: &[u8; ENTROPY_LEN]) -> u8 {
296 let mut h = Sha256::new();
297 h.update(b"omy/v1/recovery-checksum");
298 h.update(entropy);
299 let digest = h.finalize();
300 digest.first().map_or(0, |b| b >> 4)
302}
303
304fn word_at(idx: usize) -> &'static str {
309 WORDS.get(idx).copied().unwrap_or("academic")
310}
311
312fn index_of(word: &str) -> Option<usize> {
314 WORDS.binary_search(&word).ok()
315}
316
317fn nearest(word: &str) -> Option<&'static str> {
325 if word.len() >= 4 {
326 let p: String = word.chars().take(4).collect();
327 if let Some(hit) = WORDS.iter().find(|w| w.starts_with(&p)) {
328 return Some(hit);
329 }
330 }
331 WORDS.iter().copied().find(|w| edit_distance_le_1(word, w))
332}
333
334fn edit_distance_le_1(a: &str, b: &str) -> bool {
338 let (ab, bb) = (a.as_bytes(), b.as_bytes());
339 let (la, lb) = (ab.len(), bb.len());
340 if la.abs_diff(lb) > 1 {
341 return false;
342 }
343 let mut i = 0usize;
344 let mut j = 0usize;
345 let mut diff = 0usize;
346 while i < la && j < lb {
347 if ab.get(i) == bb.get(j) {
348 i = i.saturating_add(1);
349 j = j.saturating_add(1);
350 continue;
351 }
352 diff = diff.saturating_add(1);
353 if diff > 1 {
354 return false;
355 }
356 match la.cmp(&lb) {
357 core::cmp::Ordering::Greater => i = i.saturating_add(1),
358 core::cmp::Ordering::Less => j = j.saturating_add(1),
359 core::cmp::Ordering::Equal => {
360 i = i.saturating_add(1);
361 j = j.saturating_add(1);
362 }
363 }
364 }
365 diff.saturating_add(la.saturating_sub(i)).saturating_add(lb.saturating_sub(j)) <= 1
367}
368
369#[cfg(test)]
370#[expect(
371 clippy::unwrap_used,
372 clippy::indexing_slicing,
373 clippy::panic,
374 reason = "测试里用断言与 panic 表达失败即可"
375)]
376mod tests {
377 use super::*;
378
379 #[test]
384 fn roundtrip_preserves_entropy() {
385 for _ in 0..50 {
386 let code = RecoveryCode::generate();
387 let phrase = code.to_phrase();
388 let back = RecoveryCode::from_phrase(&phrase).expect("解析自己生成的码");
389 assert_eq!(code.entropy(), back.entropy(), "往返后熵必须一致");
390 }
391 }
392
393 #[test]
395 fn encodes_to_expected_word_count() {
396 let code = RecoveryCode::from_entropy([0x5Au8; ENTROPY_LEN]);
397 let words = code.to_words();
398 assert_eq!(words.len(), 26, "26 词 × 10 bit = 260 = 256 熵 + 4 校验");
399 for w in &words {
400 assert!(WORDS.contains(w), "生成的词必须在词表内:{w}");
401 }
402 }
403
404 #[test]
409 fn normalizes_case_and_whitespace() {
410 let code = RecoveryCode::generate();
411 let words = code.to_words();
412 let messy = words
413 .iter()
414 .enumerate()
415 .map(|(i, w)| if i % 2 == 0 { w.to_uppercase() } else { format!(" {w} ") })
416 .collect::<Vec<_>>()
417 .join(" ");
418 let back = RecoveryCode::from_phrase(&messy).expect("应容忍大小写与空白");
419 assert_eq!(code.entropy(), back.entropy());
420 }
421
422 #[test]
424 fn wrong_word_count_is_reported_precisely() {
425 let code = RecoveryCode::generate();
426 let mut words = code.to_words();
427 words.pop();
428 let err = RecoveryCode::from_words(&words).unwrap_err();
429 assert_eq!(err, ParseError::WordCount { got: 25, want: 26 });
430 assert!(err.to_string().contains("25"), "实际:{err}");
432 }
433
434 #[test]
439 fn typo_points_at_the_word_and_suggests() {
440 let code = RecoveryCode::generate();
441 let mut words: Vec<String> = code.to_words().iter().map(|s| (*s).to_owned()).collect();
442 let original = words[6].clone();
443 words[6].pop();
445 let refs: Vec<&str> = words.iter().map(String::as_str).collect();
446
447 match RecoveryCode::from_words(&refs) {
448 Err(ParseError::UnknownWord { index, suggestion, .. }) => {
449 assert_eq!(index, 6, "必须指出是第 7 个词(下标 6)");
450 assert_eq!(
451 suggestion,
452 Some(original.as_str()),
453 "前 4 字母唯一,应当能精确建议回原词"
454 );
455 }
456 Err(ParseError::Checksum) => {}
459 other => panic!("期望 UnknownWord 或 Checksum,实际 {other:?}"),
460 }
461 }
462
463 #[test]
468 fn swapped_words_are_caught_by_checksum() {
469 let mut caught = 0;
470 let mut attempts = 0;
471 for _ in 0..40 {
472 let code = RecoveryCode::generate();
473 let mut words = code.to_words();
474 if words[3] == words[9] {
475 continue; }
477 words.swap(3, 9);
478 attempts += 1;
479 match RecoveryCode::from_words(&words) {
480 Err(ParseError::Checksum) => caught += 1,
481 Ok(_) => {}
483 other => panic!("交换词序不该报这个错:{other:?}"),
484 }
485 }
486 assert!(
489 caught * 100 >= attempts * 60,
490 "校验和应抓到绝大多数词序错误,实际 {caught}/{attempts}"
491 );
492 }
493
494 #[test]
496 fn unknown_word_is_rejected() {
497 let code = RecoveryCode::generate();
498 let mut words: Vec<String> = code.to_words().iter().map(|s| (*s).to_owned()).collect();
499 words[0] = String::from("zzzzzz");
500 let refs: Vec<&str> = words.iter().map(String::as_str).collect();
501 let err = RecoveryCode::from_words(&refs).unwrap_err();
502 assert!(matches!(err, ParseError::UnknownWord { index: 0, .. }), "实际 {err:?}");
503 }
504
505 #[test]
507 fn kek_derivation_is_discriminating() {
508 let salt_a = [0x11u8; 16];
509 let salt_b = [0x22u8; 16];
510 let a = RecoveryCode::from_entropy([1u8; ENTROPY_LEN]);
511 let b = RecoveryCode::from_entropy([2u8; ENTROPY_LEN]);
512
513 let fp = |k: &Kek| *k.derive_slot_key(&[0u8; 16], 0).as_bytes();
514 assert_ne!(fp(&a.to_kek(&salt_a)), fp(&b.to_kek(&salt_a)), "不同码必须不同 KEK");
515 assert_ne!(fp(&a.to_kek(&salt_a)), fp(&a.to_kek(&salt_b)), "跨库必须不同 KEK");
516 assert_eq!(fp(&a.to_kek(&salt_a)), fp(&a.to_kek(&salt_a)), "必须可重现");
518 }
519
520 #[test]
525 fn recovery_kek_differs_from_password_path() {
526 let salt = [0x33u8; 16];
527 let code = RecoveryCode::from_entropy([7u8; ENTROPY_LEN]);
528 let phrase = code.to_phrase();
529 let as_password =
530 Kek::from_password(phrase.as_bytes(), &salt, crate::crypto::Argon2Params::TEST_WEAK)
531 .expect("argon2");
532 let fp = |k: &Kek| *k.derive_slot_key(&[0u8; 16], 0).as_bytes();
533 assert_ne!(fp(&code.to_kek(&salt)), fp(&as_password), "两条派生路径必须域分隔");
534 }
535
536 #[test]
538 fn debug_is_redacted() {
539 let code = RecoveryCode::from_entropy([0xABu8; ENTROPY_LEN]);
540 let s = format!("{code:?}");
541 assert!(!s.contains("ab"), "Debug 不得打印熵内容:{s}");
542 assert!(s.contains("redacted"));
543 }
544
545 #[test]
547 fn edit_distance_helper_is_correct() {
548 assert!(edit_distance_le_1("academic", "academic"), "相同");
549 assert!(edit_distance_le_1("acadmic", "academic"), "少一个字母");
550 assert!(edit_distance_le_1("academics", "academic"), "多一个字母");
551 assert!(edit_distance_le_1("acedemic", "academic"), "改一个字母");
552 assert!(!edit_distance_le_1("acdmic", "academic"), "差两个字母");
553 assert!(!edit_distance_le_1("zebra", "academic"), "完全不同");
554 }
555}