Skip to content

Commit

Permalink
wip: impl Prover for ssz primitive types and add some tests
Browse files Browse the repository at this point in the history
  • Loading branch information
ralexstokes committed Mar 28, 2024
1 parent af475cc commit d37cb48
Show file tree
Hide file tree
Showing 4 changed files with 111 additions and 4 deletions.
4 changes: 3 additions & 1 deletion ssz-rs/src/boolean.rs
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
use crate::{
de::{Deserialize, DeserializeError},
lib::*,
merkleization::{multiproofs::Indexed, MerkleizationError, Merkleized, Node},
merkleization::{multiproofs::Indexed, MerkleizationError, Merkleized, Node, Prover},
ser::{Serialize, SerializeError},
Serializable, SimpleSerialize,
};
Expand Down Expand Up @@ -53,6 +53,8 @@ impl Merkleized for bool {
}
}

impl Prover for bool {}

impl SimpleSerialize for bool {}

impl Indexed for bool {
Expand Down
2 changes: 1 addition & 1 deletion ssz-rs/src/merkleization/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,7 @@ use crate::{
use sha2::{Digest, Sha256};

pub use node::Node;
pub use proofs::is_valid_merkle_branch;
pub use proofs::{is_valid_merkle_branch, Proof, ProofAndWitness, Prover};

pub(crate) const BYTES_PER_CHUNK: usize = 32;
pub(crate) const BITS_PER_CHUNK: usize = BYTES_PER_CHUNK * (crate::BITS_PER_BYTE as usize);
Expand Down
101 changes: 100 additions & 1 deletion ssz-rs/src/merkleization/proofs.rs
Original file line number Diff line number Diff line change
@@ -1,6 +1,70 @@
use crate::merkleization::{MerkleizationError as Error, Node};
use crate::{
merkleization::{MerkleizationError as Error, Node},
multiproofs::{self, default_generalized_index, GeneralizedIndex},
Merkleized,
};
use sha2::{Digest, Sha256};

#[derive(Debug)]
pub struct Proof {
pub leaf: Node,
pub branch: Vec<Node>,
pub index: GeneralizedIndex,
}

impl Proof {
pub fn verify(&self, root: Node) -> Result<(), Error> {
is_valid_merkle_branch_for_generalized_index(self.leaf, &self.branch, self.index, root)
}
}

impl Default for Proof {
fn default() -> Self {

Check warning on line 22 in ssz-rs/src/merkleization/proofs.rs

View check run for this annotation

Codecov / codecov/patch

ssz-rs/src/merkleization/proofs.rs#L22

Added line #L22 was not covered by tests
Self {
leaf: Default::default(),
branch: Default::default(),
index: default_generalized_index(),

Check warning on line 26 in ssz-rs/src/merkleization/proofs.rs

View check run for this annotation

Codecov / codecov/patch

ssz-rs/src/merkleization/proofs.rs#L24-L26

Added lines #L24 - L26 were not covered by tests
}
}
}

pub type ProofAndWitness = (Proof, Node);

pub fn prove_primitive<T: Prover + ?Sized>(
data: &mut T,
index: GeneralizedIndex,
) -> Result<ProofAndWitness, Error> {
if index != default_generalized_index() {
return Err(Error::InvalidGeneralizedIndex)

Check warning on line 38 in ssz-rs/src/merkleization/proofs.rs

View check run for this annotation

Codecov / codecov/patch

ssz-rs/src/merkleization/proofs.rs#L38

Added line #L38 was not covered by tests
}

let root = data.hash_tree_root()?;
let proof = Proof { leaf: root, branch: vec![], index };
Ok((proof, root))
}

pub trait Prover: Merkleized {
fn prove(&mut self, index: GeneralizedIndex) -> Result<ProofAndWitness, Error> {
prove_primitive(self, index)
}
}

pub fn get_subtree_index(i: GeneralizedIndex) -> Result<usize, Error> {
let i_log2 = multiproofs::log_2(i).ok_or(Error::InvalidGeneralizedIndex)?;
Ok(i % 2usize.pow(i_log2))
}

pub fn is_valid_merkle_branch_for_generalized_index<T: AsRef<[u8]>>(
leaf: Node,
branch: &[T],
generalized_index: GeneralizedIndex,
root: Node,
) -> Result<(), Error> {
let index = multiproofs::log_2(generalized_index).unwrap() as usize;
let depth = get_subtree_index(generalized_index).unwrap();
is_valid_merkle_branch(leaf, branch, depth, index, root)
}

/// `is_valid_merkle_branch` verifies the Merkle proof
/// against the `root` given the other metadata.
pub fn is_valid_merkle_branch<T: AsRef<[u8]>>(
Expand Down Expand Up @@ -39,6 +103,8 @@ pub fn is_valid_merkle_branch<T: AsRef<[u8]>>(

#[cfg(test)]
mod tests {
use crate::U256;

use super::*;

fn decode_node_from_hex(hex: &str) -> Node {
Expand Down Expand Up @@ -67,4 +133,37 @@ mod tests {

assert!(is_valid_merkle_branch(leaf, &branch, depth, index, root).is_ok());
}

#[test]
fn test_proving_primitives() {
let index = default_generalized_index();

let mut data = 8u8;
let (proof, witness) = data.prove(index).unwrap();
assert_eq!(witness, data.hash_tree_root().unwrap());
let result = proof.verify(witness);
assert!(result.is_ok());

let mut data = 234238u64;
let (proof, witness) = data.prove(index).unwrap();
assert_eq!(witness, data.hash_tree_root().unwrap());
let result = proof.verify(witness);
assert!(result.is_ok());

let mut data = U256::from_str_radix(
"f8c2ed25e9c31399d4149dcaa48c51f394043a6a1297e65780a5979e3d7bb77c",
16,
)
.unwrap();
let (proof, witness) = data.prove(index).unwrap();
assert_eq!(witness, data.hash_tree_root().unwrap());
let result = proof.verify(witness);
assert!(result.is_ok());

let mut data = true;
let (proof, witness) = data.prove(index).unwrap();
assert_eq!(witness, data.hash_tree_root().unwrap());
let result = proof.verify(witness);
assert!(result.is_ok())
}
}
8 changes: 7 additions & 1 deletion ssz-rs/src/uint.rs
Original file line number Diff line number Diff line change
@@ -1,7 +1,9 @@
use crate::{
de::{Deserialize, DeserializeError},
lib::*,
merkleization::{multiproofs::Indexed, pack_bytes, MerkleizationError, Merkleized, Node},
merkleization::{
multiproofs::Indexed, pack_bytes, MerkleizationError, Merkleized, Node, Prover,
},
ser::{Serialize, SerializeError},
Serializable, SimpleSerialize, BITS_PER_BYTE,
};
Expand Down Expand Up @@ -65,6 +67,8 @@ macro_rules! define_uint {
}
}

impl Prover for $uint {}

impl SimpleSerialize for $uint {}

impl Indexed for $uint {
Expand Down Expand Up @@ -136,6 +140,8 @@ impl Merkleized for U256 {
}
}

impl Prover for U256 {}

impl SimpleSerialize for U256 {}

impl Indexed for U256 {
Expand Down

0 comments on commit d37cb48

Please sign in to comment.