add get_strorage_at to client

src/client/client.rs

@@ -29,23 +29,36 @@ impl Client {
         self.consensus.sync().await
     }
 
-    pub async fn get_balance(&mut self, address: Address) -> Result<U256> {
+    pub async fn get_balance(&mut self, address: &Address) -> Result<U256> {
         let payload = self.consensus.get_execution_payload().await?;
-        let account = self.execution.get_account(&address, &payload).await?;
+        let account = self.execution.get_account(&address, None, &payload).await?;
         Ok(account.balance)
     }
 
-    pub async fn get_nonce(&mut self, address: Address) -> Result<U256> {
+    pub async fn get_nonce(&mut self, address: &Address) -> Result<U256> {
         let payload = self.consensus.get_execution_payload().await?;
-        let account = self.execution.get_account(&address, &payload).await?;
+        let account = self.execution.get_account(&address, None, &payload).await?;
         Ok(account.nonce)
     }
 
-    pub async fn get_code(&mut self, address: Address) -> Result<Vec<u8>> {
+    pub async fn get_code(&mut self, address: &Address) -> Result<Vec<u8>> {
         let payload = self.consensus.get_execution_payload().await?;
         self.execution.get_code(&address, &payload).await
     }
 
+    pub async fn get_storage_at(&mut self, address: &Address, slot: U256) -> Result<U256> {
+        let payload = self.consensus.get_execution_payload().await?;
+        let account = self
+            .execution
+            .get_account(address, Some(&[slot]), &payload)
+            .await?;
+        let value = account.slots.get(&slot);
+        match value {
+            Some(value) => Ok(*value),
+            None => Err(eyre::eyre!("Slot Not Found")),
+        }
+    }
+
     pub fn get_header(&self) -> &Header {
         self.consensus.get_head()
     }

src/execution/execution.rs

@@ -1,9 +1,15 @@
-use ethers::prelude::{Address, H256, U256};
+use std::collections::HashMap;
+
+use ethers::abi::AbiEncode;
+use ethers::prelude::{Address, U256};
 use ethers::utils::keccak256;
+use ethers::utils::rlp::encode;
 use eyre::Result;
 
 use super::proof::{encode_account, verify_proof};
 use super::rpc::Rpc;
+use super::types::Account;
+use crate::common::utils::hex_str_to_bytes;
 use crate::consensus::types::ExecutionPayload;
 
 pub struct ExecutionClient {
@@ -19,9 +25,14 @@ impl ExecutionClient {
     pub async fn get_account(
         &self,
         address: &Address,
+        slots: Option<&[U256]>,
         payload: &ExecutionPayload,
     ) -> Result<Account> {
-        let proof = self.rpc.get_proof(&address, payload.block_number).await?;
+        let slots = slots.unwrap_or(&[]);
+        let proof = self
+            .rpc
+            .get_proof(&address, slots, payload.block_number)
+            .await?;
 
         let account_path = keccak256(address.as_bytes()).to_vec();
         let account_encoded = encode_account(&proof);
@@ -37,16 +48,39 @@ impl ExecutionClient {
             eyre::bail!("Invalid Proof");
         }
 
+        let mut slot_map = HashMap::new();
+
+        for storage_proof in proof.storage_proof {
+            let key = hex_str_to_bytes(&storage_proof.key.encode_hex())?;
+            let value = encode(&storage_proof.value).to_vec();
+
+            let key_hash = keccak256(key);
+
+            let is_valid = verify_proof(
+                &storage_proof.proof,
+                &proof.storage_hash.as_bytes().to_vec(),
+                &key_hash.to_vec(),
+                &value,
+            );
+
+            if !is_valid {
+                eyre::bail!("Invalid Proof");
+            }
+
+            slot_map.insert(storage_proof.key, storage_proof.value);
+        }
+
         Ok(Account {
             balance: proof.balance,
             nonce: proof.nonce,
             code_hash: proof.code_hash,
             storage_hash: proof.storage_hash,
+            slots: slot_map,
         })
     }
 
     pub async fn get_code(&self, address: &Address, payload: &ExecutionPayload) -> Result<Vec<u8>> {
-        let account = self.get_account(address, payload).await?;
+        let account = self.get_account(address, None, payload).await?;
         let code = self.rpc.get_code(address, payload.block_number).await?;
 
         let code_hash = keccak256(&code).into();
@@ -58,10 +92,3 @@ impl ExecutionClient {
         Ok(code)
     }
 }
-
-pub struct Account {
-    pub balance: U256,
-    pub nonce: U256,
-    pub code_hash: H256,
-    pub storage_hash: H256,
-}

src/execution/proof.rs

@@ -25,7 +25,7 @@ pub fn verify_proof(proof: &Vec<Vec<u8>>, root: &Vec<u8>, path: &Vec<u8>, value:
                     return false;
                 }
             } else {
-                expected_hash = node_list[1].clone();
+                panic!("not implemented");
             }
         } else {
             return false;
@@ -53,7 +53,3 @@ pub fn encode_account(proof: &Proof) -> Vec<u8> {
     let encoded = stream.out();
     encoded.to_vec()
 }
-
-pub fn get_account_path(addr: &Vec<u8>) -> Vec<u8> {
-    keccak256(addr).to_vec()
-}

src/execution/rpc.rs

@@ -1,4 +1,5 @@
-use ethers::prelude::Address;
+use ethers::abi::AbiEncode;
+use ethers::prelude::{Address, U256};
 use eyre::Result;
 use jsonrpsee::{
     core::client::ClientT,
@@ -21,11 +22,15 @@ impl Rpc {
         }
     }
 
-    pub async fn get_proof(&self, address: &Address, block: u64) -> Result<Proof> {
+    pub async fn get_proof(&self, address: &Address, slots: &[U256], block: u64) -> Result<Proof> {
         let client = self.client()?;
         let block_hex = u64_to_hex_string(block);
         let addr_hex = address_to_hex_string(address);
-        let params = rpc_params!(addr_hex, [""], block_hex);
+        let slots = slots
+            .iter()
+            .map(|slot| slot.encode_hex())
+            .collect::<Vec<String>>();
+        let params = rpc_params!(addr_hex, slots.as_slice(), block_hex);
         Ok(client.request("eth_getProof", params).await?)
     }
 

src/execution/types.rs

@@ -1,3 +1,5 @@
+use std::collections::HashMap;
+
 use ethers::prelude::{Address, H256, U256};
 use eyre::Result;
 use serde::de::Error;
@@ -15,6 +17,24 @@ pub struct Proof {
     pub storage_hash: H256,
     #[serde(deserialize_with = "proof_deserialize")]
     pub account_proof: Vec<Vec<u8>>,
+    pub storage_proof: Vec<StorageProof>,
+}
+
+#[derive(Deserialize, Debug)]
+#[serde(rename_all = "camelCase")]
+pub struct StorageProof {
+    pub key: U256,
+    pub value: U256,
+    #[serde(deserialize_with = "proof_deserialize")]
+    pub proof: Vec<Vec<u8>>,
+}
+
+pub struct Account {
+    pub balance: U256,
+    pub nonce: U256,
+    pub code_hash: H256,
+    pub storage_hash: H256,
+    pub slots: HashMap<U256, U256>,
 }
 
 fn proof_deserialize<'de, D>(deserializer: D) -> Result<Vec<Vec<u8>>, D::Error>

src/main.rs

@@ -1,6 +1,6 @@
 use std::str::FromStr;
 
-use ethers::prelude::Address;
+use ethers::prelude::{Address, U256};
 use eyre::Result;
 
 use client::Client;
@@ -23,13 +23,15 @@ async fn main() -> Result<()> {
     println!("synced up to slot: {}", header.slot);
 
     let address = Address::from_str("0x14f9D4aF749609c1438528C0Cce1cC3f6D411c47")?;
-    let balance = client.get_balance(address).await?;
-    let nonce = client.get_nonce(address).await?;
-    let code = client.get_code(address).await?;
+    let balance = client.get_balance(&address).await?;
+    let nonce = client.get_nonce(&address).await?;
+    let code = client.get_code(&address).await?;
+    let storage_value = client.get_storage_at(&address, U256::from(0)).await?;
 
     println!("balance: {}", balance);
     println!("nonce: {}", nonce);
     println!("code: 0x{}...", hex::encode(code[..5].to_vec()));
+    println!("value at slot 0: 0x{:x}", storage_value);
 
     Ok(())
 }