diff --git a/CHANGELOG.md b/CHANGELOG.md index 4692cf3a2..47aa8f63f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,7 +5,7 @@ - [BREAKING] Refactored how foreign account inputs are passed to `TransactionExecutor`, and upgraded Rust version to 1.86 (#1229). - [BREAKING] Add `TransactionHeader` and include it in batches and blocks (#1247). - [BREAKING] Hash keys in storage maps before insertion into the SMT (#1250). -- Add `AccountTree` and `PartialAccountTree` wrappers and enforce ID prefix uniqueness (#1254). +- Add `AccountTree` and `PartialAccountTree` wrappers and enforce ID prefix uniqueness (#1254, #1301). - Added getter for proof security level in `ProvenBatch` and `ProvenBlock` (#1259). - [BREAKING] Replaced the `ProvenBatch::new_unchecked` with the `ProvenBatch::new` method to initialize the struct with validations (#1260). - Added pretty print for `AccountCode` (#1273). diff --git a/crates/miden-lib/asm/shared/util/account_id.masm b/crates/miden-lib/asm/shared/util/account_id.masm index 5481f8d5c..23ffc4f38 100644 --- a/crates/miden-lib/asm/shared/util/account_id.masm +++ b/crates/miden-lib/asm/shared/util/account_id.masm @@ -13,7 +13,7 @@ const.ERR_ACCOUNT_ID_UNKNOWN_STORAGE_MODE=0x00020145 # Least significant byte of the account ID suffix must be zero. const.ERR_ACCOUNT_ID_LEAST_SIGNIFICANT_BYTE_MUST_BE_ZERO=0x00020144 -# Provided storage slot index is out of bounds. +# Provided storage slot index is out of bounds const.ERR_ACCOUNT_STORAGE_SLOT_INDEX_OUT_OF_BOUNDS=0x00020152 # The account ID must have storage mode public if the network flag is set. diff --git a/crates/miden-objects/src/block/account_tree.rs b/crates/miden-objects/src/block/account_tree.rs index 2ab57867d..bf706690a 100644 --- a/crates/miden-objects/src/block/account_tree.rs +++ b/crates/miden-objects/src/block/account_tree.rs @@ -55,13 +55,17 @@ impl AccountTree { /// Returns an error if: /// - the provided entries contain multiple commitments for the same account ID. /// - multiple account IDs share the same prefix. - pub fn with_entries( - entries: impl IntoIterator, - ) -> Result { + pub fn with_entries( + entries: impl IntoIterator, + ) -> Result + where + I: ExactSizeIterator, + { + let entries = entries.into_iter(); + let num_accounts = entries.len(); + let smt = Smt::with_entries( - entries - .into_iter() - .map(|(id, commitment)| (Self::id_to_smt_key(id), Word::from(commitment))), + entries.map(|(id, commitment)| (Self::id_to_smt_key(id), Word::from(commitment))), ) .map_err(|err| { let MerkleError::DuplicateValuesForIndex(leaf_idx) = err else { @@ -77,16 +81,22 @@ impl AccountTree { } })?; - for (leaf_idx, leaf) in smt.leaves() { - if leaf.num_entries() >= 2 { - // SAFETY: Since we only inserted account IDs into the SMT, it is guaranteed that - // the leaf_idx is a valid Felt as well as a valid account ID prefix. - return Err(AccountTreeError::DuplicateIdPrefix { - duplicate_prefix: AccountIdPrefix::new_unchecked( - Felt::try_from(leaf_idx.value()) - .expect("leaf index should be a valid felt"), - ), - }); + // If the number of leaves in the SMT is smaller than the number of accounts that were + // passed in, it means that at least one account ID pair ended up in the same leaf. If this + // is the case, we iterate the SMT entries to find the duplicated account ID prefix. + if smt.num_leaves() < num_accounts { + for (leaf_idx, leaf) in smt.leaves() { + if leaf.num_entries() >= 2 { + // SAFETY: Since we only inserted account IDs into the SMT, it is guaranteed + // that the leaf_idx is a valid Felt as well as a valid + // account ID prefix. + return Err(AccountTreeError::DuplicateIdPrefix { + duplicate_prefix: AccountIdPrefix::new_unchecked( + Felt::try_from(leaf_idx.value()) + .expect("leaf index should be a valid felt"), + ), + }); + } } } @@ -155,17 +165,45 @@ impl AccountTree { /// /// This is a thin wrapper around [`Smt::compute_mutations`]. See its documentation for more /// details. + /// + /// # Errors + /// + /// Returns an error if: + /// - an insertion of an account ID would violate the uniqueness of account ID prefixes in the + /// tree. pub fn compute_mutations( &self, account_commitments: impl IntoIterator, - ) -> AccountMutationSet { + ) -> Result { let mutation_set = self.smt.compute_mutations( account_commitments .into_iter() .map(|(id, commitment)| (Self::id_to_smt_key(id), Word::from(commitment))), ); - AccountMutationSet::new(mutation_set) + for id_key in mutation_set.new_pairs().keys() { + // Check if the insertion would be valid. + match self.smt.get_leaf(id_key) { + // Inserting into an empty leaf is valid. + SmtLeaf::Empty(_) => (), + SmtLeaf::Single((existing_key, _)) => { + // If the key matches the existing one, then we're updating the leaf, which is + // valid. If it does not match, then we would insert a duplicate. + if existing_key != *id_key { + return Err(AccountTreeError::DuplicateIdPrefix { + duplicate_prefix: Self::smt_key_to_id(*id_key).prefix(), + }); + } + }, + SmtLeaf::Multiple(_) => { + unreachable!( + "account tree should never contain duplicate ID prefixes and therefore never a multiple leaf" + ) + }, + } + } + + Ok(AccountMutationSet::new(mutation_set)) } // PUBLIC MUTATORS @@ -251,8 +289,10 @@ impl Default for AccountTree { // ACCOUNT MUTATION SET // ================================================================================================ -/// A newtype wrapper around a [`MutationSet`] which exists for type safety in some [`AccountTree`] -/// methods. +/// A newtype wrapper around a [`MutationSet`] for use in the [`AccountTree`]. +/// +/// It guarantees that applying the contained mutations will result in an account tree with unique +/// account ID prefixes. /// /// It is returned by and used in methods on the [`AccountTree`]. #[derive(Debug, Clone, PartialEq, Eq)] @@ -373,7 +413,9 @@ pub(super) mod tests { let mut tree = AccountTree::with_entries([(id0, digest0), (id1, digest1)]).unwrap(); - let mutations = tree.compute_mutations([(id0, digest1), (id1, digest2), (id2, digest3)]); + let mutations = tree + .compute_mutations([(id0, digest1), (id1, digest2), (id2, digest3)]) + .unwrap(); tree.apply_mutations(mutations).unwrap(); @@ -383,6 +425,21 @@ pub(super) mod tests { assert_eq!(tree.get(id2), digest3); } + #[test] + fn duplicates_in_compute_mutations() { + let [pair0, pair1] = setup_duplicate_prefix_ids(); + let id2 = AccountIdBuilder::new().build_with_seed([5; 32]); + let commitment2 = Digest::from([0, 0, 0, 99u32]); + + let tree = AccountTree::with_entries([pair0, (id2, commitment2)]).unwrap(); + + let err = tree.compute_mutations([pair1]).unwrap_err(); + + assert_matches!(err, AccountTreeError::DuplicateIdPrefix { + duplicate_prefix + } if duplicate_prefix == pair1.0.prefix()); + } + #[test] fn account_commitments() { let id0 = AccountIdBuilder::new().build_with_seed([5; 32]);