starknet_os_runner: storage proof provider trait

Author: AvivYossef-starkware

Cargo.lock

@@ -10,14 +10,19 @@ description = "Runs transactions through the Starknet OS and returns Cairo PIE a
 blockifier.workspace = true
 blockifier_reexecution.workspace = true
 indexmap.workspace = true
-starknet_api.workspace = true
-starknet_os.workspace = true
-starknet_patricia.workspace = true
 starknet-rust.workspace = true
 starknet-rust-core.workspace = true
 starknet-types-core.workspace = true
+starknet_api.workspace = true
+starknet_os.workspace = true
+starknet_patricia.workspace = true
 thiserror.workspace = true
+tokio.workspace = true
 url.workspace = true
 
+[dev-dependencies]
+blockifier = { workspace = true, features = ["testing"] }
+rstest.workspace = true
+
 [lints]
 workspace = true

crates/starknet_os_runner/Cargo.toml

@@ -2,5 +2,7 @@ pub mod errors;
 pub mod storage_proofs;
 pub mod virtual_block_executor;
 
+#[cfg(test)]
+mod storage_proofs_test;
 #[cfg(test)]
 mod virtual_block_executor_test;

crates/starknet_os_runner/src/lib.rs

@@ -4,8 +4,7 @@ use blockifier::state::cached_state::StateMaps;
 use starknet_api::block::BlockNumber;
 use starknet_api::core::{ClassHash, ContractAddress, Nonce};
 use starknet_api::hash::HashOutput;
-use starknet_api::state::StorageKey;
-use starknet_os::io::os_input::{CachedStateInput, CommitmentInfo, CommitmentInfos};
+use starknet_os::io::os_input::{CommitmentInfo, StateCommitmentInfos};
 use starknet_patricia::patricia_merkle_tree::node_data::inner_node::{
     flatten_preimages,
     Preimage,
@@ -20,9 +19,26 @@ use starknet_rust_core::types::{
     Felt,
     StorageProof as RpcStorageProof,
 };
-use starknet_types_core::felt::Felt as TypesFelt;
 
 use crate::errors::ProofProviderError;
+use crate::virtual_block_executor::VirtualBlockExecutionData;
+
+/// Provides Patricia Merkle proofs for the initial state used in transaction execution.
+///
+/// This trait abstracts the retrieval of storage proofs, which are essential for OS input
+/// generation. The proofs allow the OS to verify that the initial state values (read during
+/// execution) are consistent with the global state commitment (Patricia root).
+///
+/// The returned `StorageProofs` contains:
+/// - `proof_state`: The ambient state values (nonces, class hashes) discovered in the proof.
+/// - `commitment_infos`: The Patricia Merkle proof nodes for contracts, classes, and storage tries.
+pub trait StorageProofProvider {
+    fn get_storage_proofs(
+        &self,
+        block_number: BlockNumber,
+        execution_data: &VirtualBlockExecutionData,
+    ) -> Result<StorageProofs, ProofProviderError>;
+}
 
 /// Query parameters for fetching storage proofs from RPC.
 pub struct RpcStorageProofsQuery {
@@ -33,8 +49,11 @@ pub struct RpcStorageProofsQuery {
 
 /// Complete OS input data built from RPC proofs.
 pub struct StorageProofs {
-    pub cached_state_input: CachedStateInput,
-    pub commitment_infos: CommitmentInfos,
+    /// State information discovered in the Patricia proof (nonces, class hashes)
+    /// that might not have been explicitly read during transaction execution.
+    /// This data is required by the OS to verify the contract state leaves.
+    pub proof_state: StateMaps,
+    pub commitment_infos: StateCommitmentInfos,
 }
 
 /// Converts RPC merkle nodes (hash → MerkleNode mapping) to a PreimageMap.
@@ -55,12 +74,11 @@ impl RpcStorageProofsProvider {
     }
 
     /// Extract query parameters from StateMaps.
-    pub fn prepare_query(
-        initial_reads: &StateMaps,
-        executed_class_hashes: &HashSet<ClassHash>,
-    ) -> RpcStorageProofsQuery {
-        let class_hashes: Vec<Felt> = executed_class_hashes.iter().map(|ch| ch.0).collect();
+    pub fn prepare_query(execution_data: &VirtualBlockExecutionData) -> RpcStorageProofsQuery {
+        let class_hashes: Vec<Felt> =
+            execution_data.executed_class_hashes.iter().map(|ch| ch.0).collect();
 
+        let initial_reads = &execution_data.initial_reads;
         let contract_addresses: Vec<ContractAddress> =
             initial_reads.get_contract_addresses().into_iter().collect();
 
@@ -107,49 +125,26 @@ impl RpcStorageProofsProvider {
     /// Converts an RPC storage proof response to OS input format.
     pub fn to_storage_proofs(
         rpc_proof: &RpcStorageProof,
-        initial_reads: &StateMaps,
         contract_addresses: &[ContractAddress],
     ) -> StorageProofs {
-        let cached_state_input =
-            Self::build_cached_state_input(rpc_proof, initial_reads, contract_addresses);
+        let mut proof_state = StateMaps::default();
         let commitment_infos = Self::build_commitment_infos(rpc_proof, contract_addresses);
 
-        StorageProofs { cached_state_input, commitment_infos }
-    }
-
-    fn build_cached_state_input(
-        rpc_proof: &RpcStorageProof,
-        initial_reads: &StateMaps,
-        contract_addresses: &[ContractAddress],
-    ) -> CachedStateInput {
-        let (address_to_class_hash, address_to_nonce) = rpc_proof
-            .contracts_proof
-            .contract_leaves_data
-            .iter()
-            .zip(contract_addresses)
-            .map(|(leaf, addr)| ((*addr, ClassHash(leaf.class_hash)), (*addr, Nonce(leaf.nonce))))
-            .unzip();
-
-        let storage = initial_reads.storage.iter().fold(
-            HashMap::<ContractAddress, HashMap<StorageKey, TypesFelt>>::new(),
-            |mut acc, ((addr, key), val)| {
-                acc.entry(*addr).or_default().insert(*key, *val);
-                acc
-            },
-        );
-
-        CachedStateInput {
-            storage,
-            address_to_class_hash,
-            address_to_nonce,
-            class_hash_to_compiled_class_hash: initial_reads.compiled_class_hashes.clone(),
+        // Update proof_state with class hashes and nonces from the proof.
+        for (leaf, addr) in
+            rpc_proof.contracts_proof.contract_leaves_data.iter().zip(contract_addresses)
+        {
+            proof_state.class_hashes.insert(*addr, ClassHash(leaf.class_hash));
+            proof_state.nonces.insert(*addr, Nonce(leaf.nonce));
         }
+
+        StorageProofs { proof_state, commitment_infos }
     }
 
     fn build_commitment_infos(
         rpc_proof: &RpcStorageProof,
         contract_addresses: &[ContractAddress],
-    ) -> CommitmentInfos {
+    ) -> StateCommitmentInfos {
         let contracts_tree_root = HashOutput(rpc_proof.global_roots.contracts_tree_root);
         let classes_tree_root = HashOutput(rpc_proof.global_roots.classes_tree_root);
 
@@ -174,7 +169,7 @@ impl RpcStorageProofsProvider {
         let storage_tries_commitment_infos =
             Self::build_storage_commitment_infos(rpc_proof, contract_addresses);
 
-        CommitmentInfos {
+        StateCommitmentInfos {
             contracts_trie_commitment_info,
             classes_trie_commitment_info,
             storage_tries_commitment_infos,
@@ -219,3 +214,19 @@ impl RpcStorageProofsProvider {
             .collect()
     }
 }
+
+impl StorageProofProvider for RpcStorageProofsProvider {
+    fn get_storage_proofs(
+        &self,
+        block_number: BlockNumber,
+        execution_data: &VirtualBlockExecutionData,
+    ) -> Result<StorageProofs, ProofProviderError> {
+        let query = Self::prepare_query(execution_data);
+        let contract_addresses = query.contract_addresses.clone();
+
+        let runtime = tokio::runtime::Runtime::new().expect("Failed to create tokio runtime");
+        let rpc_proof = runtime.block_on(self.fetch_proofs(block_number, &query))?;
+
+        Ok(Self::to_storage_proofs(&rpc_proof, &contract_addresses))
+    }
+}

crates/starknet_os_runner/src/storage_proofs.rs

@@ -0,0 +1,128 @@
+use std::collections::HashSet;
+use std::env;
+
+use blockifier::context::BlockContext;
+use blockifier::state::cached_state::StateMaps;
+use rstest::{fixture, rstest};
+use starknet_api::block::BlockNumber;
+use starknet_api::core::ContractAddress;
+use starknet_api::state::StorageKey;
+use starknet_rust::providers::Provider;
+use starknet_types_core::felt::Felt;
+use url::Url;
+
+use crate::storage_proofs::{RpcStorageProofsProvider, StorageProofProvider};
+use crate::virtual_block_executor::VirtualBlockExecutionData;
+
+/// Mainnet STRK token contract address.
+const STRK_CONTRACT_ADDRESS: Felt =
+    Felt::from_hex_unchecked("0x04718f5a0fc34cc1af16a1cdee98ffb20c31f5cd61d6ab07201858f4287c938d");
+
+/// Fixture: Creates initial reads with the STRK contract and storage slot 0.
+#[fixture]
+fn initial_reads() -> (StateMaps, ContractAddress, StorageKey) {
+    let mut state_maps = StateMaps::default();
+    let contract_address = ContractAddress::try_from(STRK_CONTRACT_ADDRESS).unwrap();
+
+    // Add a storage read for slot 0 (commonly used for total_supply or similar).
+    let storage_key = StorageKey::from(0u32);
+    state_maps.storage.insert((contract_address, storage_key), Felt::ZERO);
+
+    (state_maps, contract_address, storage_key)
+}
+
+/// Fixture: Creates an RPC provider from the RPC_URL environment variable.
+#[fixture]
+fn rpc_provider() -> RpcStorageProofsProvider {
+    let rpc_url_str = env::var("NODE_URL").expect("NODE_URL environment variable must be set");
+    let rpc_url = Url::parse(&rpc_url_str).expect("Invalid RPC URL");
+    RpcStorageProofsProvider::new(rpc_url)
+}
+
+/// Sanity test that verifies storage proof fetching works with a real RPC endpoint.
+///
+/// This test is ignored by default because it requires a running RPC node.
+/// Run with: `RPC_URL=<your_rpc_url> cargo test -p starknet_os_runner -- --ignored`
+#[rstest]
+#[ignore]
+fn test_get_storage_proofs_from_rpc(
+    rpc_provider: RpcStorageProofsProvider,
+    initial_reads: (StateMaps, ContractAddress, StorageKey),
+) {
+    let (state_maps, contract_address, storage_key) = initial_reads;
+
+    // Fetch latest block number.
+    let runtime = tokio::runtime::Runtime::new().expect("Failed to create tokio runtime");
+    let block_number = runtime.block_on(async { rpc_provider.0.block_number().await }).unwrap();
+
+    let execution_data = VirtualBlockExecutionData {
+        execution_outputs: vec![],
+        block_context: BlockContext::create_for_account_testing(),
+        initial_reads: state_maps,
+        executed_class_hashes: HashSet::new(),
+    };
+
+    let result = rpc_provider.get_storage_proofs(BlockNumber(block_number), &execution_data);
+    assert!(result.is_ok(), "Failed to get storage proofs: {:?}", result.err());
+
+    let storage_proofs = result.unwrap();
+
+    // Verify contracts tree root is non-zero.
+    assert!(
+        storage_proofs.commitment_infos.contracts_trie_commitment_info.previous_root.0
+            != Felt::ZERO,
+        "Expected non-zero contracts tree root"
+    );
+
+    // Verify contracts tree commitment facts are not empty.
+    assert!(
+        !storage_proofs.commitment_infos.contracts_trie_commitment_info.commitment_facts.is_empty(),
+        "Expected non-empty contracts tree commitment facts"
+    );
+
+    // Verify the queried contract is in proof_state.
+    assert!(
+        storage_proofs.proof_state.class_hashes.contains_key(&contract_address),
+        "Expected contract address {:?} in class_hashes",
+        contract_address
+    );
+    assert!(
+        storage_proofs.proof_state.nonces.contains_key(&contract_address),
+        "Expected contract address {:?} in nonces",
+        contract_address
+    );
+
+    // Verify the queried storage is in the original execution_data (not proof_state, which only has
+    // nonces/hashes)
+    assert!(
+        execution_data.initial_reads.storage.contains_key(&(contract_address, storage_key)),
+        "Expected storage key {:?} in contract's storage",
+        storage_key
+    );
+
+    // Verify the contract has a storage trie commitment info.
+    assert!(
+        storage_proofs
+            .commitment_infos
+            .storage_tries_commitment_infos
+            .contains_key(&contract_address),
+        "Expected contract address {:?} in storage_tries_commitment_infos",
+        contract_address
+    );
+
+    // Verify the storage trie commitment facts are not empty.
+    let storage_commitment =
+        &storage_proofs.commitment_infos.storage_tries_commitment_infos[&contract_address];
+    assert!(
+        !storage_commitment.commitment_facts.is_empty(),
+        "Expected non-empty storage trie commitment facts for contract {:?}",
+        contract_address
+    );
+
+    // Verify the storage root is non-zero.
+    assert!(
+        storage_commitment.previous_root.0 != Felt::ZERO,
+        "Expected non-zero storage root for contract {:?}",
+        contract_address
+    );
+}