1use sha2::{Digest, Sha256};
4
5use super::error::TransparencyError;
6use super::types::MerkleHash;
7
8const LEAF_PREFIX: u8 = 0x00;
10const NODE_PREFIX: u8 = 0x01;
12
13pub fn hash_leaf(data: &[u8]) -> MerkleHash {
23 let mut hasher = Sha256::new();
24 hasher.update([LEAF_PREFIX]);
25 hasher.update(data);
26 let digest = hasher.finalize();
27 let mut out = [0u8; 32];
28 out.copy_from_slice(&digest);
29 MerkleHash::from_bytes(out)
30}
31
32pub fn hash_children(left: &MerkleHash, right: &MerkleHash) -> MerkleHash {
43 let mut hasher = Sha256::new();
44 hasher.update([NODE_PREFIX]);
45 hasher.update(left.as_bytes());
46 hasher.update(right.as_bytes());
47 let digest = hasher.finalize();
48 let mut out = [0u8; 32];
49 out.copy_from_slice(&digest);
50 MerkleHash::from_bytes(out)
51}
52
53pub fn verify_inclusion(
70 leaf_hash: &MerkleHash,
71 index: u64,
72 size: u64,
73 proof: &[MerkleHash],
74 root: &MerkleHash,
75) -> Result<(), TransparencyError> {
76 if size == 0 {
77 return Err(TransparencyError::InvalidProof("tree size is 0".into()));
78 }
79 if index >= size {
80 return Err(TransparencyError::InvalidProof(format!(
81 "index {index} >= size {size}"
82 )));
83 }
84
85 let (computed, _) = root_from_inclusion_proof(leaf_hash, index, size, proof)?;
86
87 if computed != *root {
88 return Err(TransparencyError::RootMismatch {
89 expected: root.to_string(),
90 actual: computed.to_string(),
91 });
92 }
93 Ok(())
94}
95
96fn root_from_inclusion_proof(
100 leaf_hash: &MerkleHash,
101 index: u64,
102 size: u64,
103 proof: &[MerkleHash],
104) -> Result<(MerkleHash, usize), TransparencyError> {
105 let expected_len = inclusion_proof_length(index, size);
106 if proof.len() != expected_len {
107 return Err(TransparencyError::InvalidProof(format!(
108 "expected {expected_len} proof elements, got {}",
109 proof.len()
110 )));
111 }
112
113 let mut hash = *leaf_hash;
114 let mut idx = index;
115 let mut level_size = size;
116 let mut pos = 0;
117
118 while level_size > 1 {
119 if pos >= proof.len() {
120 return Err(TransparencyError::InvalidProof("proof too short".into()));
121 }
122 if idx & 1 == 1 || idx + 1 == level_size {
123 if idx & 1 == 1 {
124 hash = hash_children(&proof[pos], &hash);
125 pos += 1;
126 }
127 } else {
128 hash = hash_children(&hash, &proof[pos]);
129 pos += 1;
130 }
131 idx >>= 1;
132 level_size = (level_size + 1) >> 1;
133 }
134
135 Ok((hash, pos))
136}
137
138fn inclusion_proof_length(index: u64, size: u64) -> usize {
140 if size <= 1 {
141 return 0;
142 }
143 let mut length = 0;
144 let mut idx = index;
145 let mut level_size = size;
146 while level_size > 1 {
147 if idx & 1 == 1 || idx + 1 < level_size {
148 length += 1;
149 }
150 idx >>= 1;
151 level_size = (level_size + 1) >> 1;
152 }
153 length
154}
155
156pub fn verify_consistency(
172 old_size: u64,
173 new_size: u64,
174 proof: &[MerkleHash],
175 old_root: &MerkleHash,
176 new_root: &MerkleHash,
177) -> Result<(), TransparencyError> {
178 if old_size == 0 {
179 if proof.is_empty() {
180 return Ok(());
181 }
182 return Err(TransparencyError::ConsistencyError(
183 "non-empty proof for empty old tree".into(),
184 ));
185 }
186 if old_size > new_size {
187 return Err(TransparencyError::ConsistencyError(format!(
188 "old size {old_size} > new size {new_size}"
189 )));
190 }
191 if old_size == new_size {
192 if !proof.is_empty() {
193 return Err(TransparencyError::ConsistencyError(
194 "non-empty proof for equal sizes".into(),
195 ));
196 }
197 if old_root != new_root {
198 return Err(TransparencyError::RootMismatch {
199 expected: old_root.to_string(),
200 actual: new_root.to_string(),
201 });
202 }
203 return Ok(());
204 }
205
206 let new_computed = new_root_from_consistency_proof(old_size, new_size, proof, old_root)?;
210
211 if new_computed != *new_root {
212 return Err(TransparencyError::RootMismatch {
213 expected: new_root.to_string(),
214 actual: new_computed.to_string(),
215 });
216 }
217 Ok(())
218}
219
220fn new_root_from_consistency_proof(
233 old_size: u64,
234 new_size: u64,
235 proof: &[MerkleHash],
236 old_root: &MerkleHash,
237) -> Result<MerkleHash, TransparencyError> {
238 let _ = new_size; let (mut fn_hash, mut fr_hash, start) = if old_size.is_power_of_two() {
241 (*old_root, *old_root, 0)
243 } else {
244 if proof.is_empty() {
245 return Err(TransparencyError::ConsistencyError(
246 "proof too short".into(),
247 ));
248 }
249 (proof[0], proof[0], 1)
250 };
251
252 let mut pos = start;
253
254 if !old_size.is_power_of_two() {
256 let mut bit = old_size - 1;
257 while bit > 0 {
258 if pos >= proof.len() {
259 return Err(TransparencyError::ConsistencyError(
260 "proof too short during decomposition".into(),
261 ));
262 }
263 if bit & 1 != 0 {
264 fn_hash = hash_children(&proof[pos], &fn_hash);
265 fr_hash = hash_children(&proof[pos], &fr_hash);
266 } else {
267 fr_hash = hash_children(&fr_hash, &proof[pos]);
268 }
269 pos += 1;
270 bit >>= 1;
271 }
272
273 if fn_hash != *old_root {
274 return Err(TransparencyError::RootMismatch {
275 expected: old_root.to_string(),
276 actual: fn_hash.to_string(),
277 });
278 }
279 }
280
281 while pos < proof.len() {
283 fr_hash = hash_children(&fr_hash, &proof[pos]);
284 pos += 1;
285 }
286
287 Ok(fr_hash)
288}
289
290pub fn compute_root(leaves: &[MerkleHash]) -> MerkleHash {
303 match leaves.len() {
304 0 => MerkleHash::EMPTY,
305 1 => leaves[0],
306 n => {
307 let k = largest_power_of_2_lt(n as u64) as usize;
308 let left = compute_root(&leaves[..k]);
309 let right = compute_root(&leaves[k..]);
310 hash_children(&left, &right)
311 }
312 }
313}
314
315fn largest_power_of_2_lt(n: u64) -> u64 {
317 debug_assert!(n > 1);
318 if n.is_power_of_two() {
319 n / 2
320 } else {
321 1u64 << (63 - n.leading_zeros())
322 }
323}
324
325pub fn prove_inclusion(
343 leaves: &[MerkleHash],
344 index: u64,
345) -> Result<Vec<MerkleHash>, TransparencyError> {
346 let size = leaves.len() as u64;
347 if size == 0 {
348 return Err(TransparencyError::InvalidProof("tree size is 0".into()));
349 }
350 if index >= size {
351 return Err(TransparencyError::InvalidProof(format!(
352 "index {index} >= size {size}"
353 )));
354 }
355 Ok(inclusion_path(leaves, index))
356}
357
358fn inclusion_path(leaves: &[MerkleHash], index: u64) -> Vec<MerkleHash> {
361 let n = leaves.len();
362 if n <= 1 {
363 return Vec::new();
364 }
365 let k = largest_power_of_2_lt(n as u64) as usize;
366 if (index as usize) < k {
367 let mut path = inclusion_path(&leaves[..k], index);
368 path.push(compute_root(&leaves[k..]));
369 path
370 } else {
371 let mut path = inclusion_path(&leaves[k..], index - k as u64);
372 path.push(compute_root(&leaves[..k]));
373 path
374 }
375}
376
377#[cfg(test)]
378mod tests {
379 use super::*;
380
381 #[test]
382 fn hash_leaf_domain_separation() {
383 let data = b"test data";
384 let h = hash_leaf(data);
385
386 let mut hasher = Sha256::new();
388 hasher.update([0x00]);
389 hasher.update(data);
390 let expected = hasher.finalize();
391
392 assert_eq!(h.as_bytes(), expected.as_slice());
393 }
394
395 #[test]
396 fn hash_children_domain_separation() {
397 let left = MerkleHash::from_bytes([0x11; 32]);
398 let right = MerkleHash::from_bytes([0x22; 32]);
399 let h = hash_children(&left, &right);
400
401 let mut hasher = Sha256::new();
402 hasher.update([0x01]);
403 hasher.update([0x11; 32]);
404 hasher.update([0x22; 32]);
405 let expected = hasher.finalize();
406
407 assert_eq!(h.as_bytes(), expected.as_slice());
408 }
409
410 #[test]
411 fn leaf_and_children_produce_different_hashes() {
412 let data = [0xab; 64];
413 let leaf = hash_leaf(&data);
414
415 let left = MerkleHash::from_bytes(data[..32].try_into().unwrap());
416 let right = MerkleHash::from_bytes(data[32..].try_into().unwrap());
417 let node = hash_children(&left, &right);
418
419 assert_ne!(leaf, node);
420 }
421
422 #[test]
423 fn compute_root_single_leaf() {
424 let h = MerkleHash::from_bytes([0x42; 32]);
425 assert_eq!(compute_root(&[h]), h);
426 }
427
428 #[test]
429 fn prove_inclusion_roundtrips_against_verify_inclusion() {
430 for size in 1u64..=20 {
433 let leaves: Vec<MerkleHash> = (0..size)
434 .map(|i| hash_leaf(format!("leaf-{i}").as_bytes()))
435 .collect();
436 let root = compute_root(&leaves);
437 for index in 0..size {
438 let hashes = prove_inclusion(&leaves, index).unwrap();
439 verify_inclusion(&leaves[index as usize], index, size, &hashes, &root).unwrap();
440 }
441 }
442 }
443
444 #[test]
445 fn prove_inclusion_path_fails_for_wrong_leaf() {
446 let leaves: Vec<MerkleHash> = (0..7u64)
447 .map(|i| hash_leaf(format!("leaf-{i}").as_bytes()))
448 .collect();
449 let root = compute_root(&leaves);
450 let hashes = prove_inclusion(&leaves, 3).unwrap();
451 let wrong = hash_leaf(b"not-in-the-tree");
452 assert!(verify_inclusion(&wrong, 3, 7, &hashes, &root).is_err());
453 }
454
455 #[test]
456 fn prove_inclusion_rejects_empty_tree_and_out_of_range() {
457 assert!(prove_inclusion(&[], 0).is_err());
458 let leaves = vec![hash_leaf(b"only")];
459 assert!(prove_inclusion(&leaves, 1).is_err());
460 }
461
462 #[test]
463 fn compute_root_empty() {
464 assert_eq!(compute_root(&[]), MerkleHash::EMPTY);
465 }
466
467 #[test]
468 fn compute_root_two_leaves() {
469 let a = hash_leaf(b"a");
470 let b = hash_leaf(b"b");
471 let root = compute_root(&[a, b]);
472 assert_eq!(root, hash_children(&a, &b));
473 }
474
475 #[test]
476 fn inclusion_proof_single_leaf() {
477 let leaf = hash_leaf(b"only leaf");
478 let root = leaf;
479 verify_inclusion(&leaf, 0, 1, &[], &root).unwrap();
480 }
481
482 #[test]
483 fn inclusion_proof_two_leaves() {
484 let a = hash_leaf(b"a");
485 let b = hash_leaf(b"b");
486 let root = hash_children(&a, &b);
487
488 verify_inclusion(&a, 0, 2, &[b], &root).unwrap();
490 verify_inclusion(&b, 1, 2, &[a], &root).unwrap();
492 }
493
494 #[test]
495 fn inclusion_proof_four_leaves() {
496 let leaves: Vec<MerkleHash> = (0..4u8).map(|i| hash_leaf(&[i])).collect();
497 let root = compute_root(&leaves);
498
499 let ab = hash_children(&leaves[0], &leaves[1]);
501 let cd = hash_children(&leaves[2], &leaves[3]);
502 let _ = hash_children(&ab, &cd);
503
504 verify_inclusion(&leaves[0], 0, 4, &[leaves[1], cd], &root).unwrap();
505 verify_inclusion(&leaves[1], 1, 4, &[leaves[0], cd], &root).unwrap();
506 verify_inclusion(&leaves[2], 2, 4, &[leaves[3], ab], &root).unwrap();
507 verify_inclusion(&leaves[3], 3, 4, &[leaves[2], ab], &root).unwrap();
508 }
509
510 #[test]
511 fn inclusion_proof_rejects_wrong_root() {
512 let a = hash_leaf(b"a");
513 let b = hash_leaf(b"b");
514 let _root = hash_children(&a, &b);
515 let wrong = MerkleHash::from_bytes([0xff; 32]);
516
517 let err = verify_inclusion(&a, 0, 2, &[b], &wrong);
518 assert!(err.is_err());
519 }
520
521 #[test]
522 fn inclusion_proof_three_leaves() {
523 let leaves: Vec<MerkleHash> = (0..3u8).map(|i| hash_leaf(&[i])).collect();
524 let root = compute_root(&leaves);
525
526 let ab = hash_children(&leaves[0], &leaves[1]);
527
528 verify_inclusion(&leaves[0], 0, 3, &[leaves[1], leaves[2]], &root).unwrap();
530 verify_inclusion(&leaves[1], 1, 3, &[leaves[0], leaves[2]], &root).unwrap();
532 verify_inclusion(&leaves[2], 2, 3, &[ab], &root).unwrap();
534 }
535
536 #[test]
537 fn inclusion_proof_five_leaves() {
538 let leaves: Vec<MerkleHash> = (0..5u8).map(|i| hash_leaf(&[i])).collect();
539 let root = compute_root(&leaves);
540
541 let h01 = hash_children(&leaves[0], &leaves[1]);
542 let h23 = hash_children(&leaves[2], &leaves[3]);
543 let h0123 = hash_children(&h01, &h23);
544
545 verify_inclusion(&leaves[4], 4, 5, &[h0123], &root).unwrap();
547 verify_inclusion(&leaves[0], 0, 5, &[leaves[1], h23, leaves[4]], &root).unwrap();
549 }
550
551 #[test]
552 fn inclusion_proof_seven_leaves() {
553 let leaves: Vec<MerkleHash> = (0..7u8).map(|i| hash_leaf(&[i])).collect();
554 let root = compute_root(&leaves);
555
556 let h01 = hash_children(&leaves[0], &leaves[1]);
557 let h23 = hash_children(&leaves[2], &leaves[3]);
558 let h45 = hash_children(&leaves[4], &leaves[5]);
559 let h0123 = hash_children(&h01, &h23);
560 let h456 = hash_children(&h45, &leaves[6]);
561
562 verify_inclusion(&leaves[6], 6, 7, &[h45, h0123], &root).unwrap();
564 verify_inclusion(&leaves[0], 0, 7, &[leaves[1], h23, h456], &root).unwrap();
566 }
567
568 #[test]
569 fn inclusion_proof_rejects_index_out_of_range() {
570 let a = hash_leaf(b"a");
571 let root = a;
572 let err = verify_inclusion(&a, 1, 1, &[], &root);
573 assert!(err.is_err());
574 }
575
576 #[test]
577 fn consistency_proof_same_size() {
578 let root = MerkleHash::from_bytes([0x42; 32]);
579 verify_consistency(5, 5, &[], &root, &root).unwrap();
580 }
581
582 #[test]
583 fn consistency_proof_empty_old() {
584 let new_root = MerkleHash::from_bytes([0x42; 32]);
585 let old_root = MerkleHash::EMPTY;
586 verify_consistency(0, 5, &[], &old_root, &new_root).unwrap();
587 }
588
589 #[test]
590 fn consistency_proof_2_to_4() {
591 let leaves: Vec<MerkleHash> = (0..4u8).map(|i| hash_leaf(&[i])).collect();
592 let old_root = compute_root(&leaves[..2]);
593 let new_root = compute_root(&leaves);
594 let proof = build_consistency_proof(&leaves[..2], &leaves);
595 verify_consistency(2, 4, &proof, &old_root, &new_root).unwrap();
596 }
597
598 #[test]
599 fn consistency_proof_3_to_5() {
600 let leaves: Vec<MerkleHash> = (0..5u8).map(|i| hash_leaf(&[i])).collect();
601 let old_root = compute_root(&leaves[..3]);
602 let new_root = compute_root(&leaves);
603 let proof = build_consistency_proof(&leaves[..3], &leaves);
604 verify_consistency(3, 5, &proof, &old_root, &new_root).unwrap();
605 }
606
607 #[test]
608 fn consistency_proof_4_to_8() {
609 let leaves: Vec<MerkleHash> = (0..8u8).map(|i| hash_leaf(&[i])).collect();
610 let old_root = compute_root(&leaves[..4]);
611 let new_root = compute_root(&leaves);
612 let proof = build_consistency_proof(&leaves[..4], &leaves);
613 verify_consistency(4, 8, &proof, &old_root, &new_root).unwrap();
614 }
615
616 #[test]
617 fn consistency_proof_7_to_15() {
618 let leaves: Vec<MerkleHash> = (0..15u8).map(|i| hash_leaf(&[i])).collect();
619 let old_root = compute_root(&leaves[..7]);
620 let new_root = compute_root(&leaves);
621 let proof = build_consistency_proof(&leaves[..7], &leaves);
622 verify_consistency(7, 15, &proof, &old_root, &new_root).unwrap();
623 }
624
625 #[test]
626 fn consistency_proof_1_to_4() {
627 let leaves: Vec<MerkleHash> = (0..4u8).map(|i| hash_leaf(&[i])).collect();
628 let old_root = compute_root(&leaves[..1]);
629 let new_root = compute_root(&leaves);
630 let proof = build_consistency_proof(&leaves[..1], &leaves);
631 verify_consistency(1, 4, &proof, &old_root, &new_root).unwrap();
632 }
633
634 #[test]
635 fn consistency_proof_rejects_wrong_old_root() {
636 let leaves: Vec<MerkleHash> = (0..4u8).map(|i| hash_leaf(&[i])).collect();
637 let wrong_old = MerkleHash::from_bytes([0xff; 32]);
638 let new_root = compute_root(&leaves);
639 let proof = build_consistency_proof(&leaves[..3], &leaves);
640 assert!(verify_consistency(3, 4, &proof, &wrong_old, &new_root).is_err());
641 }
642
643 #[test]
644 fn consistency_proof_rejects_wrong_new_root() {
645 let leaves: Vec<MerkleHash> = (0..4u8).map(|i| hash_leaf(&[i])).collect();
646 let old_root = compute_root(&leaves[..3]);
647 let wrong_new = MerkleHash::from_bytes([0xff; 32]);
648 let proof = build_consistency_proof(&leaves[..3], &leaves);
649 assert!(verify_consistency(3, 4, &proof, &old_root, &wrong_new).is_err());
650 }
651
652 fn build_consistency_proof(
654 old_leaves: &[MerkleHash],
655 new_leaves: &[MerkleHash],
656 ) -> Vec<MerkleHash> {
657 assert!(old_leaves.len() <= new_leaves.len());
658 subproof(old_leaves.len() as u64, new_leaves, true)
659 }
660
661 fn subproof(m: u64, leaves: &[MerkleHash], b: bool) -> Vec<MerkleHash> {
663 let n = leaves.len() as u64;
664 if m == n {
665 if b {
666 return vec![];
667 }
668 return vec![compute_root(leaves)];
669 }
670 let k = largest_power_of_2_lt(n) as usize;
671 if m <= k as u64 {
672 let mut proof = subproof(m, &leaves[..k], b);
673 proof.push(compute_root(&leaves[k..]));
674 proof
675 } else {
676 let mut proof = subproof(m - k as u64, &leaves[k..], false);
677 proof.push(compute_root(&leaves[..k]));
678 proof
679 }
680 }
681
682 #[test]
683 fn largest_pow2_lt() {
684 assert_eq!(largest_power_of_2_lt(2), 1);
685 assert_eq!(largest_power_of_2_lt(3), 2);
686 assert_eq!(largest_power_of_2_lt(4), 2);
687 assert_eq!(largest_power_of_2_lt(5), 4);
688 assert_eq!(largest_power_of_2_lt(8), 4);
689 assert_eq!(largest_power_of_2_lt(9), 8);
690 }
691}