From 256e207fd42e6ba04800e3db530cb36d3e0e0e25 Mon Sep 17 00:00:00 2001 From: petscheit Date: Thu, 25 Apr 2024 19:59:39 +0200 Subject: [PATCH] feat: add account profs --- tests/rust/Cargo.lock | 3 + tests/rust/Cargo.toml | 3 + tests/rust/src/main.rs | 160 +++++++++++++++++++++++++++++++++-------- 3 files changed, 136 insertions(+), 30 deletions(-) diff --git a/tests/rust/Cargo.lock b/tests/rust/Cargo.lock index c010017..79cae1f 100644 --- a/tests/rust/Cargo.lock +++ b/tests/rust/Cargo.lock @@ -2091,11 +2091,14 @@ dependencies = [ "alloy-transport", "alloy-transport-http", "eth-trie-proofs", + "eth_trie", + "ethereum-types", "rand", "reqwest", "serde", "serde_json", "serde_with", + "tiny-keccak", "tokio", ] diff --git a/tests/rust/Cargo.toml b/tests/rust/Cargo.toml index 1bbd252..e2c5bad 100644 --- a/tests/rust/Cargo.toml +++ b/tests/rust/Cargo.toml @@ -19,6 +19,9 @@ serde = "1.0.197" serde_with = { version = "3.7.0", features = ["hex"] } serde_json = "1.0.114" rand = "0.8.4" +tiny-keccak = "2.0.2" +eth_trie = "0.4.0" +ethereum-types = "0.14.1" [dev-dependencies] tokio = { version = "1", features = ["full", "test-util"] } \ No newline at end of file diff --git a/tests/rust/src/main.rs b/tests/rust/src/main.rs index 926ceee..b074ab3 100644 --- a/tests/rust/src/main.rs +++ b/tests/rust/src/main.rs @@ -1,33 +1,39 @@ -use alloy_primitives::{B256, BlockNumber, U256}; +use alloy_network::Ethereum; + +use alloy_primitives::private::alloy_rlp; +use alloy_primitives::{B256, U256}; use alloy_provider::{Provider, ProviderBuilder, RootProvider}; +use alloy_rpc_client::RpcClient; +use alloy_rpc_types::{Block, BlockTransactions, Transaction}; use alloy_transport::{RpcError, TransportErrorKind}; use alloy_transport_http::Http; -use alloy_network::Ethereum; -use alloy_primitives::private::alloy_rlp; +use eth_trie::MemoryDB; +use eth_trie::{EthTrie, Trie, TrieError as EthTrieError}; +use eth_trie_proofs::tx_receipt_trie::TxReceiptsMptHandler; use eth_trie_proofs::tx_trie::TxsMptHandler; -use alloy_rpc_client::RpcClient; -use alloy_rpc_types::BlockTransactions; +use eth_trie_proofs::Error as TrieError; +use ethereum_types::H256; +use rand::Rng; use reqwest::Client; use serde::Serialize; use serde_with::serde_as; -use alloy_primitives::hex::FromHex; use std::fs::File; use std::io::Write; -use eth_trie_proofs::tx_receipt_trie::TxReceiptsMptHandler; -use eth_trie_proofs::Error as TrieError; -use rand::Rng; +use std::sync::Arc; +use tiny_keccak::{Hasher, Keccak}; #[derive(Debug)] enum Error { Trie(TrieError), Transport(TransportErrorKind), Rpc(RpcError), + EthTrie(EthTrieError), } struct Fetcher { provider: RootProvider>, tx: TxsMptHandler, - receipt: TxReceiptsMptHandler + receipt: TxReceiptsMptHandler, } #[derive(Debug, Serialize)] @@ -36,6 +42,8 @@ enum ProofType { TxProof, #[serde(rename = "receipt_proof")] ReceiptProof, + #[serde(rename = "account_proof")] + AccountProof, } #[serde_with::serde_as] @@ -52,28 +60,35 @@ pub struct MptProof { impl Fetcher { fn new(rpc_url: &str) -> Result { - let http = Http::::new(rpc_url.clone().to_string().parse().unwrap()); + let http = Http::::new(rpc_url.to_string().parse().unwrap()); let provider = ProviderBuilder::<_, Ethereum>::new() .provider(RootProvider::new(RpcClient::new(http, true))); let tx = TxsMptHandler::new(rpc_url)?; let receipt = TxReceiptsMptHandler::new(rpc_url)?; - Ok(Fetcher { provider, tx, receipt }) + Ok(Fetcher { + provider, + tx, + receipt, + }) } - async fn generate_block_tx_proofs(&mut self, block_number: u64) -> Result, Error> { + async fn generate_block_tx_proofs( + &mut self, + block_number: u64, + ) -> Result, Error> { let tx_count = self.get_block_txs_count(block_number).await?; self.tx.build_tx_tree_from_block(block_number).await?; let mut proofs = vec![]; for i in 0..tx_count { let trie_proof = self.tx.get_proof(i)?; let root = self.tx.get_root().unwrap(); - + // ensure the proof is valid self.tx.verify_proof(i, trie_proof.clone())?; - + proofs.push(MptProof { proof: trie_proof, - key: generate_key_le(i), + key: generate_key_from_index(i), root, kind: ProofType::TxProof, }); @@ -82,9 +97,14 @@ impl Fetcher { Ok(proofs) } - async fn generate_block_receipt_proofs(&mut self, block_number: u64)-> Result, Error> { + async fn generate_block_receipt_proofs( + &mut self, + block_number: u64, + ) -> Result, Error> { let tx_count = self.get_block_txs_count(block_number).await?; - self.receipt.build_tx_receipts_tree_from_block(block_number).await?; + self.receipt + .build_tx_receipts_tree_from_block(block_number) + .await?; let mut proofs = vec![]; for i in 0..tx_count { let trie_proof = self.receipt.get_proof(i)?; @@ -92,10 +112,10 @@ impl Fetcher { // ensure the proof is valid self.receipt.verify_proof(i, trie_proof.clone())?; - + proofs.push(MptProof { proof: trie_proof, - key: generate_key_le(i), + key: generate_key_from_index(i), root, kind: ProofType::ReceiptProof, }); @@ -104,22 +124,85 @@ impl Fetcher { Ok(proofs) } + async fn get_account_proofs(&mut self, block_number: u64) -> Result, Error> { + let block = self.get_block(block_number).await?; + + // fetch all txs from the block + let txs = self.get_block_txs(block_number).await?; + + // retrieve sender and receiver addresses + let mut addresses = vec![]; + for tx in txs { + addresses.push(tx.from); + if let Some(to) = tx.to { + addresses.push(to); + } + } + + let mut proofs = vec![]; + let memdb = Arc::new(MemoryDB::new(true)); + let trie = EthTrie::new(memdb.clone()); + + for address in addresses { + let proof: Vec<_> = self + .provider + .get_proof(address, vec![], Some(block_number.into())) + .await? + .account_proof + .iter() + .map(|x| x.as_ref().to_vec()) + .collect(); + + // ensure the proof is valid + trie.verify_proof( + H256::from_slice(block.header.state_root.as_slice()), + generate_key_from_address(address.0.as_slice()).as_slice(), + proof.clone(), + )?; + + proofs.push(MptProof { + proof, + key: generate_key_from_address(address.0.as_slice()), + root: block.header.state_root, + kind: ProofType::AccountProof, + }); + } + + Ok(proofs) + } + + async fn get_block(&self, block_number: u64) -> Result { + match self.provider.get_block(block_number.into(), false).await { + Ok(Some(block)) => Ok(block), + _ => panic!(), + } + } + + async fn get_block_txs(&self, block_number: u64) -> Result, Error> { + let block = self.provider.get_block(block_number.into(), true).await?; + match block.unwrap().transactions { + BlockTransactions::Full(txs) => Ok(txs), + _ => panic!(), + } + } + async fn get_block_txs_count(&self, block_number: u64) -> Result { let block = self.provider.get_block(block_number.into(), false).await?; match block.unwrap().transactions { - BlockTransactions::Hashes(hashes) => { - Ok(hashes.len() as u64) - }, - _ => panic!() + BlockTransactions::Hashes(hashes) => Ok(hashes.len() as u64), + _ => panic!(), } } - async fn generate_random_block_proofs(&mut self, num_blocks: u32) -> Result, Error> { + async fn generate_random_block_proofs( + &mut self, + num_blocks: u32, + ) -> Result, Error> { let mut rng = rand::thread_rng(); let mut proofs = vec![]; // Fetch the current block number - let current_block_number: u64 = self.provider.get_block_number().await?.into(); + let current_block_number: u64 = self.provider.get_block_number().await?; for _ in 0..num_blocks { // As there is a bug in the underlying alloy library, we need to start with the byzantium hardfork. https://github.com/alloy-rs/alloy/issues/630 let block_number = rng.gen_range(4370000..current_block_number); @@ -135,6 +218,9 @@ impl Fetcher { Err(Error::Trie(TrieError::TxNotFound)) => continue, // found block no txs Err(e) => return Err(e), } + + let account_proofs = self.get_account_proofs(block_number).await?; + proofs.extend(account_proofs); } Ok(proofs) @@ -144,12 +230,13 @@ impl Fetcher { #[tokio::main] async fn main() -> Result<(), Error> { // Create a new Fetcher instance - let mut fetcher = Fetcher::new("https://ethereum-rpc.publicnode.com")?; + let mut fetcher = + Fetcher::new("")?; // Generate random block proofs - let proofs = fetcher.generate_random_block_proofs(25).await?; + let proofs = fetcher.generate_random_block_proofs(5).await?; + // let proofs = fetcher.get_account_proofs(19733390).await?; - // Export the proofs to a JSON file export_batch(proofs).unwrap(); Ok(()) @@ -168,8 +255,15 @@ fn export_batch(proofs: Vec) -> Result<(), Error> { Ok(()) } +fn generate_key_from_address(address: &[u8]) -> Vec { + let mut hasher = Keccak::v256(); + hasher.update(address); + let mut output = [0u8; 32]; + hasher.finalize(&mut output); + output.to_vec() +} -fn generate_key_le(index: u64) -> Vec { +fn generate_key_from_index(index: u64) -> Vec { alloy_rlp::encode(U256::from(index)) } @@ -179,6 +273,12 @@ impl From for Error { } } +impl From for Error { + fn from(error: EthTrieError) -> Self { + Error::EthTrie(error) + } +} + impl From for Error { fn from(error: TransportErrorKind) -> Self { Error::Transport(error)