|
| 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