use sha2::{Digest, Sha256};
#[derive(Debug, thiserror::Error)]
pub enum MerkleAccumulatorError {
#[error("Invalid accumulator proof")]
InvalidProof,
#[error("Hash mismatch")]
HashMismatch,
#[error("Empty accumulator")]
EmptyAccumulator,
}
pub type MerkleAccumulatorResult<T> = Result<T, MerkleAccumulatorError>;
#[derive(Clone, Debug)]
pub struct Leaf {
pub hash: [u8; 32],
pub index: u64,
pub data: Vec<u8>,
}
#[derive(Clone, Debug)]
pub enum MerkleNode {
Leaf(Leaf),
Internal {
hash: [u8; 32],
left: Box<MerkleNode>,
right: Box<MerkleNode>,
},
Empty,
}
impl MerkleNode {
pub fn hash(&self) -> [u8; 32] {
match self {
MerkleNode::Leaf(leaf) => leaf.hash,
MerkleNode::Internal { hash, .. } => *hash,
MerkleNode::Empty => Self::empty_hash(),
}
}
pub fn empty_hash() -> [u8; 32] {
let digest = Sha256::digest(b"");
digest.into()
}
pub fn compute_internal_hash(left_hash: [u8; 32], right_hash: [u8; 32]) -> [u8; 32] {
let mut hasher = Sha256::new();
hasher.update(left_hash);
hasher.update(right_hash);
hasher.finalize().into()
}
}
#[derive(Clone, Debug)]
pub struct MerkleAccumulator {
root: MerkleNode,
num_leaves: u64,
}
impl Default for MerkleAccumulator {
fn default() -> Self {
Self {
root: MerkleNode::Empty,
num_leaves: 0,
}
}
}
impl MerkleAccumulator {
pub fn new() -> Self {
Self {
root: MerkleNode::Empty,
num_leaves: 0,
}
}
pub fn from_leaves(leaves: &[[u8; 32]]) -> Self {
if leaves.is_empty() {
return Self::new();
}
let leaf_nodes: Vec<MerkleNode> = leaves
.iter()
.enumerate()
.map(|(i, &hash)| {
MerkleNode::Leaf(Leaf {
hash,
index: i as u64,
data: Vec::new(),
})
})
.collect();
let root = Self::build_tree(&leaf_nodes);
Self {
root,
num_leaves: leaves.len() as u64,
}
}
fn build_tree(nodes: &[MerkleNode]) -> MerkleNode {
if nodes.is_empty() {
return MerkleNode::Empty;
}
if nodes.len() == 1 {
return nodes[0].clone();
}
let mut current_level = nodes.to_vec();
while current_level.len() > 1 {
let mut next_level = Vec::new();
for i in (0..current_level.len()).step_by(2) {
let left = current_level[i].clone();
let right = if i + 1 < current_level.len() {
current_level[i + 1].clone()
} else {
current_level[i].clone()
};
let left_hash = left.hash();
let right_hash = right.hash();
let internal_hash = MerkleNode::compute_internal_hash(left_hash, right_hash);
next_level.push(MerkleNode::Internal {
hash: internal_hash,
left: Box::new(left),
right: Box::new(right),
});
}
current_level = next_level;
}
current_level[0].clone()
}
pub fn root_hash(&self) -> [u8; 32] {
self.root.hash()
}
pub fn num_leaves(&self) -> u64 {
self.num_leaves
}
pub fn verify_proof(&self, leaf_hash: [u8; 32], index: u64, proof: &[MerkleProofItem]) -> bool {
if proof.is_empty() {
return self.root_hash() == leaf_hash && self.num_leaves == 1;
}
let computed_root = Self::compute_root_from_leaf(leaf_hash, index, proof);
computed_root == self.root_hash()
}
fn compute_root_from_leaf(
leaf_hash: [u8; 32],
index: u64,
proof: &[MerkleProofItem],
) -> [u8; 32] {
let mut current_hash = leaf_hash;
let mut _current_index = index;
for item in proof {
match item {
MerkleProofItem::Left { hash } => {
current_hash = MerkleNode::compute_internal_hash(*hash, current_hash);
}
MerkleProofItem::Right { hash } => {
current_hash = MerkleNode::compute_internal_hash(current_hash, *hash);
}
}
_current_index /= 2;
}
current_hash
}
}
#[derive(Clone, Debug)]
pub enum MerkleProofItem {
Left { hash: [u8; 32] },
Right { hash: [u8; 32] },
}
#[derive(Clone, Debug)]
pub struct StateProof {
pub address: [u8; 32],
pub resource_type: String,
pub exists: bool,
pub data: Option<Vec<u8>>,
pub accumulator_proof: Vec<MerkleProofItem>,
pub version: u64,
pub leaf_hash: [u8; 32],
}
impl StateProof {
pub fn new(
address: [u8; 32],
resource_type: String,
exists: bool,
data: Option<Vec<u8>>,
accumulator_proof: Vec<MerkleProofItem>,
version: u64,
leaf_hash: [u8; 32],
) -> Self {
Self {
address,
resource_type,
exists,
data,
accumulator_proof,
version,
leaf_hash,
}
}
pub fn compute_leaf_hash(&self) -> [u8; 32] {
let mut hasher = Sha256::new();
hasher.update(b"APTOS::STATE::LEAF");
hasher.update(self.address);
hasher.update(self.resource_type.as_bytes());
if self.exists {
hasher.update(b"EXISTS");
if let Some(data) = &self.data {
hasher.update(data);
}
} else {
hasher.update(b"NOT_EXISTS");
}
hasher.finalize().into()
}
pub fn verify(&self, expected_root: [u8; 32]) -> bool {
let leaf_hash = self.compute_leaf_hash();
if leaf_hash != self.leaf_hash {
return false;
}
let computed_root = MerkleAccumulator::compute_root_from_leaf(
leaf_hash,
0, &self.accumulator_proof,
);
computed_root == expected_root
}
}
#[derive(Clone, Debug)]
pub struct TransactionProof {
pub version: u64,
pub hash: [u8; 32],
pub state_change_hash: [u8; 32],
pub event_root_hash: [u8; 32],
pub state_checkpoint_hash: Option<[u8; 32]>,
pub epoch: u64,
pub round: u64,
pub ledger_proof: Vec<MerkleProofItem>,
}
impl TransactionProof {
#[allow(clippy::too_many_arguments)]
pub fn new(
version: u64,
hash: [u8; 32],
state_change_hash: [u8; 32],
event_root_hash: [u8; 32],
state_checkpoint_hash: Option<[u8; 32]>,
epoch: u64,
round: u64,
ledger_proof: Vec<MerkleProofItem>,
) -> Self {
Self {
version,
hash,
state_change_hash,
event_root_hash,
state_checkpoint_hash,
epoch,
round,
ledger_proof,
}
}
}
#[derive(Clone, Debug)]
pub struct EventProof {
pub hash: [u8; 32],
pub index: u64,
pub event_proof: Vec<MerkleProofItem>,
}
impl EventProof {
pub fn new(hash: [u8; 32], index: u64, event_proof: Vec<MerkleProofItem>) -> Self {
Self {
hash,
index,
event_proof,
}
}
pub fn verify(&self, event_root: [u8; 32]) -> bool {
let computed_root =
MerkleAccumulator::compute_root_from_leaf(self.hash, self.index, &self.event_proof);
computed_root == event_root
}
}
#[derive(Clone, Debug)]
pub struct LedgerProof {
pub version: u64,
pub root_hash: [u8; 32],
pub chain_id: u64,
pub epoch: u64,
pub proof: Vec<MerkleProofItem>,
}
impl LedgerProof {
pub fn new(
version: u64,
root_hash: [u8; 32],
chain_id: u64,
epoch: u64,
proof: Vec<MerkleProofItem>,
) -> Self {
Self {
version,
root_hash,
chain_id,
epoch,
proof,
}
}
pub fn verify(&self) -> bool {
self.root_hash != [0u8; 32]
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_merkle_accumulator_empty() {
let acc = MerkleAccumulator::new();
assert_eq!(acc.num_leaves(), 0);
assert_ne!(acc.root_hash(), [0u8; 32]); }
#[test]
fn test_merkle_accumulator_single_leaf() {
let leaves = [[1u8; 32]];
let acc = MerkleAccumulator::from_leaves(&leaves);
assert_eq!(acc.num_leaves(), 1);
assert_eq!(acc.root_hash(), leaves[0]);
}
#[test]
fn test_merkle_accumulator_two_leaves() {
let leaves = [[1u8; 32], [2u8; 32]];
let acc = MerkleAccumulator::from_leaves(&leaves);
assert_eq!(acc.num_leaves(), 2);
assert_ne!(acc.root_hash(), [0u8; 32]);
}
#[test]
fn test_merkle_node_compute_internal_hash() {
let left = [1u8; 32];
let right = [2u8; 32];
let hash = MerkleNode::compute_internal_hash(left, right);
assert_ne!(hash, [0u8; 32]);
}
#[test]
fn test_state_proof_leaf_hash() {
let proof = StateProof::new(
[1u8; 32],
"CSV::Seal".to_string(),
true,
Some(vec![1, 2, 3]),
vec![],
100,
[0u8; 32],
);
let hash = proof.compute_leaf_hash();
assert_eq!(hash.len(), 32);
}
#[test]
fn test_ledger_proof_verification() {
let proof = LedgerProof::new(100, [1u8; 32], 1, 1, vec![]);
assert!(proof.verify());
}
}