feat(katana): limit number of proof keys (#2814)

crates/dojo/test-utils/src/sequencer.rs

@@ -125,6 +125,7 @@ pub fn get_default_test_config(sequencing: SequencingConfig) -> Config {
         max_connections: DEFAULT_RPC_MAX_CONNECTIONS,
         apis: HashSet::from([ApiKind::Starknet, ApiKind::Dev, ApiKind::Saya, ApiKind::Torii]),
         max_event_page_size: Some(100),
+        max_proof_keys: Some(100),
     };
 
     Config { sequencing, rpc, dev, chain: chain.into(), ..Default::default() }

crates/katana/cli/src/args.rs

@@ -213,6 +213,7 @@ impl NodeArgs {
                 max_connections: self.server.max_connections,
                 cors_origins: self.server.http_cors_origins.clone(),
                 max_event_page_size: Some(self.server.max_event_page_size),
+                max_proof_keys: Some(self.server.max_proof_keys),
             }
         }
 

crates/katana/cli/src/options.rs

@@ -12,6 +12,7 @@ use std::net::IpAddr;
 use clap::Args;
 use katana_node::config::execution::{DEFAULT_INVOCATION_MAX_STEPS, DEFAULT_VALIDATION_MAX_STEPS};
 use katana_node::config::metrics::{DEFAULT_METRICS_ADDR, DEFAULT_METRICS_PORT};
+use katana_node::config::rpc::DEFAULT_RPC_MAX_PROOF_KEYS;
 #[cfg(feature = "server")]
 use katana_node::config::rpc::{
     DEFAULT_RPC_ADDR, DEFAULT_RPC_MAX_CONNECTIONS, DEFAULT_RPC_MAX_EVENT_PAGE_SIZE,
@@ -107,6 +108,12 @@ pub struct ServerOptions {
     #[arg(default_value_t = DEFAULT_RPC_MAX_EVENT_PAGE_SIZE)]
     #[serde(default = "default_page_size")]
     pub max_event_page_size: u64,
+
+    /// Maximum keys for requesting storage proofs.
+    #[arg(long = "rpc.max-proof-keys", value_name = "SIZE")]
+    #[arg(default_value_t = DEFAULT_RPC_MAX_PROOF_KEYS)]
+    #[serde(default = "default_proof_keys")]
+    pub max_proof_keys: u64,
 }
 
 #[cfg(feature = "server")]
@@ -118,6 +125,7 @@ impl Default for ServerOptions {
             max_connections: DEFAULT_RPC_MAX_CONNECTIONS,
             http_cors_origins: Vec::new(),
             max_event_page_size: DEFAULT_RPC_MAX_EVENT_PAGE_SIZE,
+            max_proof_keys: DEFAULT_RPC_MAX_PROOF_KEYS,
         }
     }
 }
@@ -380,6 +388,11 @@ fn default_page_size() -> u64 {
     DEFAULT_RPC_MAX_EVENT_PAGE_SIZE
 }
 
+#[cfg(feature = "server")]
+fn default_proof_keys() -> u64 {
+    katana_node::config::rpc::DEFAULT_RPC_MAX_PROOF_KEYS
+}
+
 #[cfg(feature = "server")]
 fn default_metrics_addr() -> IpAddr {
     DEFAULT_METRICS_ADDR

crates/katana/core/src/backend/mod.rs

@@ -252,7 +252,6 @@ impl<'a, P: TrieWriter> UncommittedBlock<'a, P> {
 
     // state_commitment = hPos("STARKNET_STATE_V0", contract_trie_root, class_trie_root)
     fn compute_new_state_root(&self) -> Felt {
-        println!("ohayo im committing now");
         let class_trie_root = self
             .provider
             .trie_insert_declared_classes(self.header.number, &self.state_updates.declared_classes)

crates/katana/node/src/config/rpc.rs

@@ -10,6 +10,8 @@ pub const DEFAULT_RPC_PORT: u16 = 5050;
 
 /// Default maximmum page size for the `starknet_getEvents` RPC method.
 pub const DEFAULT_RPC_MAX_EVENT_PAGE_SIZE: u64 = 1024;
+/// Default maximmum number of keys for the `starknet_getStorageProof` RPC method.
+pub const DEFAULT_RPC_MAX_PROOF_KEYS: u64 = 100;
 
 /// List of APIs supported by Katana.
 #[derive(
@@ -29,8 +31,9 @@ pub struct RpcConfig {
     pub port: u16,
     pub max_connections: u32,
     pub apis: HashSet<ApiKind>,
-    pub max_event_page_size: Option<u64>,
     pub cors_origins: Vec<HeaderValue>,
+    pub max_event_page_size: Option<u64>,
+    pub max_proof_keys: Option<u64>,
 }
 
 impl RpcConfig {
@@ -49,6 +52,7 @@ impl Default for RpcConfig {
             max_connections: DEFAULT_RPC_MAX_CONNECTIONS,
             apis: HashSet::from([ApiKind::Starknet]),
             max_event_page_size: Some(DEFAULT_RPC_MAX_EVENT_PAGE_SIZE),
+            max_proof_keys: Some(DEFAULT_RPC_MAX_PROOF_KEYS),
         }
     }
 }

crates/katana/node/src/lib.rs

@@ -253,7 +253,10 @@ pub async fn build(mut config: Config) -> Result<Node> {
         .allow_headers([hyper::header::CONTENT_TYPE, "argent-client".parse().unwrap(), "argent-version".parse().unwrap()]);
 
     if config.rpc.apis.contains(&ApiKind::Starknet) {
-        let cfg = StarknetApiConfig { max_event_page_size: config.rpc.max_event_page_size };
+        let cfg = StarknetApiConfig {
+            max_event_page_size: config.rpc.max_event_page_size,
+            max_proof_keys: config.rpc.max_proof_keys,
+        };
 
         let api = if let Some(client) = forked_client {
             StarknetApi::new_forked(

crates/katana/rpc/rpc-types/src/error/starknet.rs

@@ -3,6 +3,7 @@ use jsonrpsee::types::error::CallError;
 use jsonrpsee::types::ErrorObject;
 use katana_pool::validation::error::InvalidTransactionError;
 use katana_pool::PoolError;
+use katana_primitives::block::BlockNumber;
 use katana_primitives::event::ContinuationTokenError;
 use katana_provider::error::ProviderError;
 use serde::Serialize;
@@ -77,12 +78,24 @@ pub enum StarknetApiError {
     UnsupportedContractClassVersion,
     #[error("An unexpected error occured")]
     UnexpectedError { reason: String },
-    #[error("Too many storage keys requested")]
-    ProofLimitExceeded,
     #[error("Too many keys provided in a filter")]
     TooManyKeysInFilter,
     #[error("Failed to fetch pending transactions")]
     FailedToFetchPendingTransactions,
+    #[error("The node doesn't support storage proofs for blocks that are too far in the past")]
+    StorageProofNotSupported {
+        /// The oldest block whose storage proof can be obtained.
+        oldest_block: BlockNumber,
+        /// The block of the storage proof that is being requested.
+        requested_block: BlockNumber,
+    },
+    #[error("Proof limit exceeded")]
+    ProofLimitExceeded {
+        /// The limit for the total number of keys that can be specified in a single request.
+        limit: u64,
+        /// The total number of keys that is being requested.
+        total: u64,
+    },
 }
 
 impl StarknetApiError {
@@ -103,6 +116,7 @@ impl StarknetApiError {
             StarknetApiError::FailedToFetchPendingTransactions => 38,
             StarknetApiError::ContractError { .. } => 40,
             StarknetApiError::TransactionExecutionError { .. } => 41,
+            StarknetApiError::StorageProofNotSupported { .. } => 42,
             StarknetApiError::InvalidContractClass => 50,
             StarknetApiError::ClassAlreadyDeclared => 51,
             StarknetApiError::InvalidTransactionNonce { .. } => 52,
@@ -117,7 +131,7 @@ impl StarknetApiError {
             StarknetApiError::UnsupportedTransactionVersion => 61,
             StarknetApiError::UnsupportedContractClassVersion => 62,
             StarknetApiError::UnexpectedError { .. } => 63,
-            StarknetApiError::ProofLimitExceeded => 10000,
+            StarknetApiError::ProofLimitExceeded { .. } => 1000,
         }
     }
 
@@ -128,8 +142,10 @@ impl StarknetApiError {
     pub fn data(&self) -> Option<serde_json::Value> {
         match self {
             StarknetApiError::ContractError { .. }
-            | StarknetApiError::UnexpectedError { .. }
             | StarknetApiError::PageSizeTooBig { .. }
+            | StarknetApiError::UnexpectedError { .. }
+            | StarknetApiError::ProofLimitExceeded { .. }
+            | StarknetApiError::StorageProofNotSupported { .. }
             | StarknetApiError::TransactionExecutionError { .. } => Some(serde_json::json!(self)),
 
             StarknetApiError::InvalidTransactionNonce { reason }
@@ -284,7 +300,6 @@ mod tests {
     #[case(StarknetApiError::InvalidMessageSelector, 21, "Invalid message selector")]
     #[case(StarknetApiError::NonAccount, 58, "Sender address in not an account contract")]
     #[case(StarknetApiError::InvalidTxnIndex, 27, "Invalid transaction index in a block")]
-    #[case(StarknetApiError::ProofLimitExceeded, 10000, "Too many storage keys requested")]
     #[case(StarknetApiError::TooManyKeysInFilter, 34, "Too many keys provided in a filter")]
     #[case(StarknetApiError::ContractClassSizeIsTooLarge, 57, "Contract class size is too large")]
     #[case(StarknetApiError::FailedToFetchPendingTransactions, 38, "Failed to fetch pending transactions")]
@@ -372,6 +387,30 @@ mod tests {
          	"max_allowed": 500
         }),
     )]
+    #[case(
+    	StarknetApiError::StorageProofNotSupported {
+     		oldest_block: 10,
+       		requested_block: 9
+     	},
+      	42,
+       	"The node doesn't support storage proofs for blocks that are too far in the past",
+        json!({
+        	"oldest_block": 10,
+         	"requested_block": 9
+        }),
+    )]
+    #[case(
+    	StarknetApiError::ProofLimitExceeded {
+     		limit: 5,
+       		total: 10
+     	},
+      	1000,
+       	"Proof limit exceeded",
+        json!({
+        	"limit": 5,
+         	"total": 10
+        }),
+    )]
     fn test_starknet_api_error_to_error_conversion_data_some(
         #[case] starknet_error: StarknetApiError,
         #[case] expected_code: i32,

crates/katana/rpc/rpc-types/src/trie.rs

@@ -6,7 +6,7 @@ use katana_trie::bonsai::BitSlice;
 use katana_trie::{MultiProof, Path, ProofNode};
 use serde::{Deserialize, Serialize};
 
-#[derive(Debug, Serialize, Deserialize)]
+#[derive(Debug, Default, Serialize, Deserialize)]
 pub struct ContractStorageKeys {
     #[serde(rename = "contract_address")]
     pub address: ContractAddress,

crates/katana/rpc/rpc/src/starknet/config.rs

@@ -4,4 +4,9 @@ pub struct StarknetApiConfig {
     ///
     /// If `None`, the maximum chunk size is bounded by [`u64::MAX`].
     pub max_event_page_size: Option<u64>,
+
+    /// The max keys whose proofs can be requested for from the `getStorageProof` method.
+    ///
+    /// If `None`, the maximum keys size is bounded by [`u64::MAX`].
+    pub max_proof_keys: Option<u64>,
 }

crates/katana/rpc/rpc/src/starknet/mod.rs

@@ -1139,6 +1139,18 @@ where
         self.on_io_blocking_task(move |this| {
             let provider = this.inner.backend.blockchain.provider();
 
+            // Check if the total number of keys requested exceeds the RPC limit.
+            if let Some(limit) = this.inner.config.max_proof_keys {
+                let total_keys = class_hashes.as_ref().map(|v| v.len()).unwrap_or(0)
+                    + contract_addresses.as_ref().map(|v| v.len()).unwrap_or(0)
+                    + contracts_storage_keys.as_ref().map(|v| v.len()).unwrap_or(0);
+
+                let total_keys = total_keys as u64;
+                if total_keys > limit {
+                    return Err(StarknetApiError::ProofLimitExceeded { limit, total: total_keys });
+                }
+            }
+
             let state = this.state(&block_id)?;
             let block_hash = provider.latest_hash()?;
 
@@ -1185,7 +1197,6 @@ where
 
             let classes_tree_root = state.classes_root()?;
             let contracts_tree_root = state.contracts_root()?;
-
             let global_roots = GlobalRoots { block_hash, classes_tree_root, contracts_tree_root };
 
             Ok(GetStorageProofResponse {

crates/katana/rpc/rpc/tests/proofs.rs

@@ -1,11 +1,14 @@
 use std::path::PathBuf;
 
+use assert_matches::assert_matches;
 use dojo_test_utils::sequencer::{get_default_test_config, TestSequencer};
 use jsonrpsee::http_client::HttpClientBuilder;
+use katana_node::config::rpc::DEFAULT_RPC_MAX_PROOF_KEYS;
 use katana_node::config::SequencingConfig;
 use katana_primitives::block::BlockIdOrTag;
-use katana_primitives::hash;
+use katana_primitives::class::ClassHash;
 use katana_primitives::hash::StarkHash;
+use katana_primitives::{hash, Felt};
 use katana_rpc_api::starknet::StarknetApiClient;
 use katana_rpc_types::trie::GetStorageProofResponse;
 use katana_trie::bitvec::view::AsBits;
@@ -17,6 +20,58 @@ use starknet::macros::short_string;
 
 mod common;
 
+#[tokio::test]
+async fn proofs_limit() {
+    use jsonrpsee::core::Error;
+    use jsonrpsee::types::error::CallError;
+    use serde_json::json;
+
+    let sequencer =
+        TestSequencer::start(get_default_test_config(SequencingConfig::default())).await;
+
+    // We need to use the jsonrpsee client because `starknet-rs` doesn't yet support RPC 0
+    let client = HttpClientBuilder::default().build(sequencer.url()).unwrap();
+
+    // Because we're using the default configuration for instantiating the node, the RPC limit is
+    // set to 100. The total keys is 35 + 35 + 35 = 105.
+
+    // Generate dummy keys
+    let mut classes = Vec::new();
+    let mut contracts = Vec::new();
+    let mut storages = Vec::new();
+
+    for i in 0..35 {
+        storages.push(Default::default());
+        classes.push(ClassHash::from(i as u64));
+        contracts.push(Felt::from(i as u64).into());
+    }
+
+    let err = client
+        .get_storage_proof(
+            BlockIdOrTag::Tag(BlockTag::Latest),
+            Some(classes),
+            Some(contracts),
+            Some(storages),
+        )
+        .await
+        .expect_err("rpc should enforce limit");
+
+    assert_matches!(err, Error::Call(CallError::Custom(e)) => {
+        assert_eq!(e.code(), 1000);
+        assert_eq!(&e.message(), &"Proof limit exceeded");
+
+        let expected_data = json!({
+            "total": 105,
+            "limit": DEFAULT_RPC_MAX_PROOF_KEYS,
+        });
+
+        let actual_data = e.data().expect("must have data");
+        let actual_data = serde_json::to_value(actual_data).unwrap();
+
+        assert_eq!(actual_data, expected_data);
+    });
+}
+
 #[tokio::test]
 async fn classes_proofs() {
     let sequencer =