feat: Enable storage map querying (0xPolygonMiden#598)

* feat: Enable storage map querying

* reviews: Clarify doc comments; error messages;  iterate over requests instead of retreived accounts; README suggestions

CHANGELOG.md

@@ -14,6 +14,7 @@
 - [BREAKING] Added support for new two `Felt` account ID (#591).
 - [BREAKING] Inverted `TransactionInputs.missing_unauthenticated_notes` to `found_missing_notes` (#509).
 - [BREAKING] Remove store's `ListXXX` endpoints which were intended for test purposes (#608).
+- [BREAKING] Added support for storage maps on `GetAccountProofs` endpoint (#598).
 
 ## v0.6.0 (2024-11-05)
 

crates/proto/src/domain/accounts.rs

@@ -8,6 +8,7 @@ use miden_objects::{
     Digest,
 };
 
+use super::try_convert;
 use crate::{
     errors::{ConversionError, MissingFieldHelper},
     generated as proto,
@@ -93,6 +94,60 @@ impl From<&AccountInfo> for proto::account::AccountInfo {
     }
 }
 
+// ACCOUNT STORAGE REQUEST
+// ================================================================================================
+
+/// Represents a request for an account proof alongside specific storage data.
+pub struct AccountProofRequest {
+    pub account_id: AccountId,
+    pub storage_requests: Vec<StorageMapKeysProof>,
+}
+
+impl TryInto<AccountProofRequest> for proto::requests::get_account_proofs_request::AccountRequest {
+    type Error = ConversionError;
+
+    fn try_into(self) -> Result<AccountProofRequest, Self::Error> {
+        let proto::requests::get_account_proofs_request::AccountRequest {
+            account_id,
+            storage_requests,
+        } = self;
+
+        Ok(AccountProofRequest {
+            account_id: account_id
+                .clone()
+                .ok_or(proto::requests::get_account_proofs_request::AccountRequest::missing_field(
+                    stringify!(account_id),
+                ))?
+                .try_into()?,
+            storage_requests: try_convert(storage_requests)?,
+        })
+    }
+}
+
+/// Represents a request for an account's storage map values and its proof of existence.
+pub struct StorageMapKeysProof {
+    /// Index of the storage map
+    pub storage_index: u8,
+    /// List of requested keys in the map
+    pub storage_keys: Vec<Digest>,
+}
+
+impl TryInto<StorageMapKeysProof> for proto::requests::get_account_proofs_request::StorageRequest {
+    type Error = ConversionError;
+
+    fn try_into(self) -> Result<StorageMapKeysProof, Self::Error> {
+        let proto::requests::get_account_proofs_request::StorageRequest {
+            storage_slot_index,
+            map_keys,
+        } = self;
+
+        Ok(StorageMapKeysProof {
+            storage_index: storage_slot_index.try_into()?,
+            storage_keys: try_convert(map_keys)?,
+        })
+    }
+}
+
 // ACCOUNT INPUT RECORD
 // ================================================================================================
 

crates/proto/src/generated/requests.rs

@@ -142,12 +142,16 @@ pub struct GetAccountStateDeltaRequest {
     #[prost(fixed32, tag = "3")]
     pub to_block_num: u32,
 }
+/// Request message to get account proofs.
 #[derive(Clone, PartialEq, ::prost::Message)]
 pub struct GetAccountProofsRequest {
-    /// List of account IDs to get states.
+    /// A list of account requests, including map keys + values.
     #[prost(message, repeated, tag = "1")]
-    pub account_ids: ::prost::alloc::vec::Vec<super::account::AccountId>,
-    /// Optional flag to include header and account code in the response. `false` by default.
+    pub account_requests: ::prost::alloc::vec::Vec<
+        get_account_proofs_request::AccountRequest,
+    >,
+    /// Optional flag to include account headers and account code in the response. If false, storage
+    /// requests are also ignored. False by default.
     #[prost(bool, optional, tag = "2")]
     pub include_headers: ::core::option::Option<bool>,
     /// Account code commitments corresponding to the last-known `AccountCode` for requested
@@ -157,3 +161,27 @@ pub struct GetAccountProofsRequest {
     #[prost(message, repeated, tag = "3")]
     pub code_commitments: ::prost::alloc::vec::Vec<super::digest::Digest>,
 }
+/// Nested message and enum types in `GetAccountProofsRequest`.
+pub mod get_account_proofs_request {
+    /// Represents per-account requests where each account ID has its own list of
+    /// (storage_slot_index, map_keys) pairs.
+    #[derive(Clone, PartialEq, ::prost::Message)]
+    pub struct AccountRequest {
+        /// The account ID for this request.
+        #[prost(message, optional, tag = "1")]
+        pub account_id: ::core::option::Option<super::super::account::AccountId>,
+        /// List of storage requests for this account.
+        #[prost(message, repeated, tag = "2")]
+        pub storage_requests: ::prost::alloc::vec::Vec<StorageRequest>,
+    }
+    /// Represents a storage slot index and the associated map keys.
+    #[derive(Clone, PartialEq, ::prost::Message)]
+    pub struct StorageRequest {
+        /// Storage slot index (\[0..255\])
+        #[prost(uint32, tag = "1")]
+        pub storage_slot_index: u32,
+        /// A list of map keys (Digests) associated with this storage slot.
+        #[prost(message, repeated, tag = "2")]
+        pub map_keys: ::prost::alloc::vec::Vec<super::super::digest::Digest>,
+    }
+}

crates/proto/src/generated/responses.rs

@@ -209,8 +209,21 @@ pub struct AccountStateHeader {
     /// Values of all account storage slots (max 255).
     #[prost(bytes = "vec", tag = "2")]
     pub storage_header: ::prost::alloc::vec::Vec<u8>,
-    /// Account code, returned only when none of the request's code commitments match with the
-    /// current one.
+    /// Account code, returned only when none of the request's code commitments match
+    /// the current one.
     #[prost(bytes = "vec", optional, tag = "3")]
     pub account_code: ::core::option::Option<::prost::alloc::vec::Vec<u8>>,
+    /// Storage slots information for this account
+    #[prost(message, repeated, tag = "4")]
+    pub storage_maps: ::prost::alloc::vec::Vec<StorageSlotMapProof>,
+}
+/// Represents a single storage slot with the reuqested keys and their respective values.
+#[derive(Clone, PartialEq, ::prost::Message)]
+pub struct StorageSlotMapProof {
+    /// The storage slot index (\[0..255\]).
+    #[prost(uint32, tag = "1")]
+    pub storage_slot: u32,
+    /// Merkle proof of the map value
+    #[prost(bytes = "vec", tag = "2")]
+    pub smt_proof: ::prost::alloc::vec::Vec<u8>,
 }

crates/rpc-proto/proto/requests.proto

@@ -126,11 +126,34 @@ message GetAccountStateDeltaRequest {
     fixed32 to_block_num = 3;
 }
 
+// Request message to get account proofs.
 message GetAccountProofsRequest {
-    // List of account IDs to get states.
-    repeated account.AccountId account_ids = 1;
-    // Optional flag to include header and account code in the response. `false` by default.
+    // Represents per-account requests where each account ID has its own list of 
+    // (storage_slot_index, map_keys) pairs.
+    message AccountRequest {
+        // The account ID for this request.
+        account.AccountId account_id = 1;
+
+        // List of storage requests for this account.
+        repeated StorageRequest storage_requests = 2;
+    }
+
+    // Represents a storage slot index and the associated map keys.
+    message StorageRequest {
+        // Storage slot index ([0..255])
+        uint32 storage_slot_index = 1;
+
+        // A list of map keys (Digests) associated with this storage slot.
+        repeated digest.Digest map_keys = 2;
+    }
+
+    // A list of account requests, including map keys + values.
+    repeated AccountRequest account_requests = 1;
+
+    // Optional flag to include account headers and account code in the response. If false, storage
+    // requests are also ignored. False by default.
     optional bool include_headers = 2;
+
     // Account code commitments corresponding to the last-known `AccountCode` for requested 
     // accounts. Responses will include only the ones that are not known to the caller.
     // These are not associated with a specific account but rather, they will be matched against

crates/rpc-proto/proto/responses.proto

@@ -180,9 +180,23 @@ message AccountProofsResponse {
 message AccountStateHeader {
     // Account header.
     account.AccountHeader header = 1;
+
     // Values of all account storage slots (max 255).
     bytes storage_header = 2;
-    // Account code, returned only when none of the request's code commitments match with the
-    // current one.
+
+    // Account code, returned only when none of the request's code commitments match
+    // the current one.
     optional bytes account_code = 3;
+
+    // Storage slots information for this account
+    repeated StorageSlotMapProof storage_maps = 4;
+}
+
+// Represents a single storage slot with the reuqested keys and their respective values.
+message StorageSlotMapProof {
+    // The storage slot index ([0..255]).
+    uint32 storage_slot = 1;
+
+    // Merkle proof of the map value
+    bytes smt_proof = 2;
 }

crates/rpc/src/server/api.rs

@@ -265,13 +265,19 @@ impl api_server::Api for RpcApi {
 
         debug!(target: COMPONENT, ?request);
 
-        if request.account_ids.len() > MAX_NUM_FOREIGN_ACCOUNTS as usize {
+        if request.account_requests.len() > MAX_NUM_FOREIGN_ACCOUNTS as usize {
             return Err(Status::invalid_argument(format!(
                 "Too many accounts requested: {}, limit: {MAX_NUM_FOREIGN_ACCOUNTS}",
-                request.account_ids.len()
+                request.account_requests.len()
             )));
         }
 
+        if request.account_requests.len() < request.code_commitments.len() {
+            return Err(Status::invalid_argument(
+                "The number of code commitments should not exceed the number of requested accounts.",
+            ));
+        }
+
         self.store.clone().get_account_proofs(request).await
     }
 }

crates/store/README.md

@@ -2,6 +2,8 @@
 
 The **Store** maintains the state of the chain. It serves as the "source of truth" for the chain - i.e., if it is not in
 the store, the node does not consider it to be part of the chain.
+Incoming requests to the store are trusted because they are validated in the RPC component.
+
 **Store** is one of components of the [Miden node](..).
 
 ## Architecture
@@ -16,7 +18,7 @@ The Store can be installed and run as part of [Miden node](../README.md#installi
 
 ## API
 
-The **Store** serves connections using the [gRPC protocol](https://grpc.io) on a port, set in the previously mentioned configuration file.
+The **Store** serves connections using the [gRPC protocol](https://grpc.io) on a port, set in the previously mentioned configuration file. 
 Here is a brief description of supported methods.
 
 ### ApplyBlock

crates/store/src/server/api.rs

@@ -2,7 +2,10 @@ use std::{collections::BTreeSet, sync::Arc};
 
 use miden_node_proto::{
     convert,
-    domain::{accounts::AccountInfo, notes::NoteAuthenticationInfo},
+    domain::{
+        accounts::{AccountInfo, AccountProofRequest},
+        notes::NoteAuthenticationInfo,
+    },
     errors::ConversionError,
     generated::{
         self,
@@ -473,25 +476,25 @@ impl api_server::Api for StoreApi {
         &self,
         request: Request<GetAccountProofsRequest>,
     ) -> Result<Response<GetAccountProofsResponse>, Status> {
-        let request = request.into_inner();
-        if request.account_ids.len() < request.code_commitments.len() {
-            return Err(Status::invalid_argument(
-                "The number of code commitments should not exceed the number of requested accounts.",
-            ));
-        }
-
         debug!(target: COMPONENT, ?request);
+        let GetAccountProofsRequest {
+            account_requests,
+            include_headers,
+            code_commitments,
+        } = request.into_inner();
 
-        let include_headers = request.include_headers.unwrap_or_default();
-        let account_ids: Vec<AccountId> = read_account_ids(&request.account_ids)?;
-        let request_code_commitments: BTreeSet<RpoDigest> = try_convert(request.code_commitments)
-            .map_err(|err| {
-            Status::invalid_argument(format!("Invalid code commitment: {}", err))
-        })?;
+        let include_headers = include_headers.unwrap_or_default();
+        let request_code_commitments: BTreeSet<RpoDigest> = try_convert(code_commitments)
+            .map_err(|err| Status::invalid_argument(format!("Invalid code commitment: {}", err)))?;
+
+        let account_requests: Vec<AccountProofRequest> =
+            try_convert(account_requests).map_err(|err| {
+                Status::invalid_argument(format!("Invalid account proofs request: {}", err))
+            })?;
 
         let (block_num, infos) = self
             .state
-            .get_account_proofs(account_ids, request_code_commitments, include_headers)
+            .get_account_proofs(account_requests, request_code_commitments, include_headers)
             .await?;
 
         Ok(Response::new(GetAccountProofsResponse {