diff --git a/CHANGELOG.md b/CHANGELOG.md index a586319f4..fb1f8476d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -14,6 +14,7 @@ - Renamed card's names in the `miden-network-monitor` binary ([#1441](https://github.com/0xMiden/miden-node/pull/1441)). - Improved tracing in `miden-network-monitor` binary ([#1366](https://github.com/0xMiden/miden-node/pull/1366)). - Integrated RPC stack with Validator component for transaction validation ([#1457](https://github.com/0xMiden/miden-node/pull/1457)). +- Add partial storage map queries to RPC ([#1428](https://github.com/0xMiden/miden-node/pull/1428)). - Added explorer status to the `miden-network-monitor` binary ([#1450](https://github.com/0xMiden/miden-node/pull/1450)). - Added `GetLimits` endpoint to the RPC server ([#1410](https://github.com/0xMiden/miden-node/pull/1410)). diff --git a/Cargo.lock b/Cargo.lock index e911e94d0..8b0a45e45 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -5134,9 +5134,9 @@ dependencies = [ [[package]] name = "simd-adler32" -version = "0.3.7" +version = "0.3.8" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d66dc143e6b11c1eddc06d5c423cfc97062865baf299914ab64caa38182078fe" +checksum = "e320a6c5ad31d271ad523dcf3ad13e2767ad8b1cb8f047f75a8aeaf8da139da2" [[package]] name = "siphasher" @@ -5882,9 +5882,9 @@ dependencies = [ [[package]] name = "tower-http" -version = "0.6.7" +version = "0.6.8" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9cf146f99d442e8e68e585f5d798ccd3cad9a7835b917e09728880a862706456" +checksum = "d4e6559d53cc268e5031cd8429d05415bc4cb4aefc4aa5d6cc35fbf5b924a1f8" dependencies = [ "bitflags 2.10.0", "bytes", diff --git a/crates/proto/src/domain/account.rs b/crates/proto/src/domain/account.rs index 98fba5961..e5039a582 100644 --- a/crates/proto/src/domain/account.rs +++ b/crates/proto/src/domain/account.rs @@ -15,7 +15,8 @@ use miden_protocol::account::{ use miden_protocol::asset::{Asset, AssetVault}; use miden_protocol::block::BlockNumber; use miden_protocol::block::account_tree::AccountWitness; -use miden_protocol::crypto::merkle::SparseMerklePath; +use miden_protocol::crypto::merkle::smt::{SmtForest, SmtProof}; +use miden_protocol::crypto::merkle::{MerkleError, SparseMerklePath}; use miden_protocol::note::{NoteExecutionMode, NoteTag}; use miden_protocol::utils::{Deserializable, DeserializationError, Serializable}; use thiserror::Error; @@ -192,7 +193,14 @@ impl TryFrom fn try_from( value: proto::rpc::account_storage_details::AccountStorageMapDetails, ) -> Result { - use proto::rpc::account_storage_details::account_storage_map_details::map_entries::StorageMapEntry; + use proto::rpc::account_storage_details::account_storage_map_details::{ + all_map_entries::StorageMapEntry, + map_entries_with_proofs::StorageMapEntryWithProof, + AllMapEntries, + MapEntriesWithProofs, + Entries as ProtoEntries, + }; + let proto::rpc::account_storage_details::AccountStorageMapDetails { slot_name, too_many_entries, @@ -204,26 +212,49 @@ impl TryFrom let entries = if too_many_entries { StorageMapEntries::LimitExceeded } else { - let map_entries = if let Some(entries) = entries { - entries - .entries - .into_iter() - .map(|entry| { - let key = entry - .key - .ok_or(StorageMapEntry::missing_field(stringify!(key)))? - .try_into()?; - let value = entry - .value - .ok_or(StorageMapEntry::missing_field(stringify!(value)))? - .try_into()?; - Ok((key, value)) - }) - .collect::, ConversionError>>()? - } else { - Vec::new() - }; - StorageMapEntries::Entries(map_entries) + match entries { + None => { + return Err(proto::rpc::account_storage_details::AccountStorageMapDetails::missing_field(stringify!(entries))); + }, + Some(ProtoEntries::AllEntries(AllMapEntries { entries })) => { + let entries = entries + .into_iter() + .map(|entry| { + let key = entry + .key + .ok_or(StorageMapEntry::missing_field(stringify!(key)))? + .try_into()?; + let value = entry + .value + .ok_or(StorageMapEntry::missing_field(stringify!(value)))? + .try_into()?; + Ok((key, value)) + }) + .collect::, ConversionError>>()?; + StorageMapEntries::AllEntries(entries) + }, + Some(ProtoEntries::EntriesWithProofs(MapEntriesWithProofs { entries })) => { + let proofs = entries + .into_iter() + .map(|entry| { + let _key: Word = entry + .key + .ok_or(StorageMapEntryWithProof::missing_field(stringify!(key)))? + .try_into()?; + let _value: Word = entry + .value + .ok_or(StorageMapEntryWithProof::missing_field(stringify!(value)))? + .try_into()?; + let smt_opening = entry.proof.ok_or( + StorageMapEntryWithProof::missing_field(stringify!(proof)), + )?; + let smt_proof = SmtProof::try_from(smt_opening)?; + Ok(smt_proof) + }) + .collect::, ConversionError>>()?; + StorageMapEntries::EntriesWithProofs(proofs) + }, + } }; Ok(Self { slot_name, entries }) @@ -256,6 +287,7 @@ impl TryFrom for proto::rpc::AccountVaultDetails { /// returning all entries in a single RPC response creates performance issues. In such cases, /// the `LimitExceeded` variant indicates to the client to use the `SyncStorageMaps` endpoint /// instead. -#[derive(Debug, Clone, PartialEq, Eq)] +#[derive(Debug, Clone, PartialEq)] pub enum StorageMapEntries { /// The map has too many entries to return inline. /// Clients must use `SyncStorageMaps` endpoint instead. LimitExceeded, - /// The storage map entries (key-value pairs), up to `MAX_RETURN_ENTRIES`. - /// TODO: For partial responses, also include Merkle proofs and inner SMT nodes. - Entries(Vec<(Word, Word)>), + /// All storage map entries (key-value pairs) without proofs. + /// Used when all entries are requested for small maps. + AllEntries(Vec<(Word, Word)>), + + /// Specific entries with their SMT proofs for client-side verification. + /// Used when specific keys are requested from the storage map. + EntriesWithProofs(Vec), } -/// Details about an account storage map slot. -#[derive(Debug, Clone, PartialEq, Eq)] +#[derive(Debug, Clone, PartialEq)] pub struct AccountStorageMapDetails { pub slot_name: StorageSlotName, pub entries: StorageMapEntries, @@ -457,46 +492,126 @@ impl AccountStorageMapDetails { /// Maximum number of storage map entries that can be returned in a single response. pub const MAX_RETURN_ENTRIES: usize = 1000; + /// Creates storage map details with all entries from the storage map. + /// + /// If the storage map has too many entries (> `MAX_RETURN_ENTRIES`), + /// returns `LimitExceeded` variant. + pub fn from_all_entries(slot_name: StorageSlotName, storage_map: &StorageMap) -> Self { + if storage_map.num_entries() > Self::MAX_RETURN_ENTRIES { + Self { + slot_name, + entries: StorageMapEntries::LimitExceeded, + } + } else { + let entries = Vec::from_iter(storage_map.entries().map(|(k, v)| (*k, *v))); + Self { + slot_name, + entries: StorageMapEntries::AllEntries(entries), + } + } + } + + /// Creates storage map details based on the requested slot data. + /// + /// Handles both "all entries" and "specific keys" requests. + /// Returns `LimitExceeded` if too many entries. pub fn new(slot_name: StorageSlotName, slot_data: SlotData, storage_map: &StorageMap) -> Self { match slot_data { SlotData::All => Self::from_all_entries(slot_name, storage_map), - SlotData::MapKeys(keys) => Self::from_specific_keys(slot_name, &keys[..], storage_map), + SlotData::MapKeys(keys) => { + if keys.len() > Self::MAX_RETURN_ENTRIES { + Self { + slot_name, + entries: StorageMapEntries::LimitExceeded, + } + } else { + // Query specific keys from the storage map - returns all entries without proofs + // For proofs, use from_specific_keys with SmtForest + let mut entries = Vec::with_capacity(keys.len()); + for key in keys { + let value = storage_map.get(&key); + entries.push((key, value)); + } + Self { + slot_name, + entries: StorageMapEntries::AllEntries(entries), + } + } + }, } } - fn from_all_entries(slot_name: StorageSlotName, storage_map: &StorageMap) -> Self { - if storage_map.num_entries() > Self::MAX_RETURN_ENTRIES { + /// Creates storage map details from forest-queried entries. + /// + /// Returns `LimitExceeded` if too many entries. + pub fn from_forest_entries(slot_name: StorageSlotName, entries: Vec<(Word, Word)>) -> Self { + if entries.len() > Self::MAX_RETURN_ENTRIES { Self { slot_name, entries: StorageMapEntries::LimitExceeded, } } else { - let map_entries = Vec::from_iter(storage_map.entries().map(|(k, v)| (*k, *v))); Self { slot_name, - entries: StorageMapEntries::Entries(map_entries), + entries: StorageMapEntries::AllEntries(entries), } } } - fn from_specific_keys( + /// Creates storage map details with SMT proofs for specific keys. + /// + /// This method queries the forest for specific keys and returns proofs that + /// enable client-side verification of the values. + /// + /// Returns `LimitExceeded` if too many keys, or `MerkleError` if the forest + /// doesn't contain sufficient data. + pub fn from_specific_keys( slot_name: StorageSlotName, keys: &[Word], - storage_map: &StorageMap, - ) -> Self { + storage_forest: &SmtForest, + smt_root: Word, + ) -> Result { if keys.len() > Self::MAX_RETURN_ENTRIES { + return Ok(Self { + slot_name, + entries: StorageMapEntries::LimitExceeded, + }); + } + + // Collect SMT proofs for each key + let mut proofs = Vec::with_capacity(keys.len()); + + for key in keys { + let proof = storage_forest.open(smt_root, *key)?; + proofs.push(proof); + } + + Ok(Self { + slot_name, + entries: StorageMapEntries::EntriesWithProofs(proofs), + }) + } + + /// Creates storage map details from pre-computed SMT proofs. + /// + /// Use this when the caller has already obtained the proofs from an `SmtForest`. + /// Returns `LimitExceeded` if too many proofs are provided. + pub fn from_proofs(slot_name: StorageSlotName, proofs: Vec) -> Self { + if proofs.len() > Self::MAX_RETURN_ENTRIES { Self { slot_name, entries: StorageMapEntries::LimitExceeded, } } else { - // TODO For now, we return all entries instead of specific keys with proofs - Self::from_all_entries(slot_name, storage_map) + Self { + slot_name, + entries: StorageMapEntries::EntriesWithProofs(proofs), + } } } } -#[derive(Debug, Clone, PartialEq, Eq)] +#[derive(Debug, Clone, PartialEq)] pub struct AccountStorageDetails { pub header: AccountStorageHeader, pub map_details: Vec, @@ -659,36 +774,62 @@ impl From for proto::rpc::account_storage_details::AccountStorageMapDetails { fn from(value: AccountStorageMapDetails) -> Self { - use proto::rpc::account_storage_details::account_storage_map_details; + use proto::rpc::account_storage_details::account_storage_map_details::{ + AllMapEntries, + Entries as ProtoEntries, + MapEntriesWithProofs, + }; let AccountStorageMapDetails { slot_name, entries } = value; - match entries { - StorageMapEntries::LimitExceeded => Self { - slot_name: slot_name.to_string(), - too_many_entries: true, - entries: Some(account_storage_map_details::MapEntries { entries: Vec::new() }), + let (too_many_entries, proto_entries) = match entries { + StorageMapEntries::LimitExceeded => (true, None), + StorageMapEntries::AllEntries(entries) => { + let all = AllMapEntries { + entries: Vec::from_iter(entries.into_iter().map(|(key, value)| { + proto::rpc::account_storage_details::account_storage_map_details::all_map_entries::StorageMapEntry { + key: Some(key.into()), + value: Some(value.into()), + } + })), + }; + (false, Some(ProtoEntries::AllEntries(all))) }, - StorageMapEntries::Entries(map_entries) => { - let entries = Some(account_storage_map_details::MapEntries { - entries: Vec::from_iter(map_entries.into_iter().map(|(key, value)| { - account_storage_map_details::map_entries::StorageMapEntry { + StorageMapEntries::EntriesWithProofs(proofs) => { + use miden_protocol::crypto::merkle::smt::SmtLeaf; + + let with_proofs = MapEntriesWithProofs { + entries: Vec::from_iter(proofs.into_iter().map(|proof| { + // Get key/value from the leaf before consuming the proof + let (key, value) = match proof.leaf() { + SmtLeaf::Empty(_) => { + (miden_protocol::EMPTY_WORD, miden_protocol::EMPTY_WORD) + }, + SmtLeaf::Single((k, v)) => (*k, *v), + SmtLeaf::Multiple(entries) => entries.iter().next().map_or( + (miden_protocol::EMPTY_WORD, miden_protocol::EMPTY_WORD), + |(k, v)| (*k, *v), + ), + }; + let smt_opening = proto::primitives::SmtOpening::from(proof); + proto::rpc::account_storage_details::account_storage_map_details::map_entries_with_proofs::StorageMapEntryWithProof { key: Some(key.into()), value: Some(value.into()), + proof: Some(smt_opening), } })), - }); - - Self { - slot_name: slot_name.to_string(), - too_many_entries: false, - entries, - } + }; + (false, Some(ProtoEntries::EntriesWithProofs(with_proofs))) }, + }; + + Self { + slot_name: slot_name.to_string(), + too_many_entries, + entries: proto_entries, } } } - // ACCOUNT WITNESS // ================================================================================================ @@ -951,3 +1092,173 @@ pub enum NetworkAccountError { fn get_account_id_tag_prefix(id: AccountId) -> AccountPrefix { (id.prefix().as_u64() >> 34) as AccountPrefix } + +#[cfg(test)] +mod tests { + use miden_protocol::crypto::merkle::EmptySubtreeRoots; + use miden_protocol::crypto::merkle::smt::SMT_DEPTH; + + use super::*; + + fn word_from_u32(arr: [u32; 4]) -> Word { + Word::from(arr) + } + + fn test_slot_name() -> StorageSlotName { + StorageSlotName::new("miden::test::storage::slot").unwrap() + } + + fn empty_smt_root() -> Word { + *EmptySubtreeRoots::entry(SMT_DEPTH, 0) + } + + #[test] + fn account_storage_map_details_from_forest_entries() { + let slot_name = test_slot_name(); + let entries = vec![ + (word_from_u32([1, 2, 3, 4]), word_from_u32([5, 6, 7, 8])), + (word_from_u32([9, 10, 11, 12]), word_from_u32([13, 14, 15, 16])), + ]; + + let details = + AccountStorageMapDetails::from_forest_entries(slot_name.clone(), entries.clone()); + + assert_eq!(details.slot_name, slot_name); + assert_eq!(details.entries, StorageMapEntries::AllEntries(entries)); + } + + #[test] + fn account_storage_map_details_from_forest_entries_limit_exceeded() { + let slot_name = test_slot_name(); + // Create more entries than MAX_RETURN_ENTRIES + let entries: Vec<_> = (0..=AccountStorageMapDetails::MAX_RETURN_ENTRIES) + .map(|i| { + let key = word_from_u32([i as u32, 0, 0, 0]); + let value = word_from_u32([0, 0, 0, i as u32]); + (key, value) + }) + .collect(); + + let details = AccountStorageMapDetails::from_forest_entries(slot_name.clone(), entries); + + assert_eq!(details.slot_name, slot_name); + assert_eq!(details.entries, StorageMapEntries::LimitExceeded); + } + + #[test] + fn account_storage_map_details_from_specific_keys() { + let slot_name = test_slot_name(); + + // Create an SmtForest and populate it with some data + let mut forest = SmtForest::new(); + let entries = [ + (word_from_u32([1, 0, 0, 0]), word_from_u32([10, 0, 0, 0])), + (word_from_u32([2, 0, 0, 0]), word_from_u32([20, 0, 0, 0])), + (word_from_u32([3, 0, 0, 0]), word_from_u32([30, 0, 0, 0])), + ]; + + // Insert entries into the forest starting from an empty root + let smt_root = forest.batch_insert(empty_smt_root(), entries.iter().copied()).unwrap(); + + // Query specific keys + let keys = vec![word_from_u32([1, 0, 0, 0]), word_from_u32([3, 0, 0, 0])]; + + let details = AccountStorageMapDetails::from_specific_keys( + slot_name.clone(), + &keys, + &forest, + smt_root, + ) + .unwrap(); + + assert_eq!(details.slot_name, slot_name); + match details.entries { + StorageMapEntries::EntriesWithProofs(proofs) => { + use miden_protocol::crypto::merkle::smt::SmtLeaf; + + assert_eq!(proofs.len(), 2); + + // Helper to extract key-value from any leaf type + let get_value = |proof: &SmtProof, expected_key: Word| -> Word { + match proof.leaf() { + SmtLeaf::Single((k, v)) if *k == expected_key => *v, + SmtLeaf::Multiple(entries) => entries + .iter() + .find(|(k, _)| *k == expected_key) + .map_or(miden_protocol::EMPTY_WORD, |(_, v)| *v), + _ => miden_protocol::EMPTY_WORD, + } + }; + + let first_key = word_from_u32([1, 0, 0, 0]); + let second_key = word_from_u32([3, 0, 0, 0]); + let first_value = get_value(&proofs[0], first_key); + let second_value = get_value(&proofs[1], second_key); + + assert_eq!(first_value, word_from_u32([10, 0, 0, 0])); + assert_eq!(second_value, word_from_u32([30, 0, 0, 0])); + }, + StorageMapEntries::LimitExceeded | StorageMapEntries::AllEntries(_) => { + panic!("Expected EntriesWithProofs") + }, + } + } + + #[test] + fn account_storage_map_details_from_specific_keys_nonexistent_returns_proof() { + let slot_name = test_slot_name(); + + // Create an SmtForest with one entry so the root is tracked + let mut forest = SmtForest::new(); + let entries = [(word_from_u32([1, 0, 0, 0]), word_from_u32([10, 0, 0, 0]))]; + let smt_root = forest.batch_insert(empty_smt_root(), entries.iter().copied()).unwrap(); + + // Query a key that doesn't exist in the tree - should return a proof + // (the proof will show non-membership or point to an adjacent leaf) + let keys = vec![word_from_u32([99, 0, 0, 0])]; + + let details = AccountStorageMapDetails::from_specific_keys( + slot_name.clone(), + &keys, + &forest, + smt_root, + ) + .unwrap(); + + match details.entries { + StorageMapEntries::EntriesWithProofs(proofs) => { + // We got a proof for the non-existent key + assert_eq!(proofs.len(), 1); + // The proof exists and can be used to verify non-membership + }, + StorageMapEntries::LimitExceeded | StorageMapEntries::AllEntries(_) => { + panic!("Expected EntriesWithProofs") + }, + } + } + + #[test] + fn account_storage_map_details_from_specific_keys_limit_exceeded() { + let slot_name = test_slot_name(); + let mut forest = SmtForest::new(); + + // Create a forest with some data to get a valid root + let entries = [(word_from_u32([1, 0, 0, 0]), word_from_u32([10, 0, 0, 0]))]; + let smt_root = forest.batch_insert(empty_smt_root(), entries.iter().copied()).unwrap(); + + // Create more keys than MAX_RETURN_ENTRIES + let keys: Vec<_> = (0..=AccountStorageMapDetails::MAX_RETURN_ENTRIES) + .map(|i| word_from_u32([i as u32, 0, 0, 0])) + .collect(); + + let details = AccountStorageMapDetails::from_specific_keys( + slot_name.clone(), + &keys, + &forest, + smt_root, + ) + .unwrap(); + + assert_eq!(details.entries, StorageMapEntries::LimitExceeded); + } +} diff --git a/crates/proto/src/generated/rpc.rs b/crates/proto/src/generated/rpc.rs index f9a59e39c..755009e2c 100644 --- a/crates/proto/src/generated/rpc.rs +++ b/crates/proto/src/generated/rpc.rs @@ -233,25 +233,53 @@ pub mod account_storage_details { /// Storage slot name. #[prost(string, tag = "1")] pub slot_name: ::prost::alloc::string::String, - /// A flag that is set to `true` if the number of to-be-returned entries in the - /// storage map would exceed a threshold. This indicates to the user that `SyncStorageMaps` - /// endpoint should be used to get all storage map data. + /// True when the number of entries exceeds the response limit. + /// When set, clients should use the `SyncStorageMaps` endpoint. #[prost(bool, tag = "2")] pub too_many_entries: bool, - /// By default we provide all storage entries. - #[prost(message, optional, tag = "3")] - pub entries: ::core::option::Option, + /// The map entries (with or without proofs). Empty when too_many_entries is true. + #[prost(oneof = "account_storage_map_details::Entries", tags = "3, 4")] + pub entries: ::core::option::Option, } /// Nested message and enum types in `AccountStorageMapDetails`. pub mod account_storage_map_details { - /// Wrapper for repeated storage map entries + /// Wrapper for repeated storage map entries including their proofs. + /// Used when specific keys are requested to enable client-side verification. #[derive(Clone, PartialEq, ::prost::Message)] - pub struct MapEntries { + pub struct MapEntriesWithProofs { #[prost(message, repeated, tag = "1")] - pub entries: ::prost::alloc::vec::Vec, + pub entries: ::prost::alloc::vec::Vec< + map_entries_with_proofs::StorageMapEntryWithProof, + >, } - /// Nested message and enum types in `MapEntries`. - pub mod map_entries { + /// Nested message and enum types in `MapEntriesWithProofs`. + pub mod map_entries_with_proofs { + /// Definition of individual storage entries including a proof. + #[derive(Clone, PartialEq, ::prost::Message)] + pub struct StorageMapEntryWithProof { + #[prost(message, optional, tag = "1")] + pub key: ::core::option::Option< + super::super::super::super::primitives::Digest, + >, + #[prost(message, optional, tag = "2")] + pub value: ::core::option::Option< + super::super::super::super::primitives::Digest, + >, + #[prost(message, optional, tag = "3")] + pub proof: ::core::option::Option< + super::super::super::super::primitives::SmtOpening, + >, + } + } + /// Wrapper for repeated storage map entries (without proofs). + /// Used when all entries are requested for small maps. + #[derive(Clone, PartialEq, ::prost::Message)] + pub struct AllMapEntries { + #[prost(message, repeated, tag = "1")] + pub entries: ::prost::alloc::vec::Vec, + } + /// Nested message and enum types in `AllMapEntries`. + pub mod all_map_entries { /// Definition of individual storage entries. #[derive(Clone, Copy, PartialEq, Eq, Hash, ::prost::Message)] pub struct StorageMapEntry { @@ -265,6 +293,16 @@ pub mod account_storage_details { >, } } + /// The map entries (with or without proofs). Empty when too_many_entries is true. + #[derive(Clone, PartialEq, ::prost::Oneof)] + pub enum Entries { + /// All storage entries without proofs (for small maps or full requests). + #[prost(message, tag = "3")] + AllEntries(AllMapEntries), + /// Specific entries with their SMT proofs (for partial requests). + #[prost(message, tag = "4")] + EntriesWithProofs(MapEntriesWithProofs), + } } } /// List of nullifiers to return proofs for. diff --git a/crates/store/src/db/mod.rs b/crates/store/src/db/mod.rs index 53f7c2397..bd171c773 100644 --- a/crates/store/src/db/mod.rs +++ b/crates/store/src/db/mod.rs @@ -7,7 +7,7 @@ use diesel::{Connection, RunQueryDsl, SqliteConnection}; use miden_node_proto::domain::account::{AccountInfo, AccountSummary, NetworkAccountPrefix}; use miden_node_proto::generated as proto; use miden_protocol::Word; -use miden_protocol::account::{AccountHeader, AccountId, AccountStorage}; +use miden_protocol::account::{AccountHeader, AccountId, AccountStorageHeader}; use miden_protocol::asset::{Asset, AssetVaultKey}; use miden_protocol::block::{BlockHeader, BlockNoteIndex, BlockNumber, ProvenBlock}; use miden_protocol::crypto::merkle::SparseMerklePath; @@ -435,23 +435,6 @@ impl Db { .await } - /// Reconstructs account storage at a specific block from the database - /// - /// This method queries the decomposed storage tables and reconstructs the full - /// `AccountStorage` with SMT backing for Map slots. - // TODO split querying the header from the content - #[instrument(level = "debug", target = COMPONENT, skip_all, ret(level = "debug"), err)] - pub async fn select_account_storage_at_block( - &self, - account_id: AccountId, - block_num: BlockNumber, - ) -> Result { - self.transact("Get account storage at block", move |conn| { - queries::select_account_storage_at_block(conn, account_id, block_num) - }) - .await - } - /// Queries vault assets at a specific block #[instrument(level = "debug", target = COMPONENT, skip_all, ret(level = "debug"), err)] pub async fn select_account_vault_at_block( @@ -478,16 +461,17 @@ impl Db { .await } - /// Queries the account header for a specific account at a specific block number. + /// Queries the account header and storage header for a specific account at a block. /// + /// Returns both in a single query to avoid querying the database twice. /// Returns `None` if the account doesn't exist at that block. - pub async fn select_account_header_at_block( + pub async fn select_account_header_with_storage_header_at_block( &self, account_id: AccountId, block_num: BlockNumber, - ) -> Result> { - self.transact("Get account header at block", move |conn| { - queries::select_account_header_at_block(conn, account_id, block_num) + ) -> Result> { + self.transact("Get account header with storage header at block", move |conn| { + queries::select_account_header_with_storage_header_at_block(conn, account_id, block_num) }) .await } diff --git a/crates/store/src/db/models/queries/accounts.rs b/crates/store/src/db/models/queries/accounts.rs index 8f3deae7e..e20d73154 100644 --- a/crates/store/src/db/models/queries/accounts.rs +++ b/crates/store/src/db/models/queries/accounts.rs @@ -50,8 +50,7 @@ use crate::errors::DatabaseError; mod at_block; pub(crate) use at_block::{ - select_account_header_at_block, - select_account_storage_at_block, + select_account_header_with_storage_header_at_block, select_account_vault_at_block, }; diff --git a/crates/store/src/db/models/queries/accounts/at_block.rs b/crates/store/src/db/models/queries/accounts/at_block.rs index 021714abe..13ed84177 100644 --- a/crates/store/src/db/models/queries/accounts/at_block.rs +++ b/crates/store/src/db/models/queries/accounts/at_block.rs @@ -1,5 +1,3 @@ -use std::collections::BTreeMap; - use diesel::prelude::Queryable; use diesel::query_dsl::methods::SelectDsl; use diesel::{ @@ -10,16 +8,7 @@ use diesel::{ RunQueryDsl, SqliteConnection, }; -use miden_protocol::account::{ - AccountHeader, - AccountId, - AccountStorage, - AccountStorageHeader, - StorageMap, - StorageSlot, - StorageSlotName, - StorageSlotType, -}; +use miden_protocol::account::{AccountHeader, AccountId, AccountStorageHeader}; use miden_protocol::asset::Asset; use miden_protocol::block::BlockNumber; use miden_protocol::utils::{Deserializable, Serializable}; @@ -54,14 +43,14 @@ struct AccountHeaderDataRaw { /// /// # Returns /// -/// * `Ok(Some(AccountHeader))` - The account header if found +/// * `Ok(Some((AccountHeader, AccountStorageHeader)))` - The headers if found /// * `Ok(None)` - If account doesn't exist at that block /// * `Err(DatabaseError)` - If there's a database error -pub(crate) fn select_account_header_at_block( +pub(crate) fn select_account_header_with_storage_header_at_block( conn: &mut SqliteConnection, account_id: AccountId, block_num: BlockNumber, -) -> Result, DatabaseError> { +) -> Result, DatabaseError> { use schema::accounts; let account_id_bytes = account_id.to_bytes(); @@ -93,14 +82,13 @@ pub(crate) fn select_account_header_at_block( return Ok(None); }; - let storage_commitment = match storage_header_blob { - Some(blob) => { - let header = AccountStorageHeader::read_from_bytes(&blob)?; - header.to_commitment() - }, - None => Word::default(), + let storage_header = match &storage_header_blob { + Some(blob) => AccountStorageHeader::read_from_bytes(blob)?, + None => AccountStorageHeader::new(Vec::new())?, }; + let storage_commitment = storage_header.to_commitment(); + let code_commitment = code_commitment_bytes .map(|bytes| Word::read_from_bytes(&bytes)) .transpose()? @@ -113,13 +101,10 @@ pub(crate) fn select_account_header_at_block( .transpose()? .unwrap_or(Word::default()); - Ok(Some(AccountHeader::new( - account_id, - nonce, - vault_root, - storage_commitment, - code_commitment, - ))) + let account_header = + AccountHeader::new(account_id, nonce, vault_root, storage_commitment, code_commitment); + + Ok(Some((account_header, storage_header))) } // ACCOUNT VAULT @@ -180,89 +165,3 @@ pub(crate) fn select_account_vault_at_block( Ok(assets) } - -// ACCOUNT STORAGE -// ================================================================================================ - -/// Returns account storage at a given block by reading from `accounts.storage_header` -/// (which contains the `AccountStorageHeader`) and reconstructing full storage from -/// map values in `account_storage_map_values` table. -pub(crate) fn select_account_storage_at_block( - conn: &mut SqliteConnection, - account_id: AccountId, - block_num: BlockNumber, -) -> Result { - use schema::account_storage_map_values as t; - - let account_id_bytes = account_id.to_bytes(); - let block_num_sql = block_num.to_raw_sql(); - - // Query storage header blob for this account at or before this block - let storage_blob: Option> = - SelectDsl::select(schema::accounts::table, schema::accounts::storage_header) - .filter(schema::accounts::account_id.eq(&account_id_bytes)) - .filter(schema::accounts::block_num.le(block_num_sql)) - .order(schema::accounts::block_num.desc()) - .limit(1) - .first(conn) - .optional()? - .flatten(); - - let Some(blob) = storage_blob else { - // No storage means empty storage - return Ok(AccountStorage::new(Vec::new())?); - }; - - // Deserialize the AccountStorageHeader from the blob - let header = AccountStorageHeader::read_from_bytes(&blob)?; - - // Query all map values for this account up to and including this block. - // For each (slot_name, key), we need the latest value at or before block_num. - // First, get all entries up to block_num - let map_values: Vec<(i64, String, Vec, Vec)> = - SelectDsl::select(t::table, (t::block_num, t::slot_name, t::key, t::value)) - .filter(t::account_id.eq(&account_id_bytes).and(t::block_num.le(block_num_sql))) - .order((t::slot_name.asc(), t::key.asc(), t::block_num.desc())) - .load(conn)?; - - // For each (slot_name, key) pair, keep only the latest entry (highest block_num) - let mut latest_map_entries: BTreeMap<(StorageSlotName, Word), Word> = BTreeMap::new(); - - for (_, slot_name_str, key_bytes, value_bytes) in map_values { - let slot_name: StorageSlotName = slot_name_str.parse().map_err(|_| { - DatabaseError::DataCorrupted(format!("Invalid slot name: {slot_name_str}")) - })?; - let key = Word::read_from_bytes(&key_bytes)?; - let value = Word::read_from_bytes(&value_bytes)?; - - // Only insert if we haven't seen this (slot_name, key) yet - // (since results are ordered by block_num desc, first one is latest) - latest_map_entries.entry((slot_name, key)).or_insert(value); - } - - // Group entries by slot name - let mut map_entries_by_slot: BTreeMap> = BTreeMap::new(); - for ((slot_name, key), value) in latest_map_entries { - map_entries_by_slot.entry(slot_name).or_default().push((key, value)); - } - - // Reconstruct StorageSlots from header slots + map entries - let mut slots = Vec::new(); - for slot_header in header.slots() { - let slot = match slot_header.slot_type() { - StorageSlotType::Value => { - // For value slots, the header value IS the slot value - StorageSlot::with_value(slot_header.name().clone(), slot_header.value()) - }, - StorageSlotType::Map => { - // For map slots, reconstruct from map entries - let entries = map_entries_by_slot.remove(slot_header.name()).unwrap_or_default(); - let storage_map = StorageMap::with_entries(entries)?; - StorageSlot::with_map(slot_header.name().clone(), storage_map) - }, - }; - slots.push(slot); - } - - Ok(AccountStorage::new(slots)?) -} diff --git a/crates/store/src/db/models/queries/accounts/tests.rs b/crates/store/src/db/models/queries/accounts/tests.rs index 6f9f5b075..8532d448f 100644 --- a/crates/store/src/db/models/queries/accounts/tests.rs +++ b/crates/store/src/db/models/queries/accounts/tests.rs @@ -1,7 +1,16 @@ //! Tests for the `accounts` module, specifically for account storage and historical queries. +use std::collections::BTreeMap; + use diesel::query_dsl::methods::SelectDsl; -use diesel::{Connection, OptionalExtension, QueryDsl, RunQueryDsl}; +use diesel::{ + BoolExpressionMethods, + Connection, + ExpressionMethods, + OptionalExtension, + QueryDsl, + RunQueryDsl, +}; use diesel_migrations::MigrationHarness; use miden_node_utils::fee::test_fee_params; use miden_protocol::account::auth::PublicKeyCommitment; @@ -13,20 +22,27 @@ use miden_protocol::account::{ AccountDelta, AccountId, AccountIdVersion, + AccountStorage, + AccountStorageHeader, AccountStorageMode, AccountType, + StorageMap, StorageSlot, StorageSlotName, + StorageSlotType, }; use miden_protocol::block::{BlockAccountUpdate, BlockHeader, BlockNumber}; use miden_protocol::crypto::dsa::ecdsa_k256_keccak::SecretKey; -use miden_protocol::utils::Serializable; +use miden_protocol::utils::{Deserializable, Serializable}; use miden_protocol::{EMPTY_WORD, Felt, Word}; use miden_standards::account::auth::AuthRpoFalcon512; use miden_standards::code_builder::CodeBuilder; use super::*; use crate::db::migrations::MIGRATIONS; +use crate::db::models::conv::SqlTypeConvert; +use crate::db::schema; +use crate::errors::DatabaseError; fn setup_test_db() -> SqliteConnection { let mut conn = @@ -37,6 +53,80 @@ fn setup_test_db() -> SqliteConnection { conn } +/// Test helper: reconstructs account storage at a given block from DB. +/// +/// Reads `accounts.storage_header` and `account_storage_map_values` to reconstruct +/// the full `AccountStorage` at the specified block. +fn test_select_account_storage_at_block( + conn: &mut SqliteConnection, + account_id: AccountId, + block_num: BlockNumber, +) -> Result { + use schema::account_storage_map_values as t; + + let account_id_bytes = account_id.to_bytes(); + let block_num_sql = block_num.to_raw_sql(); + + // Query storage header blob for this account at or before this block + let storage_blob: Option> = + SelectDsl::select(schema::accounts::table, schema::accounts::storage_header) + .filter(schema::accounts::account_id.eq(&account_id_bytes)) + .filter(schema::accounts::block_num.le(block_num_sql)) + .order(schema::accounts::block_num.desc()) + .limit(1) + .first(conn) + .optional()? + .flatten(); + + let Some(blob) = storage_blob else { + return Ok(AccountStorage::new(Vec::new())?); + }; + + let header = AccountStorageHeader::read_from_bytes(&blob)?; + + // Query all map values for this account up to and including this block. + let map_values: Vec<(i64, String, Vec, Vec)> = + SelectDsl::select(t::table, (t::block_num, t::slot_name, t::key, t::value)) + .filter(t::account_id.eq(&account_id_bytes).and(t::block_num.le(block_num_sql))) + .order((t::slot_name.asc(), t::key.asc(), t::block_num.desc())) + .load(conn)?; + + // For each (slot_name, key) pair, keep only the latest entry + let mut latest_map_entries: BTreeMap<(StorageSlotName, Word), Word> = BTreeMap::new(); + for (_, slot_name_str, key_bytes, value_bytes) in map_values { + let slot_name: StorageSlotName = slot_name_str.parse().map_err(|_| { + DatabaseError::DataCorrupted(format!("Invalid slot name: {slot_name_str}")) + })?; + let key = Word::read_from_bytes(&key_bytes)?; + let value = Word::read_from_bytes(&value_bytes)?; + latest_map_entries.entry((slot_name, key)).or_insert(value); + } + + // Group entries by slot name + let mut map_entries_by_slot: BTreeMap> = BTreeMap::new(); + for ((slot_name, key), value) in latest_map_entries { + map_entries_by_slot.entry(slot_name).or_default().push((key, value)); + } + + // Reconstruct StorageSlots from header slots + map entries + let mut slots = Vec::new(); + for slot_header in header.slots() { + let slot = match slot_header.slot_type() { + StorageSlotType::Value => { + StorageSlot::with_value(slot_header.name().clone(), slot_header.value()) + }, + StorageSlotType::Map => { + let entries = map_entries_by_slot.remove(slot_header.name()).unwrap_or_default(); + let storage_map = StorageMap::with_entries(entries)?; + StorageSlot::with_map(slot_header.name().clone(), storage_map) + }, + }; + slots.push(slot); + } + + Ok(AccountStorage::new(slots)?) +} + fn create_test_account_with_storage() -> (Account, AccountId) { // Create a simple public account with one value storage slot let account_id = AccountId::dummy( @@ -112,8 +202,9 @@ fn test_select_account_header_at_block_returns_none_for_nonexistent() { ); // Query for a non-existent account - let result = select_account_header_at_block(&mut conn, account_id, block_num) - .expect("Query should succeed"); + let result = + select_account_header_with_storage_header_at_block(&mut conn, account_id, block_num) + .expect("Query should succeed"); assert!(result.is_none(), "Should return None for non-existent account"); } @@ -138,9 +229,10 @@ fn test_select_account_header_at_block_returns_correct_header() { upsert_accounts(&mut conn, &[account_update], block_num).expect("upsert_accounts failed"); // Query the account header - let header = select_account_header_at_block(&mut conn, account_id, block_num) - .expect("Query should succeed") - .expect("Header should exist"); + let (header, _storage_header) = + select_account_header_with_storage_header_at_block(&mut conn, account_id, block_num) + .expect("Query should succeed") + .expect("Header should exist"); assert_eq!(header.id(), account_id, "Account ID should match"); assert_eq!(header.nonce(), account.nonce(), "Nonce should match"); @@ -174,16 +266,18 @@ fn test_select_account_header_at_block_historical_query() { upsert_accounts(&mut conn, &[account_update_1], block_num_1).expect("First upsert failed"); // Query at block 1 - should return the account - let header_1 = select_account_header_at_block(&mut conn, account_id, block_num_1) - .expect("Query should succeed") - .expect("Header should exist at block 1"); + let (header_1, _) = + select_account_header_with_storage_header_at_block(&mut conn, account_id, block_num_1) + .expect("Query should succeed") + .expect("Header should exist at block 1"); assert_eq!(header_1.nonce(), nonce_1, "Nonce at block 1 should match"); // Query at block 2 - should return the same account (most recent before block 2) - let header_2 = select_account_header_at_block(&mut conn, account_id, block_num_2) - .expect("Query should succeed") - .expect("Header should exist at block 2"); + let (header_2, _) = + select_account_header_with_storage_header_at_block(&mut conn, account_id, block_num_2) + .expect("Query should succeed") + .expect("Header should exist at block 2"); assert_eq!(header_2.nonce(), nonce_1, "Nonce at block 2 should match block 1"); } @@ -242,7 +336,7 @@ fn test_select_account_storage_at_block_returns_storage() { upsert_accounts(&mut conn, &[account_update], block_num).expect("upsert_accounts failed"); // Query storage - let storage = select_account_storage_at_block(&mut conn, account_id, block_num) + let storage = test_select_account_storage_at_block(&mut conn, account_id, block_num) .expect("Query should succeed"); assert_eq!( @@ -396,8 +490,9 @@ fn test_upsert_accounts_updates_is_latest_flag() { ); // Verify historical query returns first update - let storage_at_block_1 = select_account_storage_at_block(&mut conn, account_id, block_num_1) - .expect("Failed to query storage at block 1"); + let storage_at_block_1 = + test_select_account_storage_at_block(&mut conn, account_id, block_num_1) + .expect("Failed to query storage at block 1"); assert_eq!( storage_at_block_1.to_commitment(), diff --git a/crates/store/src/errors.rs b/crates/store/src/errors.rs index 42a0fe32d..09c5798a3 100644 --- a/crates/store/src/errors.rs +++ b/crates/store/src/errors.rs @@ -131,6 +131,12 @@ pub enum DatabaseError { SqlValueConversion(#[from] DatabaseTypeConversionError), #[error("Not implemented: {0}")] NotImplemented(String), + #[error("storage root not found for account {account_id}, slot {slot_name}, block {block_num}")] + StorageRootNotFound { + account_id: AccountId, + slot_name: String, + block_num: BlockNumber, + }, } impl DatabaseError { diff --git a/crates/store/src/inner_forest/mod.rs b/crates/store/src/inner_forest/mod.rs index 5778091af..058e2c5d8 100644 --- a/crates/store/src/inner_forest/mod.rs +++ b/crates/store/src/inner_forest/mod.rs @@ -1,11 +1,12 @@ use std::collections::BTreeMap; +use miden_node_proto::domain::account::{AccountStorageMapDetails, StorageMapEntries}; use miden_protocol::account::delta::{AccountDelta, AccountStorageDelta, AccountVaultDelta}; use miden_protocol::account::{AccountId, NonFungibleDeltaAction, StorageSlotName}; use miden_protocol::asset::{Asset, FungibleAsset}; use miden_protocol::block::BlockNumber; -use miden_protocol::crypto::merkle::EmptySubtreeRoots; -use miden_protocol::crypto::merkle::smt::{SMT_DEPTH, SmtForest}; +use miden_protocol::crypto::merkle::smt::{SMT_DEPTH, SmtForest, SmtProof}; +use miden_protocol::crypto::merkle::{EmptySubtreeRoots, MerkleError}; use miden_protocol::{EMPTY_WORD, Word}; #[cfg(test)] @@ -24,6 +25,10 @@ pub(crate) struct InnerForest { /// Populated during block import for all storage map slots. storage_roots: BTreeMap<(AccountId, StorageSlotName, BlockNumber), Word>, + /// Maps (`account_id`, `slot_name`, `block_num`) to all key-value entries in that storage map. + /// Accumulated from deltas - each block's entries include all entries up to that point. + storage_entries: BTreeMap<(AccountId, StorageSlotName, BlockNumber), BTreeMap>, + /// Maps (`account_id`, `block_num`) to vault SMT root. /// Tracks asset vault versions across all blocks with structural sharing. vault_roots: BTreeMap<(AccountId, BlockNumber), Word>, @@ -34,6 +39,7 @@ impl InnerForest { Self { forest: SmtForest::new(), storage_roots: BTreeMap::new(), + storage_entries: BTreeMap::new(), vault_roots: BTreeMap::new(), } } @@ -82,6 +88,47 @@ impl InnerForest { .map_or_else(Self::empty_smt_root, |(_, root)| *root) } + /// Opens a storage map and returns storage map details with SMT proofs for the given keys. + /// + /// Returns `None` if no storage root is tracked for this account/slot/block combination. + /// Returns a `MerkleError` if the forest doesn't contain sufficient data for the proofs. + pub(crate) fn open_storage_map( + &self, + account_id: AccountId, + slot_name: StorageSlotName, + block_num: BlockNumber, + keys: &[Word], + ) -> Option> { + let root = *self.storage_roots.get(&(account_id, slot_name.clone(), block_num))?; + + let proofs: Result, MerkleError> = + keys.iter().map(|key| self.forest.open(root, *key)).collect(); + + Some(proofs.map(|p| AccountStorageMapDetails::from_proofs(slot_name, p))) + } + + /// Returns all key-value entries for a specific account storage slot at a block. + /// + /// Returns `None` if no entries are tracked for this account/slot/block combination. + /// Returns an error if there are too many entries to return. + pub(crate) fn storage_map_entries( + &self, + account_id: AccountId, + slot_name: StorageSlotName, + block_num: BlockNumber, + ) -> Option { + let entries = self.storage_entries.get(&(account_id, slot_name.clone(), block_num))?; + if entries.len() > AccountStorageMapDetails::MAX_RETURN_ENTRIES { + return Some(AccountStorageMapDetails { + slot_name, + entries: StorageMapEntries::LimitExceeded, + }); + } + let entries = Vec::from_iter(entries.iter().map(|(k, v)| (*k, *v))); + + Some(AccountStorageMapDetails::from_forest_entries(slot_name.clone(), entries)) + } + // PUBLIC INTERFACE // -------------------------------------------------------------------------------------------- @@ -223,7 +270,7 @@ impl InnerForest { /// Updates the forest with storage map changes from a delta. /// /// Processes storage map slot deltas, building SMTs for each modified slot - /// and tracking the new roots. + /// and tracking the new roots and accumulated entries. fn update_account_storage( &mut self, block_num: BlockNumber, @@ -240,27 +287,49 @@ impl InnerForest { self.get_storage_root(account_id, slot_name, parent_block) }; - let entries: Vec<_> = + let delta_entries: Vec<_> = map_delta.entries().iter().map(|(key, value)| ((*key).into(), *value)).collect(); - if entries.is_empty() { + if delta_entries.is_empty() { continue; } let updated_root = self .forest - .batch_insert(prev_root, entries.iter().copied()) + .batch_insert(prev_root, delta_entries.iter().copied()) .expect("forest insertion should succeed"); self.storage_roots .insert((account_id, slot_name.clone(), block_num), updated_root); + // Accumulate entries: start from parent block's entries or empty for full state + let mut accumulated_entries = if is_full_state { + BTreeMap::new() + } else { + self.storage_entries + .get(&(account_id, slot_name.clone(), parent_block)) + .cloned() + .unwrap_or_default() + }; + + // Apply delta entries (insert or remove if value is EMPTY_WORD) + for (key, value) in &delta_entries { + if *value == EMPTY_WORD { + accumulated_entries.remove(key); + } else { + accumulated_entries.insert(*key, *value); + } + } + + self.storage_entries + .insert((account_id, slot_name.clone(), block_num), accumulated_entries); + tracing::debug!( target: crate::COMPONENT, %account_id, %block_num, ?slot_name, - entries = entries.len(), + delta_entries = delta_entries.len(), "Updated storage map in forest" ); } diff --git a/crates/store/src/state.rs b/crates/store/src/state.rs index c7f935c38..d4ff7c45f 100644 --- a/crates/store/src/state.rs +++ b/crates/store/src/state.rs @@ -18,13 +18,15 @@ use miden_node_proto::domain::account::{ AccountStorageMapDetails, AccountVaultDetails, NetworkAccountPrefix, + SlotData, StorageMapRequest, }; use miden_node_proto::domain::batch::BatchInputs; use miden_node_utils::ErrorReport; use miden_node_utils::formatting::format_array; +use miden_protocol::Word; +use miden_protocol::account::AccountId; use miden_protocol::account::delta::AccountUpdateDetails; -use miden_protocol::account::{AccountId, StorageSlotContent}; use miden_protocol::block::account_tree::{AccountTree, AccountWitness, account_id_to_smt_key}; use miden_protocol::block::nullifier_tree::{NullifierTree, NullifierWitness}; use miden_protocol::block::{BlockHeader, BlockInputs, BlockNumber, Blockchain, ProvenBlock}; @@ -39,7 +41,6 @@ use miden_protocol::crypto::merkle::smt::{ use miden_protocol::note::{NoteDetails, NoteId, NoteScript, Nullifier}; use miden_protocol::transaction::{OutputNote, PartialBlockchain}; use miden_protocol::utils::Serializable; -use miden_protocol::{AccountError, Word}; use tokio::sync::{Mutex, RwLock, oneshot}; use tracing::{info, info_span, instrument}; @@ -1003,6 +1004,10 @@ impl State { /// /// This method queries the database to fetch the account state and processes the detail /// request to return only the requested information. + /// + /// For specific key queries (`SlotData::MapKeys`), the forest is used to provide SMT proofs. + /// Returns an error if the forest doesn't have data for the requested slot. + /// All-entries queries (`SlotData::All`) use the forest to return all entries. async fn fetch_public_account_details( &self, account_id: AccountId, @@ -1022,11 +1027,12 @@ impl State { // Validate block exists in the blockchain before querying the database self.validate_block_exists(block_num).await?; - let account_header = - self.db - .select_account_header_at_block(account_id, block_num) - .await? - .ok_or(DatabaseError::AccountAtBlockHeightNotFoundInDb(account_id, block_num))?; + // Query account header and storage header together in a single DB call + let (account_header, storage_header) = self + .db + .select_account_header_with_storage_header_at_block(account_id, block_num) + .await? + .ok_or(DatabaseError::AccountAtBlockHeightNotFoundInDb(account_id, block_num))?; let account_code = match code_commitment { Some(commitment) if commitment == account_header.code_commitment() => None, @@ -1050,25 +1056,31 @@ impl State { None => AccountVaultDetails::empty(), }; - // TODO: don't load the entire store at once, load what is required - let store = self.db.select_account_storage_at_block(account_id, block_num).await?; - let storage_header = store.to_header(); let mut storage_map_details = Vec::::with_capacity(storage_requests.len()); - for StorageMapRequest { slot_name, slot_data } in storage_requests { - let Some(slot) = store.slots().iter().find(|s| s.name() == &slot_name) else { - continue; - }; + // Use forest for storage map queries + let forest_guard = self.forest.read().await; - let storage_map = match slot.content() { - StorageSlotContent::Map(map) => map, - StorageSlotContent::Value(_) => { - return Err(AccountError::StorageSlotNotMap(slot_name).into()); - }, + for StorageMapRequest { slot_name, slot_data } in storage_requests { + let details = match &slot_data { + SlotData::MapKeys(keys) => forest_guard + .open_storage_map(account_id, slot_name.clone(), block_num, keys) + .ok_or_else(|| DatabaseError::StorageRootNotFound { + account_id, + slot_name: slot_name.to_string(), + block_num, + })? + .map_err(DatabaseError::MerkleError)?, + SlotData::All => forest_guard + .storage_map_entries(account_id, slot_name.clone(), block_num) + .ok_or_else(|| DatabaseError::StorageRootNotFound { + account_id, + slot_name: slot_name.to_string(), + block_num, + })?, }; - let details = AccountStorageMapDetails::new(slot_name, slot_data, storage_map); storage_map_details.push(details); } diff --git a/proto/proto/rpc.proto b/proto/proto/rpc.proto index dccf44020..2918af848 100644 --- a/proto/proto/rpc.proto +++ b/proto/proto/rpc.proto @@ -318,8 +318,22 @@ message AccountVaultDetails { // Account storage details for AccountProofResponse message AccountStorageDetails { message AccountStorageMapDetails { - // Wrapper for repeated storage map entries - message MapEntries { + // Wrapper for repeated storage map entries including their proofs. + // Used when specific keys are requested to enable client-side verification. + message MapEntriesWithProofs { + // Definition of individual storage entries including a proof. + message StorageMapEntryWithProof { + primitives.Digest key = 1; + primitives.Digest value = 2; + primitives.SmtOpening proof = 3; + } + + repeated StorageMapEntryWithProof entries = 1; + } + + // Wrapper for repeated storage map entries (without proofs). + // Used when all entries are requested for small maps. + message AllMapEntries { // Definition of individual storage entries. message StorageMapEntry { primitives.Digest key = 1; @@ -332,13 +346,18 @@ message AccountStorageDetails { // Storage slot name. string slot_name = 1; - // A flag that is set to `true` if the number of to-be-returned entries in the - // storage map would exceed a threshold. This indicates to the user that `SyncStorageMaps` - // endpoint should be used to get all storage map data. + // True when the number of entries exceeds the response limit. + // When set, clients should use the `SyncStorageMaps` endpoint. bool too_many_entries = 2; - // By default we provide all storage entries. - MapEntries entries = 3; + // The map entries (with or without proofs). Empty when too_many_entries is true. + oneof entries { + // All storage entries without proofs (for small maps or full requests). + AllMapEntries all_entries = 3; + + // Specific entries with their SMT proofs (for partial requests). + MapEntriesWithProofs entries_with_proofs = 4; + } } // Account storage header (storage slot info for up to 256 slots)