From b6759bac6ff1b40e7327a8b3dd3ee46a98292649 Mon Sep 17 00:00:00 2001 From: Uncle Stinky Date: Sat, 3 Oct 2026 20:01:25 +0300 Subject: [PATCH] optimize merkle tree prefixed computations Signed-off-by: Uncle Stinky --- pallets/exodus/src/benchmarking.rs | 20 +- pallets/exodus/src/error.rs | 2 + pallets/exodus/src/impls/dkg.rs | 18 +- pallets/exodus/src/impls/frost.rs | 34 +- pallets/exodus/src/impls/merkle.rs | 44 +- pallets/exodus/src/impls/mod.rs | 3 + pallets/exodus/src/impls/sizes.rs | 29 + pallets/exodus/src/lib.rs | 1 + pallets/exodus/src/tests.rs | 99 ++- pallets/governance/src/tests.rs | 53 +- pallets/helpers/src/hash_chain.rs | 6 +- pallets/helpers/src/lib.rs | 7 + pallets/helpers/src/merkle_tree.rs | 1095 ++++++++++++++++------------ pallets/traits/src/exodus.rs | 17 +- pallets/traits/src/hashing.rs | 6 + 15 files changed, 867 insertions(+), 567 deletions(-) create mode 100644 pallets/exodus/src/impls/sizes.rs diff --git a/pallets/exodus/src/benchmarking.rs b/pallets/exodus/src/benchmarking.rs index 2da947c..d9f2ca6 100644 --- a/pallets/exodus/src/benchmarking.rs +++ b/pallets/exodus/src/benchmarking.rs @@ -415,10 +415,16 @@ mod benchmarks { dummy_values.insert(identifier, dummy_value.clone()); } + let identifier_size = network_curve.scalar_bytes_len(); + let merkle_tree = NetworkCurve::build_merkle_tree( - &dummy_values, - |value: &Vec| Ok(value.clone()), + identifier_size + dummy_value.len(), + &dummy_values, + |value, buffer| { + buffer.extend_from_slice(value.as_ref()); + Ok(()) + } ).expect("Merkle tree for verifying shares should be valid"); let merkle_root = merkle_tree.last().copied().unwrap(); @@ -697,10 +703,16 @@ mod benchmarks { dummy_values.insert(identifier, dummy_binding_factor.clone()); } + let identifier_size = network_curve.scalar_bytes_len(); + let merkle_tree = NetworkCurve::build_merkle_tree( - &dummy_values, - |value: &Vec| Ok(value.clone()), + identifier_size + dummy_scalar.len(), + &dummy_values, + |value, buffer| { + buffer.extend_from_slice(value.as_ref()); + Ok(()) + } ).expect("Merkle tree for binding factor should be valid"); let merkle_root = merkle_tree.last().copied().unwrap(); diff --git a/pallets/exodus/src/error.rs b/pallets/exodus/src/error.rs index 3c4b18b..8e51b3d 100644 --- a/pallets/exodus/src/error.rs +++ b/pallets/exodus/src/error.rs @@ -43,6 +43,7 @@ pub enum ExodusError { InvalidMerkleProof, InvalidIdentityElement, NoActiveAuthorities, + TreeGenerationFailed, Unknown, } @@ -88,6 +89,7 @@ impl core::fmt::Debug for ExodusError { ExodusError::InvalidMerkleProof => write!(fmt, "Invalid Merkle proof provided."), ExodusError::InvalidIdentityElement => write!(fmt, "Invalid identity element provided."), ExodusError::NoActiveAuthorities => write!(fmt, "No active authorities for signing EXODUS transactions."), + ExodusError::TreeGenerationFailed => write!(fmt, "Tree generation failed because of inconsistent indexation."), ExodusError::Unknown => write!(fmt, "Unknown error."), } } diff --git a/pallets/exodus/src/impls/dkg.rs b/pallets/exodus/src/impls/dkg.rs index 0cef728..e50d0e1 100644 --- a/pallets/exodus/src/impls/dkg.rs +++ b/pallets/exodus/src/impls/dkg.rs @@ -3,6 +3,7 @@ use rand_chacha::rand_core::{CryptoRng, RngCore}; use ghost_traits::exodus::{ DistributedKeyGeneration, IdentifierConverter, MerkleTreeBuilder, + CiphersuiteSizes, }; use crate::{AuthIndex, ExodusError, NetworkCurve}; @@ -241,6 +242,8 @@ impl DistributedKeyGeneration for NetworkCurve { round2_packages_bytes: &Self::RoundPackages, ) -> Self::Part3Result { with_ciphersuite!(self, |f| { + type InnerGroup = ::Group; + let round2_secret_package = f::keys::dkg::round2::SecretPackage::deserialize(round2_secret_package_bytes) .map_err(|_| ExodusError::DeserializationError)?; @@ -291,13 +294,24 @@ impl DistributedKeyGeneration for NetworkCurve { let key_package_bytes = key_package.serialize() .map_err(|_| ExodusError::SerializationError)?; + let identifier_size = self.scalar_bytes_len(); + let verifying_share_size = self.element_bytes_len(); + let verifying_shares_merkle_tree = Self::build_merkle_tree( + identifier_size + verifying_share_size, pubkey_package.verifying_shares(), - |verifying_share| verifying_share.serialize() - .map_err(|_| ExodusError::SerializationError) + |verifying_share, buffer| { + let serialized = ::serialize( + &verifying_share.to_element() + ).map_err(|_| ExodusError::SerializationError)?; + + buffer.extend_from_slice(serialized.as_ref()); + Ok(()) + } )?; + // TODO: make it better let (_, merkle_proof) = Self::generate_merkle_proof_from_tree( pubkey_package.verifying_shares(), diff --git a/pallets/exodus/src/impls/frost.rs b/pallets/exodus/src/impls/frost.rs index 3e3423a..67449ee 100644 --- a/pallets/exodus/src/impls/frost.rs +++ b/pallets/exodus/src/impls/frost.rs @@ -1,5 +1,5 @@ use ghost_traits::exodus::{ - IdentifierConverter, MerkleTreeBuilder, + CiphersuiteSizes, IdentifierConverter, MerkleTreeBuilder, FlexibleRoundOptimizedSchnorrThresholdSignature, }; use sp_std::{ @@ -15,29 +15,6 @@ impl FlexibleRoundOptimizedSchnorrThresholdSignature for type Packages<'a> = BTreeMap; type NoncePackages<'a> = BTreeMap; - fn header_bytes_len(&self) -> usize { 5 } - - fn signature_bytes_len(&self) -> usize { - match self { - NetworkCurve::Secp256k1 => 64, - NetworkCurve::Ed25519 => 64, - } - } - - fn element_bytes_len(&self) -> usize { - match self { - NetworkCurve::Secp256k1 => 33, - NetworkCurve::Ed25519 => 32, - } - } - - fn scalar_bytes_len(&self) -> usize { - match self { - NetworkCurve::Secp256k1 => 32, - NetworkCurve::Ed25519 => 32, - } - } - fn verify_signature_share( &self, index: AuthIndex, @@ -330,11 +307,16 @@ impl FlexibleRoundOptimizedSchnorrThresholdSignature for return Err(ExodusError::InvalidParticipantId); } + let identifier_size = self.scalar_bytes_len(); + let binding_factor_size = self.scalar_bytes_len(); + let binding_factor_tree = Self::build_merkle_tree( + identifier_size + binding_factor_size, &binding_factors_list, - |binding_factor| { + |binding_factor, buffer| { let serialized = ::serialize(binding_factor); - Ok(serialized.as_ref().to_vec()) + buffer.extend_from_slice(serialized.as_ref()); + Ok(()) } )?; diff --git a/pallets/exodus/src/impls/merkle.rs b/pallets/exodus/src/impls/merkle.rs index 009c5ed..54b252f 100644 --- a/pallets/exodus/src/impls/merkle.rs +++ b/pallets/exodus/src/impls/merkle.rs @@ -13,15 +13,16 @@ impl MerkleTreeBuilder for NetworkCurve { type Hash = ::Hash; fn build_merkle_tree( + size_hint: usize, values: &BTreeMap, V>, - serialize_fn: F, + update_preimage_buffer: F, ) -> Result, ExodusError> where C: frost_core::Ciphersuite, <::Group as Group>::Field: frost_core::Field, - F: Fn(&V) -> Result, ExodusError> + F: Fn(&V, &mut Vec) -> Result<(), ExodusError> { - use ghost_helpers::merkle_tree::generate_tree; + use ghost_helpers::merkle_tree::{generate_tree, MerkleTreeError}; type CField = <::Group as Group>::Field; @@ -34,25 +35,28 @@ impl MerkleTreeBuilder for NetworkCurve { }) .ok_or(ExodusError::IncorrectNumberOfIdentifiers)?; - generate_tree::( + let maybe_tree = generate_tree::( max_index, + size_hint, values.iter(), - |(identifier, value)| { + |(identifier, value), buffer| { let id_scalar = identifier.to_scalar(); let id_bytes = as frost_core::Field>::serialize(&id_scalar); let index = Self::convert_identifier_to_index(id_bytes.as_ref())?; - let value_bytes = serialize_fn(value).map_err(|_| ExodusError::SerializationError)?; - // NOTE: revisit - // Is it possible to use [u8; CONSTANT] where constant is - // purely dependant on provided generic. - let mut preimage = Vec::::with_capacity(id_bytes.as_ref().len() + value_bytes.len()); - preimage.extend_from_slice(id_bytes.as_ref()); - preimage.extend_from_slice(value_bytes.as_ref()); + buffer.extend_from_slice(id_bytes.as_ref()); + update_preimage_buffer(value, buffer)?; - Ok((index as usize, preimage)) - }) + Ok(index as usize) + }); + + match maybe_tree { + Ok(tree) => Ok(tree), + Err(MerkleTreeError::Generate(exodus_err)) => Err(exodus_err), + Err(MerkleTreeError::InvalidIndex) => Err(ExodusError::TreeGenerationFailed), + Err(MerkleTreeError::Empty) => Err(ExodusError::IncorrectNumberOfIdentifiers), + } } fn generate_merkle_proof_from_tree( @@ -66,7 +70,7 @@ impl MerkleTreeBuilder for NetworkCurve { <::Group as Group>::Field: frost_core::Field, F: Fn(&V) -> Result, ExodusError> { - use ghost_helpers::merkle_tree::generate_proof; + use ghost_helpers::merkle_tree::{generate_proof, MerkleTreeError}; type CField = <::Group as Group>::Field; @@ -75,7 +79,7 @@ impl MerkleTreeBuilder for NetworkCurve { .map_err(|_| ExodusError::InvalidParticipantId)?; let value = values.get(&identifier).ok_or(ExodusError::InvalidParticipantId)?; - let value_bytes = serialize_fn(&value).map_err(|_| ExodusError::SerializationError)?; + let value_bytes = serialize_fn(&value)?; let max_index = values.keys() .next_back() @@ -86,11 +90,15 @@ impl MerkleTreeBuilder for NetworkCurve { }) .ok_or(ExodusError::IncorrectNumberOfIdentifiers)?; - let proof = generate_proof::( + let proof = generate_proof::( merkle_tree, max_index, authority_index, - ); + ).map_err(|err| match err { + MerkleTreeError::Generate(exodus_err) => exodus_err, + MerkleTreeError::InvalidIndex => ExodusError::TreeGenerationFailed, + MerkleTreeError::Empty => ExodusError::IncorrectNumberOfIdentifiers, + })?; Ok((value_bytes, proof)) } diff --git a/pallets/exodus/src/impls/mod.rs b/pallets/exodus/src/impls/mod.rs index 047e986..8b1f9b0 100644 --- a/pallets/exodus/src/impls/mod.rs +++ b/pallets/exodus/src/impls/mod.rs @@ -72,3 +72,6 @@ mod ecdh; #[macro_use] mod frost; + +#[macro_use] +mod sizes; diff --git a/pallets/exodus/src/impls/sizes.rs b/pallets/exodus/src/impls/sizes.rs new file mode 100644 index 0000000..5787055 --- /dev/null +++ b/pallets/exodus/src/impls/sizes.rs @@ -0,0 +1,29 @@ +use ghost_helpers::networks::NetworkCurve; +use ghost_traits::exodus::CiphersuiteSizes; + +use crate::{AuthIndex, ExodusError}; + +impl CiphersuiteSizes for NetworkCurve { + fn header_bytes_len(&self) -> usize { 5 } + + fn signature_bytes_len(&self) -> usize { + match self { + NetworkCurve::Secp256k1 => 64, + NetworkCurve::Ed25519 => 64, + } + } + + fn element_bytes_len(&self) -> usize { + match self { + NetworkCurve::Secp256k1 => 33, + NetworkCurve::Ed25519 => 32, + } + } + + fn scalar_bytes_len(&self) -> usize { + match self { + NetworkCurve::Secp256k1 => 32, + NetworkCurve::Ed25519 => 32, + } + } +} diff --git a/pallets/exodus/src/lib.rs b/pallets/exodus/src/lib.rs index 7f43c6b..0169079 100644 --- a/pallets/exodus/src/lib.rs +++ b/pallets/exodus/src/lib.rs @@ -48,6 +48,7 @@ use ghost_traits::{ exodus::{ MerkleTreeBuilder, DistributedKeyGeneration, EllipticCurveDiffieHellman, FlexibleRoundOptimizedSchnorrThresholdSignature, EvmGovernanceRegistrar, + CiphersuiteSizes, }, networks::{ NetworkDataBasicHandler, NetworkDataInspectHandler, diff --git a/pallets/exodus/src/tests.rs b/pallets/exodus/src/tests.rs index b54870d..c52e9a4 100644 --- a/pallets/exodus/src/tests.rs +++ b/pallets/exodus/src/tests.rs @@ -1571,9 +1571,17 @@ fn test_merkle_tree_success_power_of_two_for_shares() { let indices = (1..=32).collect::>(); let shares = generate_mock_secp256k1_shares(&indices); + let identifier_len = network_curve.scalar_bytes_len(); + let secret_share_len = network_curve.element_bytes_len(); + let tree = NetworkCurve::build_merkle_tree::( + identifier_len + secret_share_len, &shares, - |share| share.serialize().map_err(|_| ExodusError::SerializationError), + |share, buffer| { + let serialized = share.serialize().map_err(|_| ExodusError::SerializationError)?; + buffer.extend_from_slice(serialized.as_ref()); + Ok(()) + } ).unwrap(); let root = *tree.last().unwrap(); @@ -1608,11 +1616,16 @@ fn test_merkle_tree_success_power_of_two_for_scalar() { let indices = (1..=32).collect::>(); let scalars = generate_mock_secp256k1_scalars(&indices); + let identifier_len = network_curve.scalar_bytes_len(); + let secret_scalar_len = network_curve.scalar_bytes_len(); + let tree = NetworkCurve::build_merkle_tree::( + identifier_len + secret_scalar_len, &scalars, - |scalar| { + |scalar, buffer| { let serialized = ::serialize(scalar); - Ok(serialized.as_ref().to_vec()) + buffer.extend_from_slice(serialized.as_ref()); + Ok(()) } ).unwrap(); @@ -1623,7 +1636,7 @@ fn test_merkle_tree_success_power_of_two_for_scalar() { &identifier.serialize(), ).unwrap(); - let (shares_bytes, proof) = + let (scalar_bytes, proof) = NetworkCurve::generate_merkle_proof_from_tree::( &scalars, &tree, @@ -1635,7 +1648,7 @@ fn test_merkle_tree_success_power_of_two_for_scalar() { ).unwrap(); let result = network_curve.verify_merkle_proof( - &shares_bytes, + &scalar_bytes, &proof, root, auth_index, @@ -1657,9 +1670,17 @@ fn test_random_sparse_combinations_stress_for_shares() { let indices: Vec = unique_indices.into_iter().collect(); let shares = generate_mock_secp256k1_shares(&indices); + let identifier_len = network_curve.scalar_bytes_len(); + let secret_share_len = network_curve.element_bytes_len(); + let tree = NetworkCurve::build_merkle_tree::( + identifier_len + secret_share_len, &shares, - |share| share.serialize().map_err(|_| ExodusError::SerializationError), + |share, buffer| { + let serialized = share.serialize().map_err(|_| ExodusError::SerializationError)?; + buffer.extend_from_slice(serialized.as_ref()); + Ok(()) + } ).unwrap(); let root = *tree.last().unwrap(); @@ -1695,11 +1716,16 @@ fn test_random_sparse_combinations_stress_for_scalars() { let indices: Vec = unique_indices.into_iter().collect(); let scalars = generate_mock_secp256k1_scalars(&indices); + let identifier_len = network_curve.scalar_bytes_len(); + let secret_scalar_len = network_curve.scalar_bytes_len(); + let tree = NetworkCurve::build_merkle_tree::( + identifier_len + secret_scalar_len, &scalars, - |scalar| { + |scalar, buffer| { let serialized = ::serialize(scalar); - Ok(serialized.as_ref().to_vec()) + buffer.extend_from_slice(serialized.as_ref()); + Ok(()) } ).unwrap(); @@ -1733,9 +1759,17 @@ fn test_participants_with_high_index_gap_for_shares() { let indices = vec![1, 1023]; let shares = generate_mock_secp256k1_shares(&indices); + let identifier_len = network_curve.scalar_bytes_len(); + let secret_share_len = network_curve.element_bytes_len(); + let tree = NetworkCurve::build_merkle_tree::( + identifier_len + secret_share_len, &shares, - |share| share.serialize().map_err(|_| ExodusError::SerializationError), + |share, buffer| { + let serialized = share.serialize().map_err(|_| ExodusError::SerializationError)?; + buffer.extend_from_slice(serialized.as_ref()); + Ok(()) + } ).unwrap(); let root = *tree.last().unwrap(); @@ -1765,11 +1799,16 @@ fn test_participants_with_high_index_gap_for_scalars() { let indices = vec![1, 1023]; let scalars = generate_mock_secp256k1_scalars(&indices); + let identifier_len = network_curve.scalar_bytes_len(); + let secret_scalar_len = network_curve.scalar_bytes_len(); + let tree = NetworkCurve::build_merkle_tree::( + identifier_len + secret_scalar_len, &scalars, - |scalar| { + |scalar, buffer| { let serialized = ::serialize(scalar); - Ok(serialized.as_ref().to_vec()) + buffer.extend_from_slice(serialized.as_ref()); + Ok(()) } ).unwrap(); @@ -1803,9 +1842,17 @@ fn test_regression_consecutive_vs_position_for_shares() { let custom_indices = vec![2, 3]; let shares = generate_mock_secp256k1_shares(&custom_indices); + let identifier_len = network_curve.scalar_bytes_len(); + let secret_share_len = network_curve.element_bytes_len(); + let tree = NetworkCurve::build_merkle_tree::( + identifier_len + secret_share_len, &shares, - |scalar| scalar.serialize().map_err(|_| ExodusError::SerializationError), + |share, buffer| { + let serialized = share.serialize().map_err(|_| ExodusError::SerializationError)?; + buffer.extend_from_slice(serialized.as_ref()); + Ok(()) + } ).unwrap(); let root = *tree.last().unwrap(); @@ -1828,11 +1875,16 @@ fn test_regression_consecutive_vs_position_for_scalars() { let custom_indices = vec![2, 3]; let scalars = generate_mock_secp256k1_scalars(&custom_indices); + let identifier_len = network_curve.scalar_bytes_len(); + let secret_scalar_len = network_curve.scalar_bytes_len(); + let tree = NetworkCurve::build_merkle_tree::( + identifier_len + secret_scalar_len, &scalars, - |scalar| { + |scalar, buffer| { let serialized = ::serialize(scalar); - Ok(serialized.as_ref().to_vec()) + buffer.extend_from_slice(serialized.as_ref()); + Ok(()) } ).unwrap(); @@ -1859,9 +1911,17 @@ fn test_attack_with_unknown_authority_index_for_shares() { let custom_indices = vec![1, 2, 3]; let shares = generate_mock_secp256k1_shares(&custom_indices); + let identifier_len = network_curve.scalar_bytes_len(); + let secret_share_len = network_curve.element_bytes_len(); + let tree = NetworkCurve::build_merkle_tree::( + identifier_len + secret_share_len, &shares, - |scalar| scalar.serialize().map_err(|_| ExodusError::SerializationError), + |share, buffer| { + let serialized = share.serialize().map_err(|_| ExodusError::SerializationError)?; + buffer.extend_from_slice(serialized.as_ref()); + Ok(()) + } ).unwrap(); let root = *tree.last().unwrap(); @@ -1884,11 +1944,16 @@ fn test_attack_with_unknown_authority_index_for_scalars() { let custom_indices = vec![1, 2, 3]; let scalars = generate_mock_secp256k1_scalars(&custom_indices); + let identifier_len = network_curve.scalar_bytes_len(); + let secret_scalar_len = network_curve.scalar_bytes_len(); + let tree = NetworkCurve::build_merkle_tree::( + identifier_len + secret_scalar_len, &scalars, - |scalar| { + |scalar, buffer| { let serialized = ::serialize(scalar); - Ok(serialized.as_ref().to_vec()) + buffer.extend_from_slice(serialized.as_ref()); + Ok(()) } ).unwrap(); diff --git a/pallets/governance/src/tests.rs b/pallets/governance/src/tests.rs index df37f4f..1d6d7a7 100644 --- a/pallets/governance/src/tests.rs +++ b/pallets/governance/src/tests.rs @@ -79,14 +79,21 @@ fn claim_happy_path() { let users = vec![user_1.clone(), user_2.clone(), user_3.clone()]; let max_index = 2u8; + let tmp_package_probe = ClaimPackage:: { + shares: user_1.allowed_shares, + merkle_proof: Default::default(), + token_id: user_1.token_id, + index: 0u32, + }; + let size_hint = tmp_package_probe.get_preimage(&user_1.evm_address, network_id).as_ref().len(); + let raw_values_complex = vec![(0u8, &user_1), (1u8, &user_2), (2u8, &user_3)]; - let merkle_tree = generate_tree::( + let merkle_tree = generate_tree::( max_index, + size_hint, raw_values_complex, - |item: (u8, &NewTestUser)| -> Result<(usize, Vec), ()> { - let (index, u) = item; - + |(index, u), buffer| { let tmp_package = ClaimPackage:: { shares: u.allowed_shares, merkle_proof: Default::default(), @@ -94,9 +101,12 @@ fn claim_happy_path() { index: index as u32 }; - Ok((index as usize, tmp_package.get_preimage(&u.evm_address, network_id).to_vec())) + let preimage_bytes = tmp_package.get_preimage(&u.evm_address, network_id); + buffer.extend_from_slice(preimage_bytes.as_ref()); + + Ok(index as usize) } - ).unwrap(); + ).ok().unwrap(); let merkle_root = *merkle_tree.last().unwrap(); @@ -116,7 +126,9 @@ fn claim_happy_path() { activated_shares: 0, }); - let proof = generate_proof::(&merkle_tree, max_index, slot_index as u8); + let proof = generate_proof::(&merkle_tree, max_index, slot_index as u8) + .ok() + .unwrap(); let claim_package = ClaimPackage:: { shares: current_user.allowed_shares, @@ -205,21 +217,33 @@ fn partial_claims_happy_path() { let max_index = 0u8; let raw_values_complex = vec![(0u8, ¤t_user)]; - let merkle_tree = generate_tree::( + let tmp_package_probe = ClaimPackage:: { + shares: current_user.allowed_shares, + merkle_proof: Default::default(), + token_id: current_user.token_id, + index: 0u32, + }; + let size_hint = tmp_package_probe.get_preimage(¤t_user.evm_address, network_id).as_ref().len(); + + let merkle_tree = generate_tree::( max_index, + size_hint, raw_values_complex, - |item: (u8, &NewTestUser)| -> Result<(usize, Vec), ()> { - let (index, u) = item; + |(index, u), buffer| { let tmp_package = ClaimPackage:: { shares: u.allowed_shares, merkle_proof: Default::default(), token_id: u.token_id, index: index as u32 }; - Ok((index as usize, tmp_package.get_preimage(&u.evm_address, network_id).to_vec())) + + let preimage_bytes = tmp_package.get_preimage(&u.evm_address, network_id); + buffer.extend_from_slice(preimage_bytes.as_ref()); + + Ok(index as usize) } - ).unwrap(); + ).ok().unwrap(); let total_shares = 1_000_000; let merkle_root = *merkle_tree.last().unwrap(); @@ -240,7 +264,10 @@ fn partial_claims_happy_path() { setup_governor(&substrate_account_b); setup_governor(&substrate_account_c); - let proof = generate_proof::(&merkle_tree, max_index, 0u8); + let proof = generate_proof::(&merkle_tree, max_index, 0u8) + .ok() + .unwrap(); + let claim_package = ClaimPackage { shares: current_user.allowed_shares, merkle_proof: BoundedVec::try_from(proof).unwrap(), diff --git a/pallets/helpers/src/hash_chain.rs b/pallets/helpers/src/hash_chain.rs index 0f6ecbd..b3d4b57 100644 --- a/pallets/helpers/src/hash_chain.rs +++ b/pallets/helpers/src/hash_chain.rs @@ -188,9 +188,9 @@ mod tests { } } - run_hasher_tests!(SubstrateBlake2Hasher, blake2_tests); - run_hasher_tests!(SubstrateKeccakHasher, keccak_tests); - run_hasher_tests!(UtxoSha2Hasher, sha2_tests); + run_hasher_tests!(SubstrateBlake2Hasher, blake2_hashing_tests); + run_hasher_tests!(SubstrateKeccakHasher, keccak_hashing_tests); + run_hasher_tests!(UtxoSha2Hasher, sha2_hashing_tests); #[test] fn test_verify_in_current_block_with_empty_ancestry() { diff --git a/pallets/helpers/src/lib.rs b/pallets/helpers/src/lib.rs index f871d87..2d6fb94 100644 --- a/pallets/helpers/src/lib.rs +++ b/pallets/helpers/src/lib.rs @@ -79,6 +79,13 @@ macro_rules! impl_ghost_hasher { combined[$hash_bytes_len..].copy_from_slice(right.as_ref()); Self::hash(&combined) } + + fn combine_hashes_prefixed(p: u8, l: Self::Hash, r: Self::Hash) -> Self::Hash { + let mut combined = [p; $hash_bytes_len * 2 + 1]; + combined[1..1 + $hash_bytes_len].copy_from_slice(l.as_ref()); + combined[1 + $hash_bytes_len..].copy_from_slice(r.as_ref()); + Self::hash(&combined) + } } }; } diff --git a/pallets/helpers/src/merkle_tree.rs b/pallets/helpers/src/merkle_tree.rs index cb109ff..7b71c68 100644 --- a/pallets/helpers/src/merkle_tree.rs +++ b/pallets/helpers/src/merkle_tree.rs @@ -6,49 +6,57 @@ use sp_std::vec::Vec; const LEAF_PREFIX: u8 = 0; const NODE_PREFIX: u8 = 1; +pub enum MerkleTreeError { + Generate(E), + InvalidIndex, + Empty, +} + fn leaf_prefixed(preimage: &[u8]) -> H::Hash where H: GhostHasher, { - let mut prefixed_data = sp_std::vec![LEAF_PREFIX; 1 + preimage.len()]; - prefixed_data[1..].copy_from_slice(preimage); - H::hash(&prefixed_data) -} + let mut buffer = Vec::with_capacity(preimage.len() + 1); + buffer.push(LEAF_PREFIX); + buffer.extend_from_slice(preimage); -fn node_prefixed(preimage1: &[u8], preimage2: &[u8]) -> H::Hash -where - H: GhostHasher, -{ - let total_len = preimage1.len() + preimage2.len(); - let mut prefixed_data = sp_std::vec![NODE_PREFIX; 1 + total_len]; - prefixed_data[1..1 + preimage1.len()].copy_from_slice(preimage1); - prefixed_data[1 + preimage1.len()..].copy_from_slice(preimage2); - H::hash(&prefixed_data) + H::hash(&buffer) } pub fn generate_tree( max_index: I, + size_hint: usize, raw_values: S, - generate_preimage: F, -) -> Result, E> + update_preimage_buffer: F, +) -> Result, MerkleTreeError> where H: GhostHasher, S: IntoIterator, I: AtLeast8BitUnsigned + UniqueSaturatedInto, - F: Fn(S::Item) -> Result<(usize, Vec), E>, + F: Fn(S::Item, &mut Vec) -> Result, { let num_of_leaves: usize = (max_index + I::one()).unique_saturated_into(); let padded_leaves = num_of_leaves.next_power_of_two(); let total_capacity = (2 * padded_leaves) - 1; + if total_capacity == 0 { return Err(MerkleTreeError::Empty); } + let empty_hash = H::empty(); let empty_leaf_hash = leaf_prefixed::(empty_hash.as_ref()); let mut merkle_tree = sp_std::vec![empty_leaf_hash; total_capacity]; + let mut buffer = Vec::with_capacity(size_hint + 1); + for item in raw_values.into_iter() { - let (index, preimage) = generate_preimage(item)?; - if index >= padded_leaves { return Ok(sp_std::vec![]); } - merkle_tree[index] = leaf_prefixed::(&preimage) + buffer.clear(); + buffer.push(LEAF_PREFIX); + + let index = update_preimage_buffer(item, &mut buffer) + .map_err(|err| MerkleTreeError::Generate(err))?; + + if index >= padded_leaves { return Err(MerkleTreeError::InvalidIndex); } + + merkle_tree[index] = H::hash(&buffer); } let mut layer_start = 0; @@ -58,10 +66,15 @@ where while current_layer_len > 1 { for i in (0..current_layer_len).step_by(2) { let layer_index = layer_start + i; - let left_bytes = merkle_tree[layer_index].as_ref(); - let right_bytes = merkle_tree[layer_index | 1].as_ref(); + let left_bytes = merkle_tree[layer_index]; + let right_bytes = merkle_tree[layer_index | 1]; - merkle_tree[write_ptr] = node_prefixed::(left_bytes, right_bytes); + let combined_hash = H::combine_hashes_prefixed( + NODE_PREFIX, + left_bytes, + right_bytes, + ); + merkle_tree[write_ptr] = combined_hash; write_ptr += 1; } layer_start += current_layer_len; @@ -71,16 +84,16 @@ where Ok(merkle_tree) } -pub fn generate_proof( +pub fn generate_proof( merkle_tree: &[H::Hash], max_index: I, index: I, -) -> Vec +) -> Result, MerkleTreeError> where H: GhostHasher, I: AtLeast8BitUnsigned + UniqueSaturatedInto, { - if index > max_index { return sp_std::vec![]; } + if index > max_index { return Err(MerkleTreeError::InvalidIndex); } let num_of_leaves: usize = (max_index + I::one()).unique_saturated_into(); let mut current_index: usize = index.unique_saturated_into(); @@ -101,7 +114,7 @@ where current_index >>= 1; } - proof + Ok(proof) } pub fn verify_tree_proof( @@ -120,14 +133,11 @@ where let mut current_hash = leaf_prefixed::(preimage); for sibling in merkle_proof.iter() { - let sibling_bytes = sibling.as_ref(); - let hash_bytes = current_hash.as_ref(); - - if current_index % 2 == 0 { - current_hash = node_prefixed::(hash_bytes, sibling_bytes); + current_hash = if current_index % 2 == 0 { + H::combine_hashes_prefixed(NODE_PREFIX, current_hash, *sibling) } else { - current_hash = node_prefixed::(sibling_bytes, hash_bytes); - } + H::combine_hashes_prefixed(NODE_PREFIX, *sibling, current_hash) + }; current_index >>= 1; } @@ -138,451 +148,582 @@ where #[cfg(test)] mod tests { use super::*; - use sp_core::H256; - - use crate::SubstrateBlake2Hasher; - - #[test] - fn test_generate_tree_perfect_power_of_two() { - let max_index = 3u32; - let raw_data = sp_std::vec![ - (0, sp_std::vec![1]), - (1, sp_std::vec![2]), - (2, sp_std::vec![3]), - (3, sp_std::vec![4]), - ]; - - let tree = generate_tree::( - max_index, - raw_data, - |(index, bytes)| Ok((index, bytes)), - ) - .unwrap(); - - assert_eq!(tree.len(), 7); - - let h0 = SubstrateBlake2Hasher::hash(&[ &[0u8], [1].as_slice() ].concat()); - let h1 = SubstrateBlake2Hasher::hash(&[ &[0u8], [2].as_slice() ].concat()); - let h2 = SubstrateBlake2Hasher::hash(&[ &[0u8], [3].as_slice() ].concat()); - let h3 = SubstrateBlake2Hasher::hash(&[ &[0u8], [4].as_slice() ].concat()); - - let mut c1 = [0u8; 65]; - c1[0] = NODE_PREFIX; - c1[1..33].copy_from_slice(h0.as_ref()); - c1[33..].copy_from_slice(h1.as_ref()); - let parent_left = SubstrateBlake2Hasher::hash(&c1); - - let mut c2 = [0u8; 65]; - c2[0] = NODE_PREFIX; - c2[1..33].copy_from_slice(h2.as_ref()); - c2[33..].copy_from_slice(h3.as_ref()); - let parent_right = SubstrateBlake2Hasher::hash(&c2); - - let mut c_root = [0u8; 65]; - c_root[0] = NODE_PREFIX; - c_root[1..33].copy_from_slice(parent_left.as_ref()); - c_root[33..].copy_from_slice(parent_right.as_ref()); - let expected_root = SubstrateBlake2Hasher::hash(&c_root); - - assert_eq!(tree.last().unwrap(), &expected_root); - } - - #[test] - fn test_generate_tree_with_padding_hashing() { - let max_index = 2u32; - let raw_data = sp_std::vec![ - (0, sp_std::vec![1]), - (1, sp_std::vec![2]), - (2, sp_std::vec![3]), - ]; - - let tree = generate_tree::( - max_index, - raw_data, - |(index, bytes)| Ok((index, bytes)), - ) - .unwrap(); - - let mut empty_data = sp_std::vec![LEAF_PREFIX]; - empty_data.extend_from_slice(&H256::zero().to_fixed_bytes()); - let empty_hash = SubstrateBlake2Hasher::hash(&empty_data); - - assert_eq!(tree.len(), 7); - assert_eq!(tree[3], empty_hash); - - let mut leaf2_data = sp_std::vec![LEAF_PREFIX]; - leaf2_data.extend_from_slice(&[3]); - let h2 = SubstrateBlake2Hasher::hash(&leaf2_data); - - let mut combined_right = [0u8; 65]; - combined_right[0] = NODE_PREFIX; - combined_right[1..33].copy_from_slice(h2.as_ref()); - combined_right[33..].copy_from_slice(empty_hash.as_ref()); - - let expected_parent_right = SubstrateBlake2Hasher::hash(&combined_right); - assert_eq!(tree[5], expected_parent_right); - } - - #[test] - fn test_generate_tree_determinism() { - let max_index = 1u32; - - let data_1 = sp_std::vec![(0, sp_std::vec![100]), (1, sp_std::vec![200])]; - let data_2 = sp_std::vec![(0, sp_std::vec![100]), (1, sp_std::vec![200])]; - let data_3 = sp_std::vec![(0, sp_std::vec![101]), (1, sp_std::vec![200])]; - - let tree_1 = - generate_tree::(max_index, data_1, |(i, b)| { - Ok((i, b)) - }) - .unwrap(); - let tree_2 = - generate_tree::(max_index, data_2, |(i, b)| { - Ok((i, b)) - }) - .unwrap(); - let tree_3 = - generate_tree::(max_index, data_3, |(i, b)| { - Ok((i, b)) - }) - .unwrap(); - - assert_eq!(tree_1.last().unwrap(), tree_2.last().unwrap()); - assert_ne!(tree_1.last().unwrap(), tree_3.last().unwrap()); - } - - #[test] - fn test_generate_tree_single_element() { - let max_index = 0u32; - let raw_data = sp_std::vec![(0, sp_std::vec![42])]; - - let tree = generate_tree::( - max_index, - raw_data, - |(index, bytes)| Ok((index, bytes)), - ) - .unwrap(); - - assert_eq!(tree.len(), 1); - assert_eq!(tree[0], SubstrateBlake2Hasher::hash(&vec![LEAF_PREFIX, 42])); - } - - #[test] - fn test_proof_and_verification_end_to_end() { - let max_index = 3u32; - let raw_data = sp_std::vec![ - (0, sp_std::vec![1]), - (1, sp_std::vec![2]), - (2, sp_std::vec![3]), - (3, sp_std::vec![4]), - ]; - - let tree = generate_tree::( - max_index, - raw_data, - |(index, bytes)| Ok((index, bytes)), - ) - .unwrap(); - - let root = *tree.last().unwrap(); - - let target_index = 2u32; - let proof = generate_proof::(&tree, max_index, target_index); - - assert_eq!(proof.len(), 2); - - let preimage = sp_std::vec![3]; - let is_valid = - verify_tree_proof::(&preimage, &proof, root, target_index); - - assert!(is_valid); - } - - #[test] - fn test_proof_index_out_of_bounds() { - let max_index = 3u32; - let raw_data = sp_std::vec![ - (0, sp_std::vec![1]), - (1, sp_std::vec![2]), - (2, sp_std::vec![3]), - (3, sp_std::vec![4]), - ]; - - let tree = generate_tree::( - max_index, - raw_data, - |(index, bytes)| Ok((index, bytes)), - ) - .unwrap(); - - let root = *tree.last().unwrap(); - - let target_index = 2u32; - let proof = generate_proof::(&tree, max_index, target_index); - - assert_eq!(proof.len(), 2); - - let preimage = sp_std::vec![3]; - let is_valid = - verify_tree_proof::(&preimage, &proof, root, target_index + 1); - - assert!(!is_valid); - } - - #[test] - fn test_proof_contains_correct_siblings() { - let max_index = 3u32; - let raw_data = sp_std::vec![ - (0, sp_std::vec![10]), - (1, sp_std::vec![20]), - (2, sp_std::vec![30]), - (3, sp_std::vec![40]), - ]; - - let tree = generate_tree::( - max_index, - raw_data, - |(index, bytes)| Ok((index, bytes)), - ) - .unwrap(); - - let h0 = tree[0]; - let h1 = tree[1]; - let h2 = tree[2]; - let h3 = tree[3]; - - let parent_left = tree[4]; // hash(h0 + h1) - let parent_right = tree[5]; // hash(h2 + h3) - - let proof_for_0 = generate_proof::(&tree, max_index, 0); - assert_eq!(proof_for_0[0], h1); - assert_eq!(proof_for_0[1], parent_right); - - let proof_for_1 = generate_proof::(&tree, max_index, 1); - assert_eq!(proof_for_1[0], h0); - assert_eq!(proof_for_1[1], parent_right); - - let proof_for_2 = generate_proof::(&tree, max_index, 2); - assert_eq!(proof_for_2[0], h3); - assert_eq!(proof_for_2[1], parent_left); - - let proof_for_3 = generate_proof::(&tree, max_index, 3); - assert_eq!(proof_for_3[0], h2); - assert_eq!(proof_for_3[1], parent_left); - } - - #[test] - fn test_proof_for_padded_tree_leaf() { - let max_index = 2u32; - let raw_data = sp_std::vec![ - (0, sp_std::vec![1]), - (1, sp_std::vec![2]), - (2, sp_std::vec![3]), - ]; - - let tree = generate_tree::( - max_index, - raw_data, - |(index, bytes)| Ok((index, bytes)), - ) - .unwrap(); - - let proof = generate_proof::(&tree, max_index, 2); - - let mut empty_data = sp_std::vec![LEAF_PREFIX]; - empty_data.extend_from_slice(&H256::zero().to_fixed_bytes()); - let empty_hash = SubstrateBlake2Hasher::hash(&empty_data); - assert_eq!(proof[0], empty_hash); - - let root = *tree.last().unwrap(); - let is_valid = - verify_tree_proof::(&sp_std::vec![3], &proof, root, 2); - assert!(is_valid); - } - - #[test] - fn test_proof_for_single_element_tree() { - let max_index = 0u32; - let raw_data = sp_std::vec![(0, sp_std::vec![99])]; - - let tree = generate_tree::( - max_index, - raw_data, - |(index, bytes)| Ok((index, bytes)), - ) - .unwrap(); - - let proof = generate_proof::(&tree, max_index, 0); - assert!(proof.is_empty()); - } - - #[test] - fn test_verify_tree_proof_with_h0_h3_and_random_values() { - let max_index = 3u32; - let leaf_0 = sp_std::vec![10u8]; - let leaf_1 = sp_std::vec![20u8]; - let leaf_2 = sp_std::vec![30u8]; - let leaf_3 = sp_std::vec![40u8]; - - let raw_data = sp_std::vec![ - (0, leaf_0.clone()), - (1, leaf_1.clone()), - (2, leaf_2.clone()), - (3, leaf_3.clone()) - ]; - - let tree = generate_tree::( - max_index, - raw_data, - |(index, bytes)| Ok((index, bytes)), - ) - .unwrap(); - - let root = *tree.last().unwrap(); - - let h0 = tree[0]; - let h1 = tree[1]; - let h2 = tree[2]; - let h3 = tree[3]; - - let proof_for_1 = generate_proof::(&tree, max_index, 1); - assert_eq!(proof_for_1[0], h0); - - let is_valid_1 = verify_tree_proof::(&leaf_1, &proof_for_1, root, 1); - assert!(is_valid_1); - - let proof_for_2 = generate_proof::(&tree, max_index, 2); - assert_eq!(proof_for_2[0], h3); - - let is_valid_2 = verify_tree_proof::(&leaf_2, &proof_for_2, root, 2); - assert!(is_valid_2); - - let proof_for_0 = generate_proof::(&tree, max_index, 0); - assert_eq!(proof_for_0[0], h1); - assert!(verify_tree_proof::( - &leaf_0, - &proof_for_0, - root, - 0 - )); - - let proof_for_3 = generate_proof::(&tree, max_index, 3); - assert_eq!(proof_for_3[0], h2); - assert!(verify_tree_proof::( - &leaf_3, - &proof_for_3, - root, - 3 - )); - - let mut bad_proof = proof_for_1.clone(); - bad_proof[0] = H256::random(); - let is_valid_bad_proof = - verify_tree_proof::(&leaf_1, &bad_proof, root, 1); - assert!(!is_valid_bad_proof); - - let bad_preimage = vec![99u8]; - let is_valid_bad_preimage = - verify_tree_proof::(&bad_preimage, &proof_for_1, root, 1); - assert!(!is_valid_bad_preimage); - } - - #[test] - fn test_verify_tree_proof_deep_tree_depth_4() { - let max_index = 15u32; - - let mut raw_data = sp_std::vec![]; - for i in 0..16 { - raw_data.push((i, sp_std::vec![i as u8])); + + use crate::{ + SubstrateBlake2Hasher, SubstrateKeccakHasher, UtxoSha2Hasher, + }; + + macro_rules! run_merkle_tree_tests { + ($hasher_type:ty, $mod_name:ident) => { + mod $mod_name { + use super::*; + + #[test] + fn test_generate_tree_perfect_power_of_two() { + let max_index = 3u32; + let size_hint = 69; + + let raw_data = sp_std::vec![ + (0, sp_std::vec![1; 69]), + (1, sp_std::vec![2; 69]), + (2, sp_std::vec![3; 69]), + (3, sp_std::vec![4; 69]), + ]; + + let tree = generate_tree::<$hasher_type, u32, _, _, ()>( + max_index, + size_hint, + raw_data, + |(index, bytes), buffer| { + buffer.extend_from_slice(&bytes); + Ok(index as usize) + } + ) + .ok() + .unwrap(); + + assert_eq!(tree.len(), 7); + + let h0 = <$hasher_type>::hash(&[ &[0u8], [1; 69].as_slice() ].concat()); + let h1 = <$hasher_type>::hash(&[ &[0u8], [2; 69].as_slice() ].concat()); + let h2 = <$hasher_type>::hash(&[ &[0u8], [3; 69].as_slice() ].concat()); + let h3 = <$hasher_type>::hash(&[ &[0u8], [4; 69].as_slice() ].concat()); + + let mut c1 = [0u8; 65]; + c1[0] = NODE_PREFIX; + c1[1..33].copy_from_slice(h0.as_ref()); + c1[33..].copy_from_slice(h1.as_ref()); + let parent_left = <$hasher_type>::hash(&c1); + + let mut c2 = [0u8; 65]; + c2[0] = NODE_PREFIX; + c2[1..33].copy_from_slice(h2.as_ref()); + c2[33..].copy_from_slice(h3.as_ref()); + let parent_right = <$hasher_type>::hash(&c2); + + let mut c_root = [0u8; 65]; + c_root[0] = NODE_PREFIX; + c_root[1..33].copy_from_slice(parent_left.as_ref()); + c_root[33..].copy_from_slice(parent_right.as_ref()); + let expected_root = <$hasher_type>::hash(&c_root); + + assert_eq!(tree.last().unwrap(), &expected_root); + } + + #[test] + fn test_generate_tree_with_padding_hashing() { + let max_index = 2u32; + let size_hint = 420; + + let raw_data = sp_std::vec![ + (0, sp_std::vec![1; size_hint]), + (1, sp_std::vec![2; size_hint]), + (2, sp_std::vec![3; size_hint]), + ]; + + let tree = generate_tree::<$hasher_type, u32, _, _, ()>( + max_index, + size_hint, + raw_data, + + |(index, bytes), buffer| { + buffer.extend_from_slice(&bytes); + Ok(index as usize) + } + ) + .ok() + .unwrap(); + + let mut empty_data = sp_std::vec![LEAF_PREFIX]; + empty_data.extend_from_slice(&<$hasher_type>::empty().to_fixed_bytes()); + let empty_hash = <$hasher_type>::hash(&empty_data); + + assert_eq!(tree.len(), 7); + assert_eq!(tree[3], empty_hash); + + let mut leaf2_data = sp_std::vec![LEAF_PREFIX]; + leaf2_data.extend_from_slice(&sp_std::vec![3; size_hint]); + let h2 = <$hasher_type>::hash(&leaf2_data); + + let mut combined_right = [0u8; 65]; + combined_right[0] = NODE_PREFIX; + combined_right[1..33].copy_from_slice(h2.as_ref()); + combined_right[33..].copy_from_slice(empty_hash.as_ref()); + + let expected_parent_right = <$hasher_type>::hash(&combined_right); + assert_eq!(tree[5], expected_parent_right); + } + + #[test] + fn test_generate_tree_determinism() { + let max_index = 1u32; + + let data_1 = sp_std::vec![(0, sp_std::vec![100]), (1, sp_std::vec![200])]; + let data_2 = sp_std::vec![(0, sp_std::vec![100]), (1, sp_std::vec![200])]; + let data_3 = sp_std::vec![(0, sp_std::vec![101]), (1, sp_std::vec![200])]; + + let tree_1 = generate_tree::<$hasher_type, u32, _, _, ()>( + max_index, + 1, + data_1, + + |(index, bytes), buffer| { + buffer.extend_from_slice(&bytes); + Ok(index as usize) + } + ) + .ok() + .unwrap(); + + let tree_2 = generate_tree::<$hasher_type, u32, _, _, ()>( + max_index, + 1, + data_2, + + |(index, bytes), buffer| { + buffer.extend_from_slice(&bytes); + Ok(index as usize) + } + ) + .ok() + .unwrap(); + + let tree_3 = generate_tree::<$hasher_type, u32, _, _, ()>( + max_index, + 1, + data_3, + + |(index, bytes), buffer| { + buffer.extend_from_slice(&bytes); + Ok(index as usize) + } + ) + .ok() + .unwrap(); + + assert_eq!(tree_1.last().unwrap(), tree_2.last().unwrap()); + assert_ne!(tree_1.last().unwrap(), tree_3.last().unwrap()); + } + + #[test] + fn test_generate_tree_single_element() { + let max_index = 0u32; + let size_hint = 1337; + let raw_data = sp_std::vec![(0, sp_std::vec![42; size_hint])]; + + let tree = generate_tree::<$hasher_type, u32, _, _, ()>( + max_index, + size_hint, + raw_data, + + |(index, bytes), buffer| { + buffer.extend_from_slice(&bytes); + Ok(index as usize) + } + ) + .ok() + .unwrap(); + + let expected_leaf_data = [ + &[LEAF_PREFIX][..], + sp_std::vec![42; size_hint].as_slice() + ].concat(); + let expected_hash = <$hasher_type>::hash(&expected_leaf_data); + + assert_eq!(tree.len(), 1); + assert_eq!(tree[0], expected_hash); + } + + #[test] + fn test_proof_and_verification_end_to_end() { + let max_index = 3u32; + let size_hint = 420; + + let raw_data = sp_std::vec![ + (0, sp_std::vec![1; size_hint]), + (1, sp_std::vec![2; size_hint]), + (2, sp_std::vec![3; size_hint]), + (3, sp_std::vec![4; size_hint]), + ]; + + let tree = generate_tree::<$hasher_type, u32, _, _, ()>( + max_index, + size_hint, + raw_data, + + |(index, bytes), buffer| { + buffer.extend_from_slice(&bytes); + Ok(index as usize) + } + ) + .ok() + .unwrap(); + + let root = *tree.last().unwrap(); + + let target_index = 2u32; + let proof = generate_proof::<$hasher_type, u32, ()>(&tree, max_index, target_index) + .ok() + .unwrap(); + + assert_eq!(proof.len(), 2); + + let preimage = sp_std::vec![3; size_hint]; + assert!(verify_tree_proof::<$hasher_type, u32>(&preimage, &proof, root, target_index)); + } + + #[test] + fn test_proof_index_out_of_bounds() { + let max_index = 3u32; + let size_hint = 420; + + let raw_data = sp_std::vec![ + (0, sp_std::vec![1; size_hint]), + (1, sp_std::vec![2; size_hint]), + (2, sp_std::vec![3; size_hint]), + (3, sp_std::vec![4; size_hint]), + ]; + + let tree = generate_tree::<$hasher_type, u32, _, _, ()>( + max_index, + size_hint, + raw_data, + + |(index, bytes), buffer| { + buffer.extend_from_slice(&bytes); + Ok(index as usize) + } + ) + .ok() + .unwrap(); + + let root = *tree.last().unwrap(); + + let target_index = 2u32; + let proof = generate_proof::<$hasher_type, u32, ()>(&tree, max_index, target_index) + .ok() + .unwrap(); + + assert_eq!(proof.len(), 2); + + let preimage = sp_std::vec![3; size_hint]; + assert!(verify_tree_proof::<$hasher_type, u32>(&preimage, &proof, root, target_index)); + assert!(!verify_tree_proof::<$hasher_type, u32>(&preimage, &proof, root, 4)); + } + + #[test] + fn test_proof_contains_correct_siblings() { + let max_index = 3u32; + let size_hint = 1337; + + let raw_data = sp_std::vec![ + (0, sp_std::vec![10; size_hint]), + (1, sp_std::vec![20; size_hint]), + (2, sp_std::vec![30; size_hint]), + (3, sp_std::vec![40; size_hint]), + ]; + + let tree = generate_tree::<$hasher_type, u32, _, _, ()>( + max_index, + size_hint, + raw_data, + + |(index, bytes), buffer| { + buffer.extend_from_slice(&bytes); + Ok(index as usize) + } + ) + .ok() + .unwrap(); + + let h0 = tree[0]; + let h1 = tree[1]; + let h2 = tree[2]; + let h3 = tree[3]; + + let parent_left = tree[4]; // hash(h0 + h1) + let parent_right = tree[5]; // hash(h2 + h3) + + let proof_for_0 = generate_proof::<$hasher_type, u32, ()>(&tree, max_index, 0) + .ok() + .unwrap(); + + assert_eq!(proof_for_0[0], h1); + assert_eq!(proof_for_0[1], parent_right); + + let proof_for_1 = generate_proof::<$hasher_type, u32, ()>(&tree, max_index, 1) + .ok() + .unwrap(); + + assert_eq!(proof_for_1[0], h0); + assert_eq!(proof_for_1[1], parent_right); + + let proof_for_2 = generate_proof::<$hasher_type, u32, ()>(&tree, max_index, 2) + .ok() + .unwrap(); + + assert_eq!(proof_for_2[0], h3); + assert_eq!(proof_for_2[1], parent_left); + + let proof_for_3 = generate_proof::<$hasher_type, u32, ()>(&tree, max_index, 3) + .ok() + .unwrap(); + + assert_eq!(proof_for_3[0], h2); + assert_eq!(proof_for_3[1], parent_left); + } + + #[test] + fn test_proof_for_padded_tree_leaf() { + let max_index = 2u32; + let size_hint = 69; + + let raw_data = sp_std::vec![ + (0, sp_std::vec![1; size_hint]), + (1, sp_std::vec![2; size_hint]), + (2, sp_std::vec![3; size_hint]), + ]; + + let tree = generate_tree::<$hasher_type, u32, _, _, ()>( + max_index, + size_hint, + raw_data, + + |(index, bytes), buffer| { + buffer.extend_from_slice(&bytes); + Ok(index as usize) + } + ) + .ok() + .unwrap(); + + let proof = generate_proof::<$hasher_type, u32, ()>(&tree, max_index, 2) + .ok() + .unwrap(); + + let mut empty_data = sp_std::vec![LEAF_PREFIX]; + empty_data.extend_from_slice(&<$hasher_type>::empty().to_fixed_bytes()); + let empty_hash = <$hasher_type>::hash(&empty_data); + assert_eq!(proof[0], empty_hash); + + let root = *tree.last().unwrap(); + assert!(verify_tree_proof::<$hasher_type, u32>(&sp_std::vec![3; size_hint], &proof, root, 2)); + } + + #[test] + fn test_proof_for_single_element_tree() { + let max_index = 0u32; + let raw_data = sp_std::vec![(0, sp_std::vec![99])]; + + let tree = generate_tree::<$hasher_type, u32, _, _, ()>( + max_index, + 1, + raw_data, + + |(index, bytes), buffer| { + buffer.extend_from_slice(&bytes); + Ok(index as usize) + } + ) + .ok() + .unwrap(); + + let proof = generate_proof::<$hasher_type, u32, ()>( + &tree, + max_index, + 0 + ) + .ok() + .unwrap(); + + assert!(proof.is_empty()); + } + + #[test] + fn test_verify_tree_proof_with_h0_h3_and_random_values() { + let max_index = 3u32; + let size_hint = 69; + + let leaf_0 = sp_std::vec![10u8; 69]; + let leaf_1 = sp_std::vec![20u8; 69]; + let leaf_2 = sp_std::vec![30u8; 69]; + let leaf_3 = sp_std::vec![40u8; 69]; + + let raw_data = sp_std::vec![ + (0, leaf_0.clone()), + (1, leaf_1.clone()), + (2, leaf_2.clone()), + (3, leaf_3.clone()) + ]; + + let tree = generate_tree::<$hasher_type, u32, _, _, ()>( + max_index, + size_hint, + raw_data, + + |(index, bytes), buffer| { + buffer.extend_from_slice(&bytes); + Ok(index as usize) + } + ) + .ok() + .unwrap(); + + let root = *tree.last().unwrap(); + + let h0 = tree[0]; + let h1 = tree[1]; + let h2 = tree[2]; + let h3 = tree[3]; + + let proof_for_1 = generate_proof::<$hasher_type, u32, ()>(&tree, max_index, 1).ok().unwrap(); + assert_eq!(proof_for_1[0], h0); + + let is_valid_1 = verify_tree_proof::<$hasher_type, u32>(&leaf_1, &proof_for_1, root, 1); + assert!(is_valid_1); + + let proof_for_2 = generate_proof::<$hasher_type, u32, ()>(&tree, max_index, 2).ok().unwrap(); + assert_eq!(proof_for_2[0], h3); + + let is_valid_2 = verify_tree_proof::<$hasher_type, u32>(&leaf_2, &proof_for_2, root, 2); + assert!(is_valid_2); + + let proof_for_0 = generate_proof::<$hasher_type, u32, ()>(&tree, max_index, 0).ok().unwrap(); + assert_eq!(proof_for_0[0], h1); + assert!(verify_tree_proof::<$hasher_type, u32>(&leaf_0, &proof_for_0, root, 0)); + + let proof_for_3 = generate_proof::<$hasher_type, u32, ()>(&tree, max_index, 3).ok().unwrap(); + assert_eq!(proof_for_3[0], h2); + assert!(verify_tree_proof::<$hasher_type, u32>(&leaf_3, &proof_for_3, root, 3)); + + let mut bad_proof = proof_for_1.clone(); + bad_proof[0] = <$hasher_type>::empty(); + let is_valid_bad_proof = + verify_tree_proof::<$hasher_type, u32>(&leaf_1, &bad_proof, root, 1); + assert!(!is_valid_bad_proof); + + let bad_preimage = vec![99u8]; + let is_valid_bad_preimage = + verify_tree_proof::<$hasher_type, u32>(&bad_preimage, &proof_for_1, root, 1); + assert!(!is_valid_bad_preimage); + } + + #[test] + fn test_verify_tree_proof_deep_tree_depth_4() { + let max_index = 15u32; + + let mut raw_data = sp_std::vec![]; + for i in 0..16 { + raw_data.push((i, sp_std::vec![i as u8])); + } + + let tree = generate_tree::<$hasher_type, u32, _, _, ()>( + max_index, + 1, + raw_data, + + |(index, bytes), buffer| { + buffer.extend_from_slice(&bytes); + Ok(index as usize) + } + ) + .ok() + .unwrap(); + + let root = *tree.last().unwrap(); + + assert_eq!(tree.len(), 31); + + let target_index = 11u32; + let preimage = vec![target_index as u8]; + + let proof = generate_proof::<$hasher_type, u32, ()>(&tree, max_index, target_index) + .ok() + .unwrap(); + + assert_eq!(proof.len(), 4); + + let h10 = tree[10]; + assert_eq!(proof[0], h10); + + assert!(verify_tree_proof::<$hasher_type, u32>(&preimage, &proof, root, target_index)); + assert!(!verify_tree_proof::<$hasher_type, u32>(&preimage, &proof, root, 4)); + } + + #[test] + fn test_verification_index_at_upper_bound() { + let max_index = 3u32; + let size_hint = 420; + + let raw_data = sp_std::vec![ + (0, sp_std::vec![1; size_hint]), + (1, sp_std::vec![2; size_hint]), + (2, sp_std::vec![3; size_hint]), + (3, sp_std::vec![4; size_hint]), + ]; + + let tree = generate_tree::<$hasher_type, u32, _, _, ()>( + max_index, + size_hint, + raw_data, + + |(index, bytes), buffer| { + buffer.extend_from_slice(&bytes); + Ok(index as usize) + } + ) + .ok() + .unwrap(); + + let root = *tree.last().unwrap(); + + let target_index = 3u32; + let proof = generate_proof::<$hasher_type, u32, ()>(&tree, max_index, target_index) + .ok() + .unwrap(); + + assert_eq!(proof.len(), 2); + + let preimage = sp_std::vec![4; size_hint]; + assert!(verify_tree_proof::<$hasher_type, u32>(&preimage, &proof, root, target_index)); + } + + #[test] + fn test_verification_fails_when_index_out_of_bounds() { + let max_index = 3u32; + let raw_data = sp_std::vec![ + (0, sp_std::vec![1]), + (1, sp_std::vec![2]), + (2, sp_std::vec![3]), + (3, sp_std::vec![4]), + ]; + + let tree = generate_tree::<$hasher_type, u32, _, _, ()>( + max_index, + 1, + raw_data, + + |(index, bytes), buffer| { + buffer.extend_from_slice(&bytes); + Ok(index as usize) + } + ) + .ok() + .unwrap(); + + let root = *tree.last().unwrap(); + let proof = generate_proof::<$hasher_type, u32, ()>(&tree, max_index, 3u32).ok().unwrap(); + + let target_index = 4u32; + let preimage = sp_std::vec![4]; + + assert!(!verify_tree_proof::<$hasher_type, u32>(&preimage, &proof, root, target_index)); + } + } } - - let tree = generate_tree::( - max_index, - raw_data, - |(index, bytes)| Ok((index, bytes)), - ) - .unwrap(); - - let root = *tree.last().unwrap(); - - assert_eq!(tree.len(), 31); - - let target_index = 11u32; - let preimage = vec![target_index as u8]; - - let proof = generate_proof::(&tree, max_index, target_index); - assert_eq!(proof.len(), 4); - - let h10 = tree[10]; - assert_eq!(proof[0], h10); - - let is_valid = - verify_tree_proof::(&preimage, &proof, root, target_index); - assert!(is_valid); - - let is_valid_wrong_index = - verify_tree_proof::(&preimage, &proof, root, 4); - assert!(!is_valid_wrong_index); } - #[test] - fn test_verification_index_at_upper_bound() { - let max_index = 3u32; - let raw_data = sp_std::vec![ - (0, sp_std::vec![1]), - (1, sp_std::vec![2]), - (2, sp_std::vec![3]), - (3, sp_std::vec![4]), - ]; - - let tree = generate_tree::( - max_index, - raw_data, - |(index, bytes)| Ok((index, bytes)), - ) - .unwrap(); - - let root = *tree.last().unwrap(); - - let target_index = 3u32; - let proof = generate_proof::(&tree, max_index, target_index); - - assert_eq!(proof.len(), 2); - - let preimage = sp_std::vec![4]; - let is_valid = - verify_tree_proof::(&preimage, &proof, root, target_index); - - assert!(is_valid); - } - - #[test] - fn test_verification_fails_when_index_out_of_bounds() { - let max_index = 3u32; - let raw_data = sp_std::vec![ - (0, sp_std::vec![1]), - (1, sp_std::vec![2]), - (2, sp_std::vec![3]), - (3, sp_std::vec![4]), - ]; - - let tree = generate_tree::( - max_index, - raw_data, - |(index, bytes)| Ok((index, bytes)), - ) - .unwrap(); - - let root = *tree.last().unwrap(); - - let proof = generate_proof::(&tree, max_index, 3u32); - - let target_index = 4u32; - let preimage = sp_std::vec![4]; - - let is_valid = - verify_tree_proof::(&preimage, &proof, root, target_index); - - assert!(!is_valid); - } + run_merkle_tree_tests!(SubstrateBlake2Hasher, blake2_merkle_tree_tests); + run_merkle_tree_tests!(SubstrateKeccakHasher, keccak_merkle_tree_tests); + run_merkle_tree_tests!(UtxoSha2Hasher, sha2_merkle_tree_tests); } diff --git a/pallets/traits/src/exodus.rs b/pallets/traits/src/exodus.rs index 6ddd7ce..0d9fe1a 100644 --- a/pallets/traits/src/exodus.rs +++ b/pallets/traits/src/exodus.rs @@ -6,7 +6,14 @@ use sp_std::{ vec::Vec, }; -pub trait IdentifierConverter { +pub trait CiphersuiteSizes { + fn header_bytes_len(&self) -> usize; + fn element_bytes_len(&self) -> usize; + fn signature_bytes_len(&self) -> usize; + fn scalar_bytes_len(&self) -> usize; +} + +pub trait IdentifierConverter: CiphersuiteSizes { fn non_zero_index(index: A) -> Result; fn convert_identifier_to_index(bytes: &[u8]) -> Result; } @@ -15,11 +22,6 @@ pub trait FlexibleRoundOptimizedSchnorrThresholdSignature: MerkleTreeBuild type Packages<'a>; type NoncePackages<'a>; - fn header_bytes_len(&self) -> usize; - fn element_bytes_len(&self) -> usize; - fn signature_bytes_len(&self) -> usize; - fn scalar_bytes_len(&self) -> usize; - fn verify_signature_share( &self, i: A, @@ -91,13 +93,14 @@ pub trait MerkleTreeBuilder: IdentifierConverter { type Hash; fn build_merkle_tree( + size_hint: usize, verifying_shares: &BTreeMap, V>, serialize_fn: F, ) -> Result, E> where C: frost_core::Ciphersuite, <::Group as Group>::Field: frost_core::Field, - F: Fn(&V) -> Result, E>; + F: Fn(&V, &mut Vec) -> Result<(), E>; fn generate_merkle_proof_from_tree( verifying_shares: &BTreeMap, V>, diff --git a/pallets/traits/src/hashing.rs b/pallets/traits/src/hashing.rs index f115b76..84a95ac 100644 --- a/pallets/traits/src/hashing.rs +++ b/pallets/traits/src/hashing.rs @@ -22,6 +22,12 @@ pub trait GhostHasher { } fn combine_hashes(left: Self::Hash, right: Self::Hash) -> Self::Hash; + fn combine_hashes_prefixed( + prefix: u8, + left: Self::Hash, + right: Self::Hash, + ) -> Self::Hash; + fn hash(data: &[u8]) -> Self::Hash; fn empty() -> Self::Hash; fn hash_len() -> usize;