starknet_os_runner: storage proof provider trait

Author: AvivYossef-starkware

Cargo.lock

@@ -17,7 +17,12 @@ 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

@@ -23,6 +23,23 @@ use starknet_rust_core::types::{
 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 {
     pub class_hashes: Vec<Felt>,
@@ -270,3 +287,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, &query)?)
+    }
+}

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
+    );
+}