refactor(stm): harden legacy decoders and key registration

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 }

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

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 }

mithril-stm/src/membership_commitment/merkle_tree/path.rs

@@ -58,14 +58,20 @@ impl<D: Digest + FixedOutput> MerklePath<D> {
         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::new();
         for i in 0..len {
+            let range_low = i
+                .checked_mul(<D as Digest>::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(<D as Digest>::output_size()))
+                .and_then(|rh| rh.checked_add(16))
+                .ok_or(MerkleTreeError::SerializationError)?;
             values.push(
                 bytes
-                    .get(
-                        16 + i * <D as Digest>::output_size()
-                            ..16 + (i + 1) * <D as Digest>::output_size(),
-                    )
+                    .get(range_low..range_high)
                     .ok_or(MerkleTreeError::SerializationError)?
                     .to_vec(),
             );
@@ -118,27 +124,53 @@ impl<D: Digest + FixedOutput> MerkleBatchPath<D> {
     /// * Indices
     fn from_bytes_legacy(bytes: &[u8]) -> StmResult<Self> {
         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::new();
         for i in 0..len_v {
+            let range_low = i
+                .checked_mul(<D as Digest>::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(<D as Digest>::output_size()))
+                .and_then(|rh| rh.checked_add(16))
+                .ok_or(MerkleTreeError::SerializationError)?;
             values.push(
-                bytes[16 + i * <D as Digest>::output_size()
-                    ..16 + (i + 1) * <D as Digest>::output_size()]
+                bytes
+                    .get(range_low..range_high)
+                    .ok_or(MerkleTreeError::SerializationError)?
                     .to_vec(),
             );
         }
-        let offset = 16 + len_v * <D as Digest>::output_size();
+        let offset = len_v
+            .checked_mul(<D as Digest>::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::new();
         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)?,

mithril-stm/src/proof_system/concatenation/aggregate_key.rs

@@ -48,10 +48,11 @@ impl<D: MembershipDigest> AggregateVerificationKeyForConcatenation<D> {
         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,

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,16 +15,22 @@ 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<RegistrationEntry>,
+    registered_keys_for_concatenation: HashSet<VerificationKeyForConcatenation>,
+    #[cfg(feature = "future_snark")]
+    registered_keys_for_snark: HashSet<VerificationKeyForSnark>,
 }
 
 impl KeyRegistration {
     /// Initialize an empty registration
     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::<RegisterError>(),
+            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::<RegisterError>(),
+            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::<RegisterError>(){
                                 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::<BTreeSet<RegistrationEntry>>();
-                assert!(retrieved_keys == keys);
+                assert!(retrieved_keys == registered_entries);
             }
         }
     }