feat: Return code for FPI (#521)

* feat: Return code for FPI

* Rename var

* CHANGELOG

* Address comments

* Check earlier

* Address comments

* Force proto rebuild

* Update crates/rpc-proto/proto/responses.proto

Co-authored-by: Mirko <[email protected]>

* Update description

* Proto rebuild

* Update description

---------

Co-authored-by: Mirko <[email protected]>

CHANGELOG.md

@@ -2,6 +2,7 @@
 
 ## v0.6.0 (TBD)
 
+- Added `AccountCode` as part of `GetAccountProofs` endpoint response (#521).
 - [BREAKING] Added `kernel_root` to block header's protobuf message definitions (#496).
 - [BREAKING] Renamed `off-chain` and `on-chain` to `private` and `public` respectively for the account storage modes (#489).
 - Optimized state synchronizations by removing unnecessary fetching and parsing of note details (#462).

crates/proto/src/generated/requests.rs

@@ -153,7 +153,11 @@ pub struct GetAccountProofsRequest {
     /// List of account IDs to get states.
     #[prost(message, repeated, tag = "1")]
     pub account_ids: ::prost::alloc::vec::Vec<super::account::AccountId>,
+    /// 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.
+    #[prost(message, repeated, tag = "2")]
+    pub code_commitments: ::prost::alloc::vec::Vec<super::digest::Digest>,
     /// Optional flag to include header in the response. `false` by default.
-    #[prost(bool, optional, tag = "2")]
+    #[prost(bool, optional, tag = "3")]
     pub include_headers: ::core::option::Option<bool>,
 }

crates/proto/src/generated/responses.rs

@@ -227,4 +227,8 @@ 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.
+    #[prost(bytes = "vec", optional, tag = "3")]
+    pub account_code: ::core::option::Option<::prost::alloc::vec::Vec<u8>>,
 }

crates/rpc-proto/proto/requests.proto

@@ -135,6 +135,9 @@ message GetAccountStateDeltaRequest {
 message GetAccountProofsRequest {
     // List of account IDs to get states.
     repeated account.AccountId account_ids = 1;
+    // 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.
+    repeated digest.Digest code_commitments = 2;
     // Optional flag to include header in the response. `false` by default.
-    optional bool include_headers = 2;
+    optional bool include_headers = 3;
 }

crates/rpc-proto/proto/responses.proto

@@ -197,4 +197,7 @@ message AccountStateHeader {
     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.
+    optional bytes account_code = 3;
 }

crates/store/src/server/api.rs

@@ -1,4 +1,4 @@
-use std::sync::Arc;
+use std::{collections::BTreeSet, sync::Arc};
 
 use miden_node_proto::{
     convert,
@@ -485,13 +485,25 @@ impl api_server::Api for StoreApi {
         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 account_ids = convert(request.account_ids);
         let include_headers = request.include_headers.unwrap_or_default();
-        let (block_num, infos) =
-            self.state.get_account_proofs(account_ids, include_headers).await?;
+        let account_ids: Vec<u64> = convert(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 (block_num, infos) = self
+            .state
+            .get_account_proofs(account_ids, request_code_commitments, include_headers)
+            .await?;
 
         Ok(Response::new(GetAccountProofsResponse {
             block_num,

crates/store/src/state.rs

@@ -5,6 +5,7 @@
 
 use std::{
     collections::{BTreeMap, BTreeSet},
+    ops::Not,
     sync::Arc,
 };
 
@@ -683,6 +684,7 @@ impl State {
     pub async fn get_account_proofs(
         &self,
         account_ids: Vec<AccountId>,
+        request_code_commitments: BTreeSet<RpoDigest>,
         include_headers: bool,
     ) -> Result<(BlockNumber, Vec<AccountProofsResponse>), DatabaseError> {
         // Lock inner state for the whole operation. We need to hold this lock to prevent the
@@ -696,8 +698,7 @@ impl State {
             let infos = self.db.select_accounts_by_ids(account_ids.clone()).await?;
 
             if account_ids.len() > infos.len() {
-                let found_ids: BTreeSet<AccountId> =
-                    infos.iter().map(|info| info.summary.account_id.into()).collect();
+                let found_ids = infos.iter().map(|info| info.summary.account_id.into()).collect();
                 return Err(DatabaseError::AccountsNotFoundInDb(
                     BTreeSet::from_iter(account_ids).difference(&found_ids).copied().collect(),
                 ));
@@ -712,6 +713,12 @@ impl State {
                             AccountStateHeader {
                                 header: Some(AccountHeader::from(&details).into()),
                                 storage_header: details.storage().get_header().to_bytes(),
+                                // Only include account code if the request did not contain it
+                                // (known by the caller)
+                                account_code: request_code_commitments
+                                    .contains(&details.code().commitment())
+                                    .not()
+                                    .then_some(details.code().to_bytes()),
                             },
                         )
                     })
@@ -724,12 +731,13 @@ impl State {
             .map(|account_id| {
                 let acc_leaf_idx = LeafIndex::new_max_depth(account_id);
                 let opening = inner_state.account_tree.open(&acc_leaf_idx);
+                let state_header = state_headers.get(&account_id).cloned();
 
                 AccountProofsResponse {
                     account_id: Some(account_id.into()),
                     account_hash: Some(opening.value.into()),
                     account_proof: Some(opening.path.into()),
-                    state_header: state_headers.get(&account_id).cloned(),
+                    state_header,
                 }
             })
             .collect();

proto/requests.proto

@@ -135,6 +135,9 @@ message GetAccountStateDeltaRequest {
 message GetAccountProofsRequest {
     // List of account IDs to get states.
     repeated account.AccountId account_ids = 1;
+    // 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.
+    repeated digest.Digest code_commitments = 2;
     // Optional flag to include header in the response. `false` by default.
-    optional bool include_headers = 2;
+    optional bool include_headers = 3;
 }

proto/responses.proto

@@ -197,4 +197,7 @@ message AccountStateHeader {
     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.
+    optional bytes account_code = 3;
 }