From c7c786649a980be5651216323a66fe78e14cd8a3 Mon Sep 17 00:00:00 2001 From: Dam Date: Mon, 20 Jul 2026 12:42:53 +0200 Subject: [PATCH 1/4] refactor(stm): updated legacy codecs --- .../membership_commitment/merkle_tree/path.rs | 58 ++++++++++++++----- 1 file changed, 45 insertions(+), 13 deletions(-) diff --git a/mithril-stm/src/membership_commitment/merkle_tree/path.rs b/mithril-stm/src/membership_commitment/merkle_tree/path.rs index 1cce05985f1..1ff1aaffa9f 100644 --- a/mithril-stm/src/membership_commitment/merkle_tree/path.rs +++ b/mithril-stm/src/membership_commitment/merkle_tree/path.rs @@ -58,14 +58,20 @@ impl MerklePath { u64_bytes.copy_from_slice(bytes.get(8..16).ok_or(MerkleTreeError::SerializationError)?); let len = usize::try_from(u64::from_be_bytes(u64_bytes)) .map_err(|_| MerkleTreeError::SerializationError)?; - let mut values = Vec::with_capacity(len); + let mut values = vec![]; for i in 0..len { + let range_low = i + .checked_mul(::output_size()) + .and_then(|rl| rl.checked_add(16)) + .ok_or(MerkleTreeError::SerializationError)?; + let range_high = i + .checked_add(1) + .and_then(|rh| rh.checked_mul(::output_size())) + .and_then(|rh| rh.checked_add(16)) + .ok_or(MerkleTreeError::SerializationError)?; values.push( bytes - .get( - 16 + i * ::output_size() - ..16 + (i + 1) * ::output_size(), - ) + .get(range_low..range_high) .ok_or(MerkleTreeError::SerializationError)? .to_vec(), ); @@ -118,27 +124,53 @@ impl MerkleBatchPath { /// * Indices fn from_bytes_legacy(bytes: &[u8]) -> StmResult { let mut u64_bytes = [0u8; 8]; - u64_bytes.copy_from_slice(&bytes[..8]); + u64_bytes.copy_from_slice(&bytes.get(..8).ok_or(MerkleTreeError::SerializationError)?); let len_v = usize::try_from(u64::from_be_bytes(u64_bytes)) .map_err(|_| MerkleTreeError::SerializationError)?; - u64_bytes.copy_from_slice(&bytes[8..16]); + u64_bytes.copy_from_slice(&bytes.get(8..16).ok_or(MerkleTreeError::SerializationError)?); let len_i = usize::try_from(u64::from_be_bytes(u64_bytes)) .map_err(|_| MerkleTreeError::SerializationError)?; - let mut values = Vec::with_capacity(len_v); + let mut values = vec![]; for i in 0..len_v { + let range_low = i + .checked_mul(::output_size()) + .and_then(|rl| rl.checked_add(16)) + .ok_or(MerkleTreeError::SerializationError)?; + let range_high = i + .checked_add(1) + .and_then(|rh| rh.checked_mul(::output_size())) + .and_then(|rh| rh.checked_add(16)) + .ok_or(MerkleTreeError::SerializationError)?; values.push( - bytes[16 + i * ::output_size() - ..16 + (i + 1) * ::output_size()] + bytes + .get(range_low..range_high) + .ok_or(MerkleTreeError::SerializationError)? .to_vec(), ); } - let offset = 16 + len_v * ::output_size(); + let offset = len_v + .checked_mul(::output_size()) + .and_then(|off| off.checked_add(16)) + .ok_or(MerkleTreeError::SerializationError)?; - let mut indices = Vec::with_capacity(len_v); + let mut indices = vec![]; for i in 0..len_i { - u64_bytes.copy_from_slice(&bytes[offset + i * 8..offset + (i + 1) * 8]); + let range_low = i + .checked_mul(8) + .and_then(|rl| rl.checked_add(offset)) + .ok_or(MerkleTreeError::SerializationError)?; + let range_high = i + .checked_add(1) + .and_then(|rh| rh.checked_mul(8)) + .and_then(|rh| rh.checked_add(offset)) + .ok_or(MerkleTreeError::SerializationError)?; + u64_bytes.copy_from_slice( + &bytes + .get(range_low..range_high) + .ok_or(MerkleTreeError::SerializationError)?, + ); indices.push( usize::try_from(u64::from_be_bytes(u64_bytes)) .map_err(|_| MerkleTreeError::SerializationError)?, From b2f9f70534a8eafb110251dc11d2a0d6ce78e40c Mon Sep 17 00:00:00 2001 From: Dam Date: Mon, 20 Jul 2026 12:45:04 +0200 Subject: [PATCH 2/4] refactor(stm): updated legacy codec for AggregateVerificationKeyForConcatenation --- .../src/membership_commitment/merkle_tree/path.rs | 12 ++++++------ .../src/proof_system/concatenation/aggregate_key.rs | 5 +++-- 2 files changed, 9 insertions(+), 8 deletions(-) diff --git a/mithril-stm/src/membership_commitment/merkle_tree/path.rs b/mithril-stm/src/membership_commitment/merkle_tree/path.rs index 1ff1aaffa9f..975e10a1fa3 100644 --- a/mithril-stm/src/membership_commitment/merkle_tree/path.rs +++ b/mithril-stm/src/membership_commitment/merkle_tree/path.rs @@ -58,7 +58,7 @@ impl MerklePath { u64_bytes.copy_from_slice(bytes.get(8..16).ok_or(MerkleTreeError::SerializationError)?); let len = usize::try_from(u64::from_be_bytes(u64_bytes)) .map_err(|_| MerkleTreeError::SerializationError)?; - let mut values = vec![]; + let mut values = Vec::new(); for i in 0..len { let range_low = i .checked_mul(::output_size()) @@ -124,15 +124,15 @@ impl MerkleBatchPath { /// * Indices fn from_bytes_legacy(bytes: &[u8]) -> StmResult { let mut u64_bytes = [0u8; 8]; - u64_bytes.copy_from_slice(&bytes.get(..8).ok_or(MerkleTreeError::SerializationError)?); + u64_bytes.copy_from_slice(bytes.get(..8).ok_or(MerkleTreeError::SerializationError)?); let len_v = usize::try_from(u64::from_be_bytes(u64_bytes)) .map_err(|_| MerkleTreeError::SerializationError)?; - u64_bytes.copy_from_slice(&bytes.get(8..16).ok_or(MerkleTreeError::SerializationError)?); + u64_bytes.copy_from_slice(bytes.get(8..16).ok_or(MerkleTreeError::SerializationError)?); let len_i = usize::try_from(u64::from_be_bytes(u64_bytes)) .map_err(|_| MerkleTreeError::SerializationError)?; - let mut values = vec![]; + let mut values = Vec::new(); for i in 0..len_v { let range_low = i .checked_mul(::output_size()) @@ -155,7 +155,7 @@ impl MerkleBatchPath { .and_then(|off| off.checked_add(16)) .ok_or(MerkleTreeError::SerializationError)?; - let mut indices = vec![]; + let mut indices = Vec::new(); for i in 0..len_i { let range_low = i .checked_mul(8) @@ -167,7 +167,7 @@ impl MerkleBatchPath { .and_then(|rh| rh.checked_add(offset)) .ok_or(MerkleTreeError::SerializationError)?; u64_bytes.copy_from_slice( - &bytes + bytes .get(range_low..range_high) .ok_or(MerkleTreeError::SerializationError)?, ); diff --git a/mithril-stm/src/proof_system/concatenation/aggregate_key.rs b/mithril-stm/src/proof_system/concatenation/aggregate_key.rs index 98d87e6970c..26c9884f0fe 100644 --- a/mithril-stm/src/proof_system/concatenation/aggregate_key.rs +++ b/mithril-stm/src/proof_system/concatenation/aggregate_key.rs @@ -48,10 +48,11 @@ impl AggregateVerificationKeyForConcatenation { let mut u64_bytes = [0u8; 8]; let size = bytes.len(); - u64_bytes.copy_from_slice(&bytes[size - 8..]); + let split = size.checked_sub(8).ok_or(MerkleTreeError::SerializationError)?; + u64_bytes.copy_from_slice(bytes.get(split..).ok_or(MerkleTreeError::SerializationError)?); let stake = u64::from_be_bytes(u64_bytes); let mt_commitment = MerkleTreeBatchCommitment::from_bytes( - bytes.get(..size - 8).ok_or(MerkleTreeError::SerializationError)?, + bytes.get(..split).ok_or(MerkleTreeError::SerializationError)?, )?; Ok(Self { mt_commitment, From 33685daf57008f878410198750e23cbd01076259 Mon Sep 17 00:00:00 2001 From: Dam Date: Mon, 20 Jul 2026 14:29:12 +0200 Subject: [PATCH 3/4] refactor(stm): updated KeyRegistration struct and register_by_entry --- .../src/protocol/key_registration/register.rs | 110 ++++++++++++++++-- 1 file changed, 99 insertions(+), 11 deletions(-) diff --git a/mithril-stm/src/protocol/key_registration/register.rs b/mithril-stm/src/protocol/key_registration/register.rs index a8a9561afb5..ff42604bcfa 100644 --- a/mithril-stm/src/protocol/key_registration/register.rs +++ b/mithril-stm/src/protocol/key_registration/register.rs @@ -1,8 +1,9 @@ +use std::collections::{BTreeSet, HashSet}; + use digest::{Digest, FixedOutput}; -use std::collections::BTreeSet; use crate::{ - Parameters, RegisterError, SignerIndex, Stake, StmResult, + Parameters, RegisterError, SignerIndex, Stake, StmResult, VerificationKeyForConcatenation, VerificationKeyProofOfPossessionForConcatenation, membership_commitment::{MerkleTree, MerkleTreeLeaf}, protocol::key_registration::ClosedRegistrationEntry, @@ -14,9 +15,12 @@ use crate::VerificationKeyForSnark; use super::RegistrationEntry; /// Key Registration -#[derive(Clone, Default, PartialEq, Eq, PartialOrd, Ord, Debug)] +#[derive(Clone, Default, PartialEq, Eq, Debug)] pub struct KeyRegistration { registration_entries: BTreeSet, + registered_keys_for_concatenation: HashSet, + #[cfg(feature = "future_snark")] + registered_keys_for_snark: HashSet, } impl KeyRegistration { @@ -24,6 +28,9 @@ impl KeyRegistration { pub fn initialize() -> Self { Self { registration_entries: Default::default(), + registered_keys_for_concatenation: Default::default(), + #[cfg(feature = "future_snark")] + registered_keys_for_snark: Default::default(), } } @@ -32,11 +39,28 @@ impl KeyRegistration { /// # Error /// The function fails when the entry is already registered. pub fn register_by_entry(&mut self, entry: &RegistrationEntry) -> StmResult<()> { - if !self.registration_entries.contains(entry) { - self.registration_entries.insert(*entry); - return Ok(()); + let vk_concatenation = entry.get_verification_key_for_concatenation(); + let is_already_registered = + self.registered_keys_for_concatenation.contains(&vk_concatenation); + + #[cfg(feature = "future_snark")] + let is_already_registered = is_already_registered + || entry + .get_verification_key_for_snark() + .is_some_and(|vk_snark| self.registered_keys_for_snark.contains(&vk_snark)); + + if is_already_registered { + return Err(RegisterError::EntryAlreadyRegistered(Box::new(*entry)).into()); } - Err(RegisterError::EntryAlreadyRegistered(Box::new(*entry)).into()) + + self.registered_keys_for_concatenation.insert(vk_concatenation); + #[cfg(feature = "future_snark")] + if let Some(vk_snark) = entry.get_verification_key_for_snark() { + self.registered_keys_for_snark.insert(vk_snark); + } + self.registration_entries.insert(*entry); + + Ok(()) } /// Registers a new signer with the given verification key proof of possession and stake. @@ -308,6 +332,66 @@ mod tests { } } + #[test] + fn register_by_entry_rejects_same_verification_key_with_different_stake() { + let mut rng = ChaCha20Rng::from_seed([0u8; 32]); + let mut kr = KeyRegistration::initialize(); + let vk_pop = VerificationKeyProofOfPossessionForConcatenation::from( + &BlsSigningKey::generate(&mut rng), + ); + + let first_entry = RegistrationEntry::new( + vk_pop, + 100, + #[cfg(feature = "future_snark")] + None, + ) + .unwrap(); + kr.register_by_entry(&first_entry) + .expect("registering a new verification key should succeed"); + + let second_entry = RegistrationEntry::new( + vk_pop, + 200, + #[cfg(feature = "future_snark")] + None, + ) + .unwrap(); + let result = kr.register_by_entry(&second_entry); + + assert!(matches!( + result.unwrap_err().downcast_ref::(), + Some(RegisterError::EntryAlreadyRegistered(_)) + )); + } + + #[cfg(feature = "future_snark")] + #[test] + fn register_by_entry_rejects_same_snark_key_with_different_concatenation_key() { + let mut rng = ChaCha20Rng::from_seed([0u8; 32]); + let mut kr = KeyRegistration::initialize(); + let schnorr_vk = + SchnorrVerificationKey::new_from_signing_key(SchnorrSigningKey::generate(&mut rng)); + + let first_vk_pop = VerificationKeyProofOfPossessionForConcatenation::from( + &BlsSigningKey::generate(&mut rng), + ); + let first_entry = RegistrationEntry::new(first_vk_pop, 100, Some(schnorr_vk)).unwrap(); + kr.register_by_entry(&first_entry) + .expect("registering a new verification key pair should succeed"); + + let second_vk_pop = VerificationKeyProofOfPossessionForConcatenation::from( + &BlsSigningKey::generate(&mut rng), + ); + let second_entry = RegistrationEntry::new(second_vk_pop, 200, Some(schnorr_vk)).unwrap(); + let result = kr.register_by_entry(&second_entry); + + assert!(matches!( + result.unwrap_err().downcast_ref::(), + Some(RegisterError::EntryAlreadyRegistered(_)) + )); + } + proptest! { #[test] fn test_keyreg(stake in vec(1..1u64 << 60, 2..=10), @@ -333,8 +417,10 @@ mod tests { VerificationKeyProofOfPossessionForConcatenation::from(&sk) }; - // Record successful registrations + // Record successful registrations, keyed by verification key since that's + // the uniqueness criterion enforced by register_by_entry let mut keys = BTreeSet::new(); + let mut registered_entries = BTreeSet::new(); for (i, &stake) in stake.iter().enumerate() { let mut pk = gen_keys[i % gen_keys.len()]; @@ -350,15 +436,17 @@ mod tests { match entry_result { Ok(entry) => { + let vk = entry.get_verification_key_for_concatenation(); let reg = kr.register_by_entry(&entry); match reg { Ok(_) => { - assert!(keys.insert(entry)); + assert!(keys.insert(vk)); + assert!(registered_entries.insert(entry)); }, Err(error) => match error.downcast_ref::(){ Some(RegisterError::EntryAlreadyRegistered(e1)) => { assert!(e1.as_ref() == &entry); - assert!(keys.contains(&entry)); + assert!(keys.contains(&vk)); }, _ => {panic!("Unexpected error: {error}")} } @@ -380,7 +468,7 @@ mod tests { let retrieved_keys = closed.closed_registration_entries.iter() .map(|entry| (*entry).clone().into()) .collect::>(); - assert!(retrieved_keys == keys); + assert!(retrieved_keys == registered_entries); } } } From 7b4d7f856bd838d87e7cefd055ce723937cb6b79 Mon Sep 17 00:00:00 2001 From: Dam Date: Wed, 22 Jul 2026 10:52:19 +0200 Subject: [PATCH 4/4] chore(stm): updated changelog and crate version --- Cargo.lock | 2 +- mithril-common/Cargo.toml | 2 +- mithril-stm/CHANGELOG.md | 7 +++++++ mithril-stm/Cargo.toml | 2 +- 4 files changed, 10 insertions(+), 3 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 5250b0222ff..f387371efac 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4816,7 +4816,7 @@ dependencies = [ [[package]] name = "mithril-stm" -version = "0.11.2" +version = "0.11.3" dependencies = [ "anyhow", "blake2 0.10.6", diff --git a/mithril-common/Cargo.toml b/mithril-common/Cargo.toml index f4d4d00ee34..06d6771ae58 100644 --- a/mithril-common/Cargo.toml +++ b/mithril-common/Cargo.toml @@ -46,7 +46,7 @@ fixed = "1.31.0" hex = { workspace = true } kes-summed-ed25519 = { version = "0.2.1", features = ["serde_enabled", "sk_clone_enabled"] } mithril-merkle-tree = { path = "../internal/mithril-merkle-tree", version = "0.1.4" } -mithril-stm = { path = "../mithril-stm", version = "0.11.2", default-features = false } +mithril-stm = { path = "../mithril-stm", version = "0.11.3", default-features = false } nom = "8.0.0" rand_chacha = { workspace = true } rand_core = { workspace = true } diff --git a/mithril-stm/CHANGELOG.md b/mithril-stm/CHANGELOG.md index eb44e089299..70704e401b2 100644 --- a/mithril-stm/CHANGELOG.md +++ b/mithril-stm/CHANGELOG.md @@ -5,6 +5,13 @@ All notable changes to this project will be documented in this file. The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). +## 0.11.3 (07-22-2026) + +### Changed + +- Updated the `from_bytes_legacy` functions for `MerklePath`, `MerkleBatchPath` and `AggregateVerificationKeyForConcatenation` +- Updated the `KeyRegistration` to track the registered key independently from the `registration_entries` + ## 0.11.2 (07-20-2026) ### Changed diff --git a/mithril-stm/Cargo.toml b/mithril-stm/Cargo.toml index f421adfff97..99805f7fccc 100644 --- a/mithril-stm/Cargo.toml +++ b/mithril-stm/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "mithril-stm" -version = "0.11.2" +version = "0.11.3" edition = { workspace = true } authors = { workspace = true } homepage = { workspace = true }