From a7a5942451f6a509f27e9f7c9c67aa3ac9dffffb Mon Sep 17 00:00:00 2001
From: Noah Citron <noah@jeff.org>
Date: Sun, 21 Aug 2022 09:13:56 -0400
Subject: [PATCH] add client module

---
 src/client.rs    | 45 +++++++++++++++++++++++++++++++++++++++++++++
 src/consensus.rs |  6 ++++--
 src/execution.rs | 22 +++++++++++++++++-----
 src/main.rs      | 27 ++++++++++++---------------
 4 files changed, 78 insertions(+), 22 deletions(-)
 create mode 100644 src/client.rs

diff --git a/src/client.rs b/src/client.rs
new file mode 100644
index 0000000..0c79985
--- /dev/null
+++ b/src/client.rs
@@ -0,0 +1,45 @@
+use ethers::prelude::{Address, U256};
+use eyre::Result;
+
+use crate::{consensus::ConsensusClient, execution::ExecutionClient, consensus_rpc::Header};
+
+pub struct Client {
+    consensus: ConsensusClient,
+    execution: ExecutionClient,
+}
+
+impl Client {
+    pub async fn new(
+        consensus_rpc: &str,
+        execution_rpc: &str,
+        checkpoint_hash: &str,
+    ) -> Result<Self> {
+        let consensus = ConsensusClient::new(consensus_rpc, checkpoint_hash).await?;
+        let execution = ExecutionClient::new(execution_rpc);
+
+        Ok(Client {
+            consensus,
+            execution,
+        })
+    }
+
+    pub async fn sync(&mut self) -> Result<()> {
+        self.consensus.sync().await
+    }
+
+    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?;
+        Ok(account.balance)
+    }
+
+    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?;
+        Ok(account.nonce)
+    }
+
+    pub fn get_header(&self) -> &Header {
+        self.consensus.get_head()
+    }
+}
diff --git a/src/consensus.rs b/src/consensus.rs
index 636794b..7e63c39 100644
--- a/src/consensus.rs
+++ b/src/consensus.rs
@@ -62,6 +62,10 @@ impl ConsensusClient {
         }
     }
 
+    pub fn get_head(&self) -> &Header {
+        &self.store.header
+    }
+
     pub async fn sync(&mut self) -> Result<()> {
         let current_period = calc_sync_period(self.store.header.slot);
         let updates = self.consensus_rpc.get_updates(current_period).await?;
@@ -85,8 +89,6 @@ impl ConsensusClient {
         self.verify_update(&mut finality_update_generic)?;
         self.apply_update(&finality_update_generic);
 
-        println!("synced up to slot: {}", self.store.header.slot);
-
         self.consensus_rpc.get_block(self.store.header.slot).await?;
 
         Ok(())
diff --git a/src/execution.rs b/src/execution.rs
index 75790bc..0a4953d 100644
--- a/src/execution.rs
+++ b/src/execution.rs
@@ -1,4 +1,4 @@
-use ethers::prelude::{Address, U256};
+use ethers::prelude::{Address, U256, H256};
 use ethers::utils::keccak256;
 use eyre::Result;
 
@@ -16,13 +16,13 @@ impl ExecutionClient {
         ExecutionClient { execution_rpc }
     }
 
-    pub async fn get_balance(&self, account: &Address, payload: &ExecutionPayload) -> Result<U256> {
+    pub async fn get_account(&self, address: &Address, payload: &ExecutionPayload) -> Result<Account> {
         let proof = self
             .execution_rpc
-            .get_proof(&account, payload.block_number)
+            .get_proof(&address, payload.block_number)
             .await?;
 
-        let account_path = keccak256(account.as_bytes()).to_vec();
+        let account_path = keccak256(address.as_bytes()).to_vec();
         let account_encoded = encode_account(&proof);
 
         let is_valid = verify_proof(
@@ -36,6 +36,18 @@ impl ExecutionClient {
             eyre::bail!("Invalid Proof");
         }
 
-        Ok(proof.balance)
+        Ok(Account {
+            balance: proof.balance,
+            nonce: proof.nonce,
+            code_hash: proof.code_hash,
+            storage_hash: proof.storage_hash,
+        })
     }
 }
+
+pub struct Account {
+    pub balance: U256,
+    pub nonce: U256,
+    pub code_hash: H256,
+    pub storage_hash: H256,
+}
diff --git a/src/main.rs b/src/main.rs
index 86f90f1..1221179 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -3,9 +3,9 @@ use std::str::FromStr;
 use ethers::prelude::Address;
 use eyre::Result;
 
-use consensus::*;
-use execution::*;
+use client::Client;
 
+pub mod client;
 pub mod consensus;
 pub mod consensus_rpc;
 pub mod execution;
@@ -15,24 +15,21 @@ pub mod utils;
 
 #[tokio::main]
 async fn main() -> Result<()> {
-    let rpc = "http://testing.prater.beacon-api.nimbus.team";
+    let consensus_rpc = "http://testing.prater.beacon-api.nimbus.team";
+    let execution_rpc = "https://eth-goerli.g.alchemy.com:443/v2/o_8Qa9kgwDPf9G8sroyQ-uQtyhyWa3ao";
     let checkpoint = "0x172128eadf1da46467f4d6a822206698e2d3f957af117dd650954780d680dc99";
-    let mut client = ConsensusClient::new(rpc, checkpoint).await?;
-
-    let rpc = "https://eth-goerli.g.alchemy.com:443/v2/o_8Qa9kgwDPf9G8sroyQ-uQtyhyWa3ao";
-    let execution = ExecutionClient::new(rpc);
 
+    let mut client = Client::new(consensus_rpc, execution_rpc, checkpoint).await?;
     client.sync().await?;
 
-    let payload = client.get_execution_payload().await?;
-    println!(
-        "verified execution block hash: {}",
-        hex::encode(&payload.block_hash)
-    );
+    let header = client.get_header();
+    println!("synced up to slot: {}", header.slot);
 
-    let addr = Address::from_str("0x25c4a76E7d118705e7Ea2e9b7d8C59930d8aCD3b")?;
-    let balance = execution.get_balance(&addr, &payload).await?;
-    println!("verified account balance: {}", balance);
+    let address = Address::from_str("0xe0Fa62CD8543473627D337fAe1212d4E639EE932")?;
+    let balance = client.get_balance(address).await?;
+    let nonce = client.get_nonce(address).await?;
+    println!("balance: {}", balance);
+    println!("nonce: {}", nonce);
 
     Ok(())
 }