From 3691369afc317cd9591219f75f90800a093f6031 Mon Sep 17 00:00:00 2001 From: Ammar Arif Date: Tue, 24 Dec 2024 20:17:14 -0500 Subject: [PATCH] fix(katana): invalid trie path conversion (#2844) fix trie path rpc conversion --- crates/katana/rpc/rpc-types/src/trie.rs | 85 ++++++++--- crates/katana/rpc/rpc/Cargo.toml | 20 +-- crates/katana/rpc/rpc/tests/proofs.rs | 132 +++++++++++++----- .../tests/test_data/test_sierra_contract.json | 1 + crates/katana/trie/src/classes.rs | 31 +++- crates/katana/trie/src/lib.rs | 33 ++--- 6 files changed, 215 insertions(+), 87 deletions(-) create mode 120000 crates/katana/rpc/rpc/tests/test_data/test_sierra_contract.json diff --git a/crates/katana/rpc/rpc-types/src/trie.rs b/crates/katana/rpc/rpc-types/src/trie.rs index 58f797cc12..a92f0b4426 100644 --- a/crates/katana/rpc/rpc-types/src/trie.rs +++ b/crates/katana/rpc/rpc-types/src/trie.rs @@ -1,9 +1,10 @@ use std::ops::{Deref, DerefMut}; use katana_primitives::contract::StorageKey; +use katana_primitives::hash::StarkHash; use katana_primitives::{ContractAddress, Felt}; -use katana_trie::bonsai::BitSlice; -use katana_trie::{MultiProof, Path, ProofNode}; +use katana_trie::bitvec::view::BitView; +use katana_trie::{BitVec, MultiProof, Path, ProofNode}; use serde::{Deserialize, Serialize}; #[derive(Debug, Default, Serialize, Deserialize)] @@ -24,7 +25,7 @@ pub struct GlobalRoots { } /// Node in the Merkle-Patricia trie. -#[derive(Debug, Serialize, Deserialize)] +#[derive(Debug, Clone, Serialize, Deserialize)] #[serde(untagged)] pub enum MerkleNode { /// Represents a path to the highest non-zero descendant node. @@ -47,6 +48,21 @@ pub enum MerkleNode { }, } +impl MerkleNode { + // Taken from `bonsai-trie`: https://github.com/madara-alliance/bonsai-trie/blob/bfc6ad47b3cb8b75b1326bf630ca16e581f194c5/src/trie/merkle_node.rs#L234-L248 + pub fn compute_hash(&self) -> Felt { + match self { + Self::Binary { left, right } => Hash::hash(left, right), + Self::Edge { child, path, length } => { + let mut length_bytes = [0u8; 32]; + length_bytes[31] = *length; + let length = Felt::from_bytes_be(&length_bytes); + Hash::hash(child, path) + length + } + } + } +} + /// The response type for `starknet_getStorageProof` method. /// /// The requested storage proofs. Note that if a requested leaf has the default value, the path to @@ -142,40 +158,63 @@ impl From for ProofNode { fn from(value: MerkleNode) -> Self { match value { MerkleNode::Binary { left, right } => Self::Binary { left, right }, - MerkleNode::Edge { path, child, .. } => Self::Edge { child, path: felt_to_path(path) }, + MerkleNode::Edge { path, child, length } => { + Self::Edge { child, path: felt_to_path(path, length) } + } } } } -fn felt_to_path(felt: Felt) -> Path { - Path(BitSlice::from_slice(&felt.to_bytes_be())[5..].to_bitvec()) +fn felt_to_path(felt: Felt, length: u8) -> Path { + let length = length as usize; + let mut bits = BitVec::new(); + + // This function converts a Felt to a Path by preserving leading zeros + // that are semantically important in the Merkle tree path representation. + // + // Example: + // For a path "0000100" (length=7): + // - As an integer/hex: 0x4 (leading zeros get truncated) + // - As a Path: [0,0,0,0,1,0,0] (leading zeros preserved) + // + // We need to preserve these leading zeros because in a Merkle tree path: + // - Each bit represents a direction (left=0, right=1) + // - The position/index of each bit matters for the path traversal + // - "0000100" and "100" would represent different paths in the tree + for bit in &felt.to_bits_be()[256 - length..] { + bits.push(*bit); + } + + Path(bits) } fn path_to_felt(path: Path) -> Felt { - let mut arr = [0u8; 32]; - let slice = &mut BitSlice::from_slice_mut(&mut arr)[5..]; - slice[..path.len()].copy_from_bitslice(&path); - Felt::from_bytes_be(&arr) + let mut bytes = [0u8; 32]; + bytes.view_bits_mut()[256 - path.len()..].copy_from_bitslice(&path); + Felt::from_bytes_be(&bytes) } #[cfg(test)] mod tests { - use katana_primitives::felt; + use katana_trie::BitVec; use super::*; - // This test is assuming that the `path` field in `MerkleNode::Edge` is already a valid trie - // path value. + // Test cases taken from `bonsai-trie` crate #[rstest::rstest] - #[case(felt!("0x1234567890abcdef"))] - #[case(felt!("0xdeadbeef"))] - #[case(Felt::MAX)] - #[case(Felt::ZERO)] - fn test_path_felt_roundtrip(#[case] path_in_felt: Felt) { - let initial_path = felt_to_path(path_in_felt); - - let converted_felt = path_to_felt(initial_path.clone()); - let path = felt_to_path(converted_felt); - assert_eq!(initial_path, path); + #[case(&[0b10101010, 0b10101010])] + #[case(&[])] + #[case(&[0b10101010])] + #[case(&[0b00000000])] + #[case(&[0b11111111])] + #[case(&[0b11111111, 0b00000000, 0b10101010, 0b10101010, 0b11111111, 0b00000000, 0b10101010, 0b10101010, 0b11111111, 0b00000000, 0b10101010, 0b10101010])] + fn path_felt_rt(#[case] input: &[u8]) { + let path = Path(BitVec::from_slice(input)); + + let converted_felt = path_to_felt(path.clone()); + let converted_path = felt_to_path(converted_felt, path.len() as u8); + + assert_eq!(path, converted_path); + assert_eq!(path.len(), converted_path.len()); } } diff --git a/crates/katana/rpc/rpc/Cargo.toml b/crates/katana/rpc/rpc/Cargo.toml index 3910bbd2c9..5d032a4a0b 100644 --- a/crates/katana/rpc/rpc/Cargo.toml +++ b/crates/katana/rpc/rpc/Cargo.toml @@ -7,11 +7,6 @@ repository.workspace = true version.workspace = true [dependencies] -anyhow.workspace = true -dojo-metrics.workspace = true -futures.workspace = true -http.workspace = true -jsonrpsee = { workspace = true, features = [ "server" ] } katana-core.workspace = true katana-executor.workspace = true katana-pool.workspace = true @@ -21,6 +16,12 @@ katana-rpc-api.workspace = true katana-rpc-types.workspace = true katana-rpc-types-builder.workspace = true katana-tasks.workspace = true + +anyhow.workspace = true +dojo-metrics.workspace = true +futures.workspace = true +http.workspace = true +jsonrpsee = { workspace = true, features = [ "server" ] } metrics.workspace = true serde_json.workspace = true starknet.workspace = true @@ -32,6 +33,11 @@ tracing.workspace = true url.workspace = true [dev-dependencies] +katana-cairo.workspace = true +katana-node.workspace = true +katana-rpc-api = { workspace = true, features = [ "client" ] } +katana-trie.workspace = true + alloy = { git = "https://github.com/alloy-rs/alloy", features = [ "contract", "network", "node-bindings", "provider-http", "providers", "signer-local" ] } alloy-primitives = { workspace = true, features = [ "serde" ] } assert_matches.workspace = true @@ -40,10 +46,6 @@ dojo-test-utils.workspace = true dojo-utils.workspace = true indexmap.workspace = true jsonrpsee = { workspace = true, features = [ "client" ] } -katana-cairo.workspace = true -katana-node.workspace = true -katana-rpc-api = { workspace = true, features = [ "client" ] } -katana-trie.workspace = true num-traits.workspace = true rand.workspace = true rstest.workspace = true diff --git a/crates/katana/rpc/rpc/tests/proofs.rs b/crates/katana/rpc/rpc/tests/proofs.rs index 07d3a35f86..ec2b4f2066 100644 --- a/crates/katana/rpc/rpc/tests/proofs.rs +++ b/crates/katana/rpc/rpc/tests/proofs.rs @@ -6,17 +6,15 @@ 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::class::ClassHash; -use katana_primitives::hash::StarkHash; -use katana_primitives::{hash, Felt}; +use katana_primitives::class::{ClassHash, CompiledClassHash}; +use katana_primitives::Felt; use katana_rpc_api::starknet::StarknetApiClient; -use katana_rpc_types::trie::GetStorageProofResponse; -use katana_trie::bitvec::view::AsBits; -use katana_trie::bonsai::BitVec; -use katana_trie::MultiProof; -use starknet::accounts::Account; +use katana_trie::{compute_classes_trie_value, ClassesMultiProof, MultiProof}; +use starknet::accounts::{Account, ConnectedAccount, SingleOwnerAccount}; use starknet::core::types::BlockTag; -use starknet::macros::short_string; +use starknet::providers::jsonrpc::HttpTransport; +use starknet::providers::JsonRpcClient; +use starknet::signers::LocalWallet; mod common; @@ -74,14 +72,96 @@ async fn proofs_limit() { #[tokio::test] async fn classes_proofs() { - let sequencer = - TestSequencer::start(get_default_test_config(SequencingConfig::default())).await; + let cfg = get_default_test_config(SequencingConfig::default()); - let provider = sequencer.provider(); + let sequencer = TestSequencer::start(cfg).await; let account = sequencer.account(); - let path: PathBuf = PathBuf::from("tests/test_data/cairo1_contract.json"); - let (contract, compiled_class_hash) = common::prepare_contract_declaration_params(&path) + let (class_hash1, compiled_class_hash1) = + declare(&account, "tests/test_data/cairo1_contract.json").await; + let (class_hash2, compiled_class_hash2) = + declare(&account, "tests/test_data/cairo_l1_msg_contract.json").await; + let (class_hash3, compiled_class_hash3) = + declare(&account, "tests/test_data/test_sierra_contract.json").await; + + // We need to use the jsonrpsee client because `starknet-rs` doesn't yet support RPC 0.8.0 + let client = HttpClientBuilder::default().build(sequencer.url()).unwrap(); + + { + let class_hash = class_hash1; + let trie_entry = compute_classes_trie_value(compiled_class_hash1); + + let proofs = client + .get_storage_proof(BlockIdOrTag::Number(1), Some(vec![class_hash]), None, None) + .await + .expect("failed to get storage proof"); + + let results = ClassesMultiProof::from(MultiProof::from(proofs.classes_proof.nodes)) + .verify(proofs.global_roots.classes_tree_root, vec![class_hash]); + + assert_eq!(vec![trie_entry], results); + } + + { + let class_hash = class_hash2; + let trie_entry = compute_classes_trie_value(compiled_class_hash2); + + let proofs = client + .get_storage_proof(BlockIdOrTag::Number(2), Some(vec![class_hash]), None, None) + .await + .expect("failed to get storage proof"); + + let results = ClassesMultiProof::from(MultiProof::from(proofs.classes_proof.nodes)) + .verify(proofs.global_roots.classes_tree_root, vec![class_hash]); + + assert_eq!(vec![trie_entry], results); + } + + { + let class_hash = class_hash3; + let trie_entry = compute_classes_trie_value(compiled_class_hash3); + + let proofs = client + .get_storage_proof(BlockIdOrTag::Number(3), Some(vec![class_hash]), None, None) + .await + .expect("failed to get storage proof"); + + let results = ClassesMultiProof::from(MultiProof::from(proofs.classes_proof.nodes)) + .verify(proofs.global_roots.classes_tree_root, vec![class_hash]); + + assert_eq!(vec![trie_entry], results); + } + + { + let class_hashes = vec![class_hash1, class_hash2, class_hash3]; + let trie_entries = vec![ + compute_classes_trie_value(compiled_class_hash1), + compute_classes_trie_value(compiled_class_hash2), + compute_classes_trie_value(compiled_class_hash3), + ]; + + let proofs = client + .get_storage_proof( + BlockIdOrTag::Tag(BlockTag::Latest), + Some(class_hashes.clone()), + None, + None, + ) + .await + .expect("failed to get storage proof"); + + let results = ClassesMultiProof::from(MultiProof::from(proofs.classes_proof.nodes)) + .verify(proofs.global_roots.classes_tree_root, class_hashes.clone()); + + assert_eq!(trie_entries, results); + } +} + +async fn declare( + account: &SingleOwnerAccount, LocalWallet>, + path: impl Into, +) -> (ClassHash, CompiledClassHash) { + let (contract, compiled_class_hash) = common::prepare_contract_declaration_params(&path.into()) .expect("failed to prepare class declaration params"); let class_hash = contract.class_hash(); @@ -91,29 +171,9 @@ async fn classes_proofs() { .await .expect("failed to send declare tx"); - dojo_utils::TransactionWaiter::new(res.transaction_hash, &provider) + dojo_utils::TransactionWaiter::new(res.transaction_hash, account.provider()) .await .expect("failed to wait on tx"); - // We need to use the jsonrpsee client because `starknet-rs` doesn't yet support RPC 0.8.0 - let client = HttpClientBuilder::default().build(sequencer.url()).unwrap(); - - let GetStorageProofResponse { global_roots, classes_proof, .. } = client - .get_storage_proof(BlockIdOrTag::Tag(BlockTag::Latest), Some(vec![class_hash]), None, None) - .await - .expect("failed to get storage proof"); - - let key: BitVec = class_hash.to_bytes_be().as_bits()[5..].to_owned(); - let value = - hash::Poseidon::hash(&short_string!("CONTRACT_CLASS_LEAF_V0"), &compiled_class_hash); - - let classes_proof = MultiProof::from(classes_proof.nodes); - - // the returned data is the list of values corresponds to the [key] - let results = classes_proof - .verify_proof::(global_roots.classes_tree_root, [key], 251) - .collect::, _>>() - .expect("failed to verify proofs"); - - assert_eq!(vec![value], results); + (class_hash, compiled_class_hash) } diff --git a/crates/katana/rpc/rpc/tests/test_data/test_sierra_contract.json b/crates/katana/rpc/rpc/tests/test_data/test_sierra_contract.json new file mode 120000 index 0000000000..e92339e522 --- /dev/null +++ b/crates/katana/rpc/rpc/tests/test_data/test_sierra_contract.json @@ -0,0 +1 @@ +../../../../contracts/build/account_with_dummy_validate.sierra.json \ No newline at end of file diff --git a/crates/katana/trie/src/classes.rs b/crates/katana/trie/src/classes.rs index 9381686ca4..c74f3dffc1 100644 --- a/crates/katana/trie/src/classes.rs +++ b/crates/katana/trie/src/classes.rs @@ -8,11 +8,32 @@ use starknet_types_core::hash::{Poseidon, StarkHash}; use crate::id::CommitId; +#[derive(Debug)] +pub struct ClassesMultiProof(pub MultiProof); + +impl ClassesMultiProof { + // TODO: maybe perform results check in this method as well. make it accept the compiled class + // hashes + pub fn verify(&self, root: Felt, class_hashes: Vec) -> Vec { + crate::verify_proof::(&self.0, root, class_hashes) + } +} + +impl From for ClassesMultiProof { + fn from(value: MultiProof) -> Self { + Self(value) + } +} + #[derive(Debug)] pub struct ClassesTrie { trie: crate::BonsaiTrie, } +///////////////////////////////////////////////////// +// ClassesTrie implementations +///////////////////////////////////////////////////// + impl ClassesTrie { const BONSAI_IDENTIFIER: &'static [u8] = b"classes"; @@ -34,9 +55,7 @@ where DB: BonsaiDatabase + BonsaiPersistentDatabase, { pub fn insert(&mut self, hash: ClassHash, compiled_hash: CompiledClassHash) { - // https://docs.starknet.io/architecture-and-concepts/network-architecture/starknet-state/#classes_trie - const CONTRACT_CLASS_LEAF_V0: Felt = short_string!("CONTRACT_CLASS_LEAF_V0"); - let value = Poseidon::hash(&CONTRACT_CLASS_LEAF_V0, &compiled_hash); + let value = compute_classes_trie_value(compiled_hash); self.trie.insert(Self::BONSAI_IDENTIFIER, hash, value) } @@ -44,3 +63,9 @@ where self.trie.commit(block.into()) } } + +pub fn compute_classes_trie_value(compiled_class_hash: CompiledClassHash) -> Felt { + // https://docs.starknet.io/architecture-and-concepts/network-architecture/starknet-state/#classes_trie + const CONTRACT_CLASS_LEAF_V0: Felt = short_string!("CONTRACT_CLASS_LEAF_V0"); + Poseidon::hash(&CONTRACT_CLASS_LEAF_V0, &compiled_class_hash) +} diff --git a/crates/katana/trie/src/lib.rs b/crates/katana/trie/src/lib.rs index 11a85fafc8..0753352c80 100644 --- a/crates/katana/trie/src/lib.rs +++ b/crates/katana/trie/src/lib.rs @@ -1,10 +1,7 @@ -use bitvec::array::BitArray; -use bitvec::order::Msb0; -use bitvec::vec::BitVec; use bitvec::view::AsBits; -pub use bonsai::{MultiProof, Path, ProofNode}; -pub use bonsai_trie::{BonsaiDatabase, BonsaiPersistentDatabase}; -use bonsai_trie::{BonsaiStorage, BonsaiStorageConfig}; +pub use bonsai::{BitVec, MultiProof, Path, ProofNode}; +use bonsai_trie::BonsaiStorage; +pub use bonsai_trie::{BonsaiDatabase, BonsaiPersistentDatabase, BonsaiStorageConfig}; use katana_primitives::class::ClassHash; use katana_primitives::Felt; use starknet_types_core::hash::{Pedersen, StarkHash}; @@ -15,7 +12,7 @@ mod contracts; mod id; mod storages; -pub use classes::ClassesTrie; +pub use classes::*; pub use contracts::ContractsTrie; pub use id::CommitId; pub use storages::StoragesTrie; @@ -26,7 +23,7 @@ pub use storages::StoragesTrie; /// having to handle how to transform the keys into the internal keys used by the trie. /// This struct is not meant to be used directly, and instead use the specific tries that have /// been derived from it, [`ClassesTrie`], [`ContractsTrie`], or [`StoragesTrie`]. -pub(crate) struct BonsaiTrie +pub struct BonsaiTrie where DB: BonsaiDatabase, Hash: StarkHash + Send + Sync, @@ -59,13 +56,8 @@ where self.storage.root_hash(id).expect("failed to get trie root") } - pub fn multiproof(&mut self, id: &[u8], mut keys: Vec) -> MultiProof { - keys.sort(); - let keys = keys - .into_iter() - .map(|key| BitArray::new(key.to_bytes_be())) - .map(|hash| hash.as_bitslice()[5..].to_owned()); - + pub fn multiproof(&mut self, id: &[u8], keys: Vec) -> MultiProof { + let keys = keys.into_iter().map(|key| key.to_bytes_be().as_bits()[5..].to_owned()); self.storage.get_multi_proof(id, keys).expect("failed to get multiproof") } } @@ -76,7 +68,7 @@ where Hash: StarkHash + Send + Sync, { pub fn insert(&mut self, id: &[u8], key: Felt, value: Felt) { - let key: BitVec = key.to_bytes_be().as_bits()[5..].to_owned(); + let key: BitVec = key.to_bytes_be().as_bits()[5..].to_owned(); self.storage.insert(id, &key, &value).unwrap(); } @@ -132,6 +124,15 @@ pub fn compute_contract_state_hash( Pedersen::hash(&hash, &CONTRACT_STATE_HASH_VERSION) } +pub fn verify_proof( + proofs: &MultiProof, + root: Felt, + keys: Vec, +) -> Vec { + let keys = keys.into_iter().map(|f| f.to_bytes_be().as_bits()[5..].to_owned()); + proofs.verify_proof::(root, keys, 251).collect::, _>>().unwrap() +} + #[cfg(test)] mod tests {