Skip to content

Commit

Permalink
Fetch contract storage proof from v0_8
Browse files Browse the repository at this point in the history
  • Loading branch information
@GMKrieger
GMKrieger committed Jan 8, 2025
1 parent a7e5cee commit dec197e
Show file tree
Hide file tree
Showing 5 changed files with 292 additions and 63 deletions.
48 changes: 32 additions & 16 deletions crates/bin/prove_block/src/lib.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,7 @@ use rpc_replay::block_context::build_block_context;
use rpc_replay::rpc_state_reader::AsyncRpcStateReader;
use rpc_replay::transactions::{starknet_rs_to_blockifier, ToBlockifierError};
use rpc_replay::utils::FeltConversionError;
use rpc_utils::{get_08_class_proofs, get_class_proofs, get_storage_proofs};
use rpc_utils::{get_08_class_proofs, get_08_contracts_proofs, get_class_proofs, get_storage_proofs};
use starknet::core::types::{BlockId, MaybePendingBlockWithTxHashes, MaybePendingBlockWithTxs, StarknetError};
use starknet::providers::{Provider, ProviderError};
use starknet_api::StarknetApiError;
Expand Down Expand Up @@ -119,11 +119,11 @@ fn compute_08_class_commitment(
updated_root: Felt,
) -> CommitmentInfo {
let previous_class_proofs: Vec<_> =
previous_class_proofs.classes_proof.clone().into_iter().map(|proof| proof.node).collect();
let class_proofs: Vec<_> = class_proofs.classes_proof.clone().into_iter().map(|proof| proof.node).collect();
previous_class_proofs.classes_proof.clone().into_iter().map(|proof| proof).collect();
let class_proofs: Vec<_> = class_proofs.classes_proof.clone().into_iter().map(|proof| proof).collect();

let previous_class_commitment_facts = format_08_commitment_facts::<PoseidonHash>(&previous_class_proofs);
let current_class_commitment_facts = format_08_commitment_facts::<PoseidonHash>(&class_proofs);
let previous_class_commitment_facts = format_08_commitment_facts::<PoseidonHash>(&[previous_class_proofs]);
let current_class_commitment_facts = format_08_commitment_facts::<PoseidonHash>(&[class_proofs]);

let class_commitment_facts: HashMap<_, _> =
previous_class_commitment_facts.into_iter().chain(current_class_commitment_facts).collect();
Expand Down Expand Up @@ -217,6 +217,15 @@ pub async fn prove_block(
.await
.expect("Failed to fetch storage proofs");

let storage_proofs_08 = get_08_contracts_proofs(&rpc_client, block_number, &tx_execution_infos, old_block_number)
.await
.expect("Failed to fetch storage proofs");

let previous_storage_proofs_08 =
get_08_contracts_proofs(&rpc_client, block_number - 1, &tx_execution_infos, old_block_number)
.await
.expect("Failed to fetch storage proofs");

let default_general_config = StarknetGeneralConfig::default();

let general_config = StarknetGeneralConfig {
Expand All @@ -233,15 +242,10 @@ pub async fn prove_block(
let mut contract_address_to_class_hash = HashMap::new();

// TODO: remove this clone()
for (contract_address, storage_proof) in storage_proofs.clone() {
for (contract_address, storage_proof) in storage_proofs_08.clone() {
let previous_storage_proof =
previous_storage_proofs.get(&contract_address).expect("failed to find previous storage proof");
let contract_storage_root = previous_storage_proof
.contract_data
.as_ref()
.map(|contract_data| contract_data.root)
.unwrap_or(Felt::ZERO)
.into();
previous_storage_proofs_08.get(&contract_address).expect("failed to find previous storage proof");
let contract_storage_root = previous_storage_proof.contracts_storage_proofs[0][0].node_hash.into();

log::debug!(
"Storage root 0x{:x} for contract 0x{:x}",
Expand Down Expand Up @@ -324,6 +328,13 @@ pub async fn prove_block(
})
.collect();

// We can extract data from any storage proof, use the one of the block hash contract
let block_hash_storage_proof_08 =
storage_proofs_08.get(&Felt::ONE).expect("there should be a storage proof for the block hash contract");
let previous_block_hash_storage_proof_08 = previous_storage_proofs_08
.get(&Felt::ONE)
.expect("there should be a storage proof for the block hash contract");

// We can extract data from any storage proof, use the one of the block hash contract
let block_hash_storage_proof =
storage_proofs.get(&Felt::ONE).expect("there should be a storage proof for the block hash contract");
Expand All @@ -348,6 +359,11 @@ pub async fn prove_block(
None => Felt252::ZERO,
};

let current_contract_trie_root_08 =
block_hash_storage_proof_08.contracts_proof.nodes[0].node.hash::<PedersenHash>();
let previous_contract_trie_root_08 =
previous_block_hash_storage_proof_08.contracts_proof.nodes[0].node.hash::<PedersenHash>();

let previous_contract_proofs: Vec<_> =
previous_storage_proofs.values().map(|proof| proof.contract_proof.clone()).collect();
let previous_state_commitment_facts = format_commitment_facts::<PedersenHash>(&previous_contract_proofs);
Expand All @@ -358,8 +374,8 @@ pub async fn prove_block(
previous_state_commitment_facts.into_iter().chain(current_state_commitment_facts).collect();

let contract_state_commitment_info = CommitmentInfo {
previous_root: previous_contract_trie_root,
updated_root: current_contract_trie_root,
previous_root: previous_contract_trie_root_08,
updated_root: current_contract_trie_root_08,
tree_height: 251,
commitment_facts: global_state_commitment_facts,
};
Expand All @@ -368,7 +384,7 @@ pub async fn prove_block(
compute_class_commitment(&previous_class_proofs, &class_proofs, previous_root, updated_root);

let contract_class_commitment_info_08 =
compute_08_class_commitment(&previous_class_proofs_08, &class_proofs_08, previous_root, updated_root);
compute_08_class_commitment(&previous_class_proofs_08, &class_proofs_08, previous_root_08, updated_root_08);

let os_input = Rc::new(StarknetOsInput {
contract_state_commitment_info,
Expand Down
85 changes: 42 additions & 43 deletions crates/bin/prove_block/src/reexecute.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,7 @@ use blockifier::transaction::objects::TransactionExecutionInfo;
use blockifier::transaction::transaction_execution::Transaction;
use blockifier::transaction::transactions::ExecutableTransaction;
use cairo_vm::Felt252;
use rpc_client::pathfinder::proofs::{ContractData, MerkleNode, PathfinderProof, PathfinderProof, TrieNode, TrieNode};
use rpc_client::pathfinder::proofs::{MerkleNode, NodeHashToNodeMappingItem, StorageProof, TrieNode};
use rpc_client::RpcClient;
use starknet::core::types::{BlockId, StarknetError};
use starknet::providers::{Provider as _, ProviderError};
Expand Down Expand Up @@ -95,8 +95,8 @@ pub(crate) struct ProverPerContractStorage {
block_id: BlockId,
contract_address: Felt252,
previous_tree_root: Felt252,
storage_proof: PathfinderProof,
previous_storage_proof: PathfinderProof,
storage_proof: StorageProof,
previous_storage_proof: StorageProof,
ongoing_storage_changes: HashMap<TreeIndex, Felt252>,
}

Expand All @@ -106,8 +106,8 @@ impl ProverPerContractStorage {
block_id: BlockId,
contract_address: Felt252,
previous_tree_root: Felt252,
storage_proof: PathfinderProof,
previous_storage_proof: PathfinderProof,
storage_proof: StorageProof,
previous_storage_proof: StorageProof,
) -> Result<Self, TreeError> {
Ok(Self {
rpc_client,
Expand Down Expand Up @@ -161,39 +161,43 @@ pub(crate) fn format_commitment_facts<H: HashFunctionType>(
}

pub(crate) fn format_08_commitment_facts<H: HashFunctionType>(
trie_nodes: &Vec<MerkleNode>,
trie_nodes: &[Vec<NodeHashToNodeMappingItem>],
) -> HashMap<Felt252, Vec<Felt252>> {
let mut facts = HashMap::new();

for node in trie_nodes {
let (key, fact_as_tuple) = match node {
MerkleNode::Binary { left, right } => {
let fact = BinaryNodeFact::new((*left).into(), (*right).into())
.expect("storage proof endpoint gave us an invalid binary node");
for nodes in trie_nodes {
for node in nodes {
let (key, fact_as_tuple) = match node.node {
MerkleNode::Binary { left, right } => {
let fact = BinaryNodeFact::new((left).into(), (right).into())
.expect("storage proof endpoint gave us an invalid binary node");

// TODO: the hash function should probably be split from the Fact trait.
// we use a placeholder for the Storage trait in the meantime.
let node_hash = Felt252::from(<BinaryNodeFact as Fact<DictStorage, H>>::hash(&fact));
let fact_as_tuple = <BinaryNodeFact as InnerNodeFact<DictStorage, H>>::to_tuple(&fact);
// TODO: the hash function should probably be split from the Fact trait.
// we use a placeholder for the Storage trait in the meantime.
let node_hash = Felt252::from(<BinaryNodeFact as Fact<DictStorage, H>>::hash(&fact));
let fact_as_tuple = <BinaryNodeFact as InnerNodeFact<DictStorage, H>>::to_tuple(&fact);

(node_hash, fact_as_tuple)
}
MerkleNode::Edge { child, path, length } => {
let len = *length;
let fact =
EdgeNodeFact::new((*child).into(), NodePath(path.to_biguint()), Length(len.try_into().unwrap()))
.expect("storage proof endpoint gave us an invalid edge node");
// TODO: the hash function should probably be split from the Fact trait.
// we use a placeholder for the Storage trait in the meantime.
let node_hash = Felt252::from(<EdgeNodeFact as Fact<DictStorage, H>>::hash(&fact));
let fact_as_tuple = <EdgeNodeFact as InnerNodeFact<DictStorage, H>>::to_tuple(&fact);
(node_hash, fact_as_tuple)
}
MerkleNode::Edge { child, path, length } => {
let fact = EdgeNodeFact::new(
(child).into(),
NodePath(path.to_biguint()),
Length(length.try_into().unwrap()),
)
.expect("storage proof endpoint gave us an invalid edge node");
// TODO: the hash function should probably be split from the Fact trait.
// we use a placeholder for the Storage trait in the meantime.
let node_hash = Felt252::from(<EdgeNodeFact as Fact<DictStorage, H>>::hash(&fact));
let fact_as_tuple = <EdgeNodeFact as InnerNodeFact<DictStorage, H>>::to_tuple(&fact);

(node_hash, fact_as_tuple)
}
};
(node_hash, fact_as_tuple)
}
};

let fact_as_tuple_of_felts: Vec<_> = fact_as_tuple.into_iter().map(Felt252::from).collect();
facts.insert(key, fact_as_tuple_of_felts);
let fact_as_tuple_of_felts: Vec<_> = fact_as_tuple.into_iter().map(Felt252::from).collect();
facts.insert(key, fact_as_tuple_of_felts);
}
}

facts
Expand All @@ -202,20 +206,15 @@ pub(crate) fn format_08_commitment_facts<H: HashFunctionType>(
impl PerContractStorage for ProverPerContractStorage {
async fn compute_commitment(&mut self) -> Result<CommitmentInfo, CommitmentInfoError> {
// TODO: error code
let contract_data = match self.storage_proof.contract_data.as_ref() {
None => &ContractData::default(),
Some(data) => data,
};
let contract_data: &[Vec<NodeHashToNodeMappingItem>] = self.storage_proof.contracts_storage_proofs.as_ref();
let updated_root = contract_data[0][0].node_hash;

let updated_root = contract_data.root;
let commitment_facts = format_08_commitment_facts::<PedersenHash>(&contract_data);

let commitment_facts = format_commitment_facts::<PedersenHash>(&contract_data.storage_proofs);

let previous_commitment_facts = match &self.previous_storage_proof.contract_data {
None => HashMap::default(),
Some(previous_contract_data) => {
format_commitment_facts::<PedersenHash>(&previous_contract_data.storage_proofs)
}
let previous_commitment_facts = if self.previous_storage_proof.contracts_storage_proofs.is_empty() {
HashMap::default()
} else {
format_08_commitment_facts::<PedersenHash>(&self.previous_storage_proof.contracts_storage_proofs)
};

let commitment_facts = commitment_facts.into_iter().chain(previous_commitment_facts.into_iter()).collect();
Expand Down
72 changes: 71 additions & 1 deletion crates/bin/prove_block/src/rpc_utils.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,8 @@ use cairo_vm::Felt252;
use num_bigint::BigInt;
use rpc_client::pathfinder::client::ClientError;
use rpc_client::pathfinder::proofs::{
ContractData, EdgePath, PathfinderClassProof, PathfinderProof, ProofVerificationError, StorageProof, TrieNode,
ContractData, ContractStorageKeysItem, EdgePath, MerkleNode, PathfinderClassProof, PathfinderProof,
ProofVerificationError, StorageProof, TrieNode,
};
use rpc_client::RpcClient;
use starknet::core::types::BlockWithTxs;
Expand Down Expand Up @@ -43,6 +44,22 @@ async fn fetch_storage_proof_for_contract(
Ok(storage_proof)
}

/// Fetches the state + storage proof for a single contract for all the specified keys.
/// This function handles the chunking of requests imposed by the RPC API and merges
/// the proofs returned from multiple calls into one.
async fn fetch_storage_proof_for_contract_08(
rpc_client: &RpcClient,
block_number: u64,
contract_address: Felt,
keys: &[Felt],
) -> Result<StorageProof, ClientError> {
let request_storage_keys = ContractStorageKeysItem { contract_address, storage_keys: keys.to_vec() };
rpc_client
.pathfinder_rpc()
.get_contract_storage_proof(block_number, &[contract_address], &[request_storage_keys])
.await
}

/// Fetches the storage proof for the specified contract and storage keys.
/// This function can fetch additional keys if required to fill gaps in the storage trie
/// that must be filled to get the OS to function. See `get_key_following_edge` for more details.
Expand All @@ -64,6 +81,7 @@ async fn get_storage_proof_for_contract<KeyIter: Iterator<Item = StorageKey>>(
}
Some(contract_data) => contract_data,
};

let additional_keys = verify_storage_proof(contract_data, &keys);

// Fetch additional proofs required to fill gaps in the storage trie that could make
Expand Down Expand Up @@ -201,6 +219,17 @@ fn get_key_following_edge(key: Felt, height: Height, edge_path: &EdgePath) -> Fe
Felt::from(new_key)
}

fn get_key_following_edge_08(key: Felt, height: Height, length: usize, path: Felt) -> Felt {
assert!(height.0 < DEFAULT_STORAGE_TREE_HEIGHT);

let shift = height.0;
let clear_mask = ((BigInt::from(1) << length) - BigInt::from(1)) << shift;
let mask = path.to_bigint() << shift;
let new_key = (key.to_bigint() & !clear_mask) | mask;

Felt::from(new_key)
}

fn merge_storage_proofs(proofs: Vec<PathfinderProof>) -> PathfinderProof {
let class_commitment = proofs[0].class_commitment;
let state_commitment = proofs[0].state_commitment;
Expand Down Expand Up @@ -248,6 +277,47 @@ pub(crate) async fn get_08_class_proofs(
rpc_client.pathfinder_rpc().get_storage_class_proof(block_number, class_hashes).await
}

pub(crate) async fn get_08_contracts_proofs(
rpc_client: &RpcClient,
block_number: u64,
tx_execution_infos: &[TransactionExecutionInfo],
old_block_number: Felt,
) -> Result<HashMap<Felt, StorageProof>, ClientError> {
let accessed_keys_by_address = {
let mut keys = get_all_accessed_keys(tx_execution_infos);
// We need to fetch the storage proof for the block hash contract
keys.entry(contract_address!("0x1")).or_default().insert(old_block_number.try_into().unwrap());
keys
};

let mut storage_proofs = HashMap::new();

for (contract_address, storage_keys) in accessed_keys_by_address {
log::info!(" Fetching proof for {}", contract_address.to_string());
let contract_address_felt = *contract_address.key();
let keys: Vec<_> = storage_keys.into_iter().map(|storage_key| *storage_key.key()).collect();

let mut storage_proof =
fetch_storage_proof_for_contract_08(rpc_client, block_number, contract_address_felt, &keys).await?;

let additional_keys = verify_storage_proof_08(&storage_proof, &keys);

// Fetch additional proofs required to fill gaps in the storage trie that could make
// the OS crash otherwise.
if !additional_keys.is_empty() {
let additional_proof =
fetch_storage_proof_for_contract_08(rpc_client, block_number, contract_address_felt, &additional_keys)
.await?;

storage_proof.contracts_storage_proofs.extend(additional_proof.contracts_storage_proofs);
}

storage_proofs.insert(contract_address_felt, storage_proof);
}

Ok(storage_proofs)
}

pub(crate) fn get_starknet_version(block_with_txs: &BlockWithTxs) -> blockifier::versioned_constants::StarknetVersion {
let starknet_version_str = &block_with_txs.starknet_version;
match starknet_version_str.as_ref() {
Expand Down
Loading

0 comments on commit dec197e

Please sign in to comment.