From dbea709596386fc57ac6ad32096ae4596c2da77e Mon Sep 17 00:00:00 2001 From: Wolfgang Welz Date: Fri, 24 Jan 2025 17:11:49 +0100 Subject: [PATCH] add account queries --- crates/steel/src/account.rs | 160 ++++++++++++++++++++++++++++++ crates/steel/src/host/db/alloy.rs | 2 +- crates/steel/src/host/mod.rs | 31 ++++++ crates/steel/src/lib.rs | 2 + crates/steel/src/verifier.rs | 31 ------ crates/steel/tests/steel.rs | 40 +++++++- 6 files changed, 231 insertions(+), 35 deletions(-) create mode 100644 crates/steel/src/account.rs diff --git a/crates/steel/src/account.rs b/crates/steel/src/account.rs new file mode 100644 index 00000000..509f5fe0 --- /dev/null +++ b/crates/steel/src/account.rs @@ -0,0 +1,160 @@ +// Copyright 2025 RISC Zero, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Types related to account queries. +pub use revm::primitives::{AccountInfo, Bytecode}; + +use crate::{state::WrapStateDb, EvmBlockHeader, GuestEvmEnv}; +use alloy_primitives::Address; +use anyhow::Result; +use revm::Database as RevmDatabase; + +/// Represents an EVM account query. +/// +/// ### Usage +/// - **Preflight calls on the Host:** To prepare the account query on the host environment and +/// build the necessary proof, use [Account::preflight]. +/// - **Calls in the Guest:** To initialize the account query in the guest, use [Account::new]. +/// +/// ### Examples +/// ```rust,no_run +/// # use risc0_steel::{Account, ethereum::EthEvmEnv}; +/// # use alloy_primitives::address; +/// +/// # #[tokio::main(flavor = "current_thread")] +/// # async fn main() -> anyhow::Result<()> { +/// let account_address = address!("F977814e90dA44bFA03b6295A0616a897441aceC"); +/// +/// // Host: +/// let url = "https://ethereum-rpc.publicnode.com".parse()?; +/// let mut env = EthEvmEnv::builder().rpc(url).build().await?; +/// let account = Account::preflight(account_address, &mut env); +/// let info = account.bytecode(true).info().await?; +/// +/// let evm_input = env.into_input().await?; +/// +/// // Guest: +/// let env = evm_input.into_env(); +/// let account = Account::new(account_address, &env); +/// let info = account.bytecode(true).info(); +/// +/// # Ok(()) +/// # } +/// ``` +pub struct Account { + address: Address, + env: E, + code: bool, +} + +impl Account { + /// Sets whether to fetch the bytecode for this account. + /// + /// If set to `true`, the bytecode will be fetched when calling [Account::info]. + pub fn bytecode(mut self, code: bool) -> Self { + self.code = code; + self + } +} + +impl<'a, H: EvmBlockHeader> Account<&'a GuestEvmEnv> { + /// Constructor for querying an Ethereum account in the guest. + pub fn new(address: Address, env: &'a GuestEvmEnv) -> Self { + Self { + address, + env, + code: false, + } + } + + /// Attempts to get the [AccountInfo] for the corresponding account and returns an error if the + /// query fails. + /// + /// In general, it's recommended to use [Account::info] unless explicit error handling is + /// required. + pub fn try_info(self) -> Result { + let mut db = WrapStateDb::new(self.env.db()); + let mut info = db.basic(self.address)?.unwrap_or_default(); + if self.code && info.code.is_none() { + let code = db.code_by_hash(info.code_hash)?; + info.code = Some(code); + } + + Ok(info) + } + + /// Gets the [AccountInfo] for the corresponding account and panics on failure. + /// + /// A convenience wrapper for [Account::try_info], panicking if the query fails. Useful when + /// success is expected. + pub fn info(self) -> AccountInfo { + self.try_info().unwrap() + } +} + +#[cfg(feature = "host")] +mod host { + use super::*; + use crate::host::HostEvmEnv; + use anyhow::Context; + use std::error::Error as StdError; + + impl<'a, D, H, C> Account<&'a mut HostEvmEnv> + where + D: RevmDatabase + Send + 'static, + ::Error: StdError + Send + Sync + 'static, + { + /// Constructor for preflighting queries to an Ethereum account on the host. + /// + /// Initializes the environment for querying account information, fetching necessary data + /// via the [Provider], and generating a storage proof for any accessed elements using + /// [EvmEnv::into_input]. + /// + /// [EvmEnv::into_input]: crate::EvmEnv::into_input + /// [EvmEnv]: crate::EvmEnv + /// [Provider]: alloy::providers::Provider + pub fn preflight(address: Address, env: &'a mut HostEvmEnv) -> Self { + Self { + address, + env, + code: false, + } + } + + /// Gets the [AccountInfo] for the corresponding account using an [EvmEnv] constructed with + /// [Account::preflight]. + /// + /// [EvmEnv]: crate::EvmEnv + pub async fn info(self) -> Result { + log::info!("Executing preflight querying account {}", &self.address); + + let mut info = self + .env + .spawn_with_db(move |db| db.basic(self.address)) + .await + .context("failed to get basic account information")? + .unwrap_or_default(); + if self.code && info.code.is_none() { + let code = self + .env + .spawn_with_db(move |db| db.code_by_hash(info.code_hash)) + .await + .context("failed to get account code by its hash")?; + info.code = Some(code); + } + + Ok(info) + } + } +} diff --git a/crates/steel/src/host/db/alloy.rs b/crates/steel/src/host/db/alloy.rs index 878231de..76507556 100644 --- a/crates/steel/src/host/db/alloy.rs +++ b/crates/steel/src/host/db/alloy.rs @@ -114,7 +114,7 @@ impl> Database for AlloyDb { config_id: B256, } +impl HostEvmEnv +where + D: Send + 'static, +{ + /// Runs the provided closure that requires mutable access to the database on a thread where + /// blocking is acceptable. + /// + /// It panics if the closure panics. + /// This function is necessary because mutable references to the database cannot be passed + /// directly to `tokio::task::spawn_blocking`. Instead, the database is temporarily taken out of + /// the `HostEvmEnv`, moved into the blocking task, and then restored after the task completes. + #[allow(dead_code)] + pub(crate) async fn spawn_with_db(&mut self, f: F) -> R + where + F: FnOnce(&mut ProofDb) -> R + Send + 'static, + R: Send + 'static, + { + // as mutable references are not possible, the DB must be moved in and out of the task + let mut db = self.db.take().unwrap(); + + let (result, db) = tokio::task::spawn_blocking(move || (f(&mut db), db)) + .await + .expect("DB execution panicked"); + + // restore the DB, so that we never return an env without a DB + self.db = Some(db); + + result + } +} + impl HostEvmEnv { /// Sets the chain ID and specification ID from the given chain spec. /// diff --git a/crates/steel/src/lib.rs b/crates/steel/src/lib.rs index 196e553f..579d6338 100644 --- a/crates/steel/src/lib.rs +++ b/crates/steel/src/lib.rs @@ -29,6 +29,7 @@ use alloy_sol_types::SolValue; use config::ChainSpec; use revm::primitives::{BlockEnv, CfgEnvWithHandlerCfg, SpecId}; +pub mod account; pub mod beacon; mod block; pub mod config; @@ -47,6 +48,7 @@ mod state; #[cfg(feature = "unstable-verifier")] mod verifier; +pub use account::Account; pub use beacon::BeaconInput; pub use block::BlockInput; pub use contract::{CallBuilder, Contract}; diff --git a/crates/steel/src/verifier.rs b/crates/steel/src/verifier.rs index 42bfb907..5cd59f82 100644 --- a/crates/steel/src/verifier.rs +++ b/crates/steel/src/verifier.rs @@ -55,7 +55,6 @@ impl<'a, H: EvmBlockHeader> SteelVerifier<&'a GuestEvmEnv> { #[cfg(feature = "host")] mod host { use super::*; - use crate::host::db::ProofDb; use crate::{history::beacon_roots, host::HostEvmEnv}; use anyhow::Context; use revm::Database; @@ -109,36 +108,6 @@ mod host { } } } - - impl HostEvmEnv - where - D: Database + Send + 'static, - { - /// Runs the provided closure that requires mutable access to the database on a thread where - /// blocking is acceptable. - /// - /// It panics if the closure panics. - /// This function is necessary because mutable references to the database cannot be passed - /// directly to `tokio::task::spawn_blocking`. Instead, the database is temporarily taken out of - /// the `HostEvmEnv`, moved into the blocking task, and then restored after the task completes. - async fn spawn_with_db(&mut self, f: F) -> R - where - F: FnOnce(&mut ProofDb) -> R + Send + 'static, - R: Send + 'static, - { - // as mutable references are not possible, the DB must be moved in and out of the task - let mut db = self.db.take().unwrap(); - - let (result, db) = tokio::task::spawn_blocking(|| (f(&mut db), db)) - .await - .expect("DB execution panicked"); - - // restore the DB, so that we never return an env without a DB - self.db = Some(db); - - result - } - } } fn validate_block_number(header: &impl EvmBlockHeader, block_number: U256) -> anyhow::Result { diff --git a/crates/steel/tests/steel.rs b/crates/steel/tests/steel.rs index 305b76a2..06cc6fdb 100644 --- a/crates/steel/tests/steel.rs +++ b/crates/steel/tests/steel.rs @@ -22,10 +22,10 @@ use alloy::{ transports::BoxTransport, uint, }; -use alloy_primitives::{address, b256, bytes, hex, Address, Bytes, U256}; +use alloy_primitives::{address, b256, bytes, hex, keccak256, Address, Bytes, U256}; use alloy_sol_types::SolCall; use common::{CallOptions, ANVIL_CHAIN_SPEC}; -use risc0_steel::{ethereum::EthEvmEnv, Contract}; +use risc0_steel::{ethereum::EthEvmEnv, Account, Contract}; use sha2::{Digest, Sha256}; use test_log::test; @@ -116,7 +116,7 @@ alloy::sol!( ); /// Returns an Anvil provider with the deployed [SteelTest] contract. -async fn test_provider() -> impl Provider { +async fn test_provider() -> impl Provider + Clone { let provider = ProviderBuilder::new() .with_recommended_fillers() .on_anvil_with_wallet_and_config(|anvil| anvil.args(["--hardfork", "cancun"])); @@ -128,6 +128,40 @@ async fn test_provider() -> impl Provider { provider } +#[test(tokio::test)] +async fn account_info() { + let provider = test_provider().await; + let mut env = EthEvmEnv::builder() + .provider(provider.clone()) + .build() + .await + .unwrap() + .with_chain_spec(&ANVIL_CHAIN_SPEC); + let address = STEEL_TEST_CONTRACT; + let preflight_info = { + let account = Account::preflight(address, &mut env); + account.bytecode(true).info().await.unwrap() + }; + + let input = env.into_input().await.unwrap(); + let env = input.into_env().with_chain_spec(&ANVIL_CHAIN_SPEC); + + let info = { + let account = Account::new(address, &env); + account.bytecode(true).info() + }; + assert_eq!(info, preflight_info, "mismatch in preflight and execution"); + + assert_eq!(info.balance, provider.get_balance(address).await.unwrap()); + assert_eq!( + info.nonce, + provider.get_transaction_count(address).await.unwrap() + ); + let code = info.code.unwrap().bytes(); + assert_eq!(code, provider.get_code_at(address).await.unwrap()); + assert_eq!(info.code_hash, keccak256(code)); +} + #[test(tokio::test)] async fn ec_recover() { let result = common::eth_call(