Skip to content

Commit e8c1980

Browse files
committed
feat(target_chains/starknet): add merkle tree utils
1 parent a1e4fc0 commit e8c1980

File tree

2 files changed

+71
-0
lines changed

2 files changed

+71
-0
lines changed

target_chains/starknet/contracts/src/lib.cairo

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,3 +3,4 @@ pub mod wormhole;
33
pub mod reader;
44
pub mod hash;
55
pub mod util;
6+
pub mod merkle_tree;
Lines changed: 70 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,70 @@
1+
use super::hash::{Hasher, HasherImpl};
2+
use super::reader::{Reader, ReaderImpl, ByteArray};
3+
use super::util::ONE_SHIFT_96;
4+
use core::cmp::{min, max};
5+
6+
const MERKLE_LEAF_PREFIX: u8 = 0;
7+
const MERKLE_NODE_PREFIX: u8 = 1;
8+
const MERKLE_EMPTY_LEAF_PREFIX: u8 = 2;
9+
10+
#[derive(Copy, Drop, Debug, Serde, PartialEq)]
11+
pub enum MerkleVerificationError {
12+
Reader: super::reader::Error,
13+
DigestMismatch,
14+
}
15+
16+
#[generate_trait]
17+
impl ResultReaderToMerkleVerification<T> of ResultReaderToMerkleVerificationTrait<T> {
18+
fn map_err(self: Result<T, pyth::reader::Error>) -> Result<T, MerkleVerificationError> {
19+
match self {
20+
Result::Ok(v) => Result::Ok(v),
21+
Result::Err(err) => Result::Err(MerkleVerificationError::Reader(err)),
22+
}
23+
}
24+
}
25+
26+
fn leaf_hash(mut reader: Reader) -> Result<u256, super::reader::Error> {
27+
let mut hasher = HasherImpl::new();
28+
hasher.push_u8(MERKLE_LEAF_PREFIX);
29+
hasher.push_reader(ref reader)?;
30+
let hash = hasher.finalize() / ONE_SHIFT_96;
31+
Result::Ok(hash)
32+
}
33+
34+
fn node_hash(a: u256, b: u256) -> u256 {
35+
let mut hasher = HasherImpl::new();
36+
hasher.push_u8(MERKLE_NODE_PREFIX);
37+
hasher.push_u160(min(a, b));
38+
hasher.push_u160(max(a, b));
39+
hasher.finalize() / ONE_SHIFT_96
40+
}
41+
42+
pub fn read_and_verify_proof(
43+
root_digest: u256, message: @ByteArray, ref reader: Reader
44+
) -> Result<(), MerkleVerificationError> {
45+
let mut message_reader = ReaderImpl::new(message.clone());
46+
let mut current_hash = leaf_hash(message_reader.clone()).map_err()?;
47+
48+
let proof_size = reader.read_u8().map_err()?;
49+
let mut i = 0;
50+
51+
let mut result = Result::Ok(());
52+
while i < proof_size {
53+
match reader.read_u160().map_err() {
54+
Result::Ok(sibling_digest) => {
55+
current_hash = node_hash(current_hash, sibling_digest);
56+
},
57+
Result::Err(err) => {
58+
result = Result::Err(err);
59+
break;
60+
},
61+
}
62+
i += 1;
63+
};
64+
result?;
65+
66+
if root_digest != current_hash {
67+
return Result::Err(MerkleVerificationError::DigestMismatch);
68+
}
69+
Result::Ok(())
70+
}

0 commit comments

Comments
 (0)