diff --git a/core/src/database/mod.rs b/core/src/database/mod.rs index 99890b41..5bba5204 100644 --- a/core/src/database/mod.rs +++ b/core/src/database/mod.rs @@ -21,25 +21,41 @@ use super::*; pub mod fork; pub mod inspector; -/// A [`ArbiterDB`] is a wrapper around a [`CacheDB`] that is used to provide -/// access to the [`environment::Environment`]'s database to multiple +/// A [`ArbiterDB`] is contains both a [`CacheDB`] that is used to provide +/// state for the [`environment::Environment`]'s as well as for multiple /// [`coprocessor::Coprocessor`]s. +/// The `logs` field is a [`BTreeMap`] to store [`ethers::types::Log`]s that can be queried from at any point. #[derive(Debug, Serialize, Deserialize)] -pub struct ArbiterDB(pub Arc>>); +pub struct ArbiterDB { + /// The state of the `ArbiterDB`. This is a `CacheDB` that is used to provide + /// a db for the `Environment` to use. + pub state: Arc>>, + + /// The logs of the `ArbiterDB`. This is a `BTreeMap` that is used to store + /// logs that can be queried from at any point. + pub logs: Arc>>>, +} +// Implement `Clone` by hand so we utilize the `Arc`'s `Clone` implementation. impl Clone for ArbiterDB { fn clone(&self) -> Self { - Self(self.0.clone()) + Self { + state: self.state.clone(), + logs: self.logs.clone(), + } } } impl ArbiterDB { /// Create a new `ArbiterDB`. pub fn new() -> Self { - Self(Arc::new(RwLock::new(CacheDB::new(EmptyDB::new())))) + Self { + state: Arc::new(RwLock::new(CacheDB::new(EmptyDB::new()))), + logs: Arc::new(RwLock::new(BTreeMap::new())), + } } - /// Write the `ArbiterDB` to a file at the given path. + /// Write the `ArbiterDB` to a file at the given path.`` pub fn write_to_file(&self, path: &str) -> io::Result<()> { // Serialize the ArbiterDB let serialized = serde_json::to_string(self)?; @@ -55,9 +71,18 @@ impl ArbiterDB { let mut file = fs::File::open(path)?; let mut contents = String::new(); file.read_to_string(&mut contents)?; + // Deserialize the content into ArbiterDB - let cache_db = serde_json::from_str(&contents)?; - Ok(Self(Arc::new(RwLock::new(cache_db)))) + #[derive(Deserialize)] + struct TempDB { + state: Option>, + logs: Option>>, + } + let temp_db: TempDB = serde_json::from_str(&contents)?; + Ok(Self { + state: Arc::new(RwLock::new(temp_db.state.unwrap_or_default())), + logs: Arc::new(RwLock::new(temp_db.logs.unwrap_or_default())), + }) } } @@ -83,11 +108,11 @@ impl Database for ArbiterDB { &mut self, address: revm::primitives::Address, ) -> Result, Self::Error> { - self.0.write().unwrap().basic(address) + self.state.write().unwrap().basic(address) } fn code_by_hash(&mut self, code_hash: B256) -> Result { - self.0.write().unwrap().code_by_hash(code_hash) + self.state.write().unwrap().code_by_hash(code_hash) } fn storage( @@ -95,11 +120,11 @@ impl Database for ArbiterDB { address: revm::primitives::Address, index: U256, ) -> Result { - self.0.write().unwrap().storage(address, index) + self.state.write().unwrap().storage(address, index) } fn block_hash(&mut self, number: U256) -> Result { - self.0.write().unwrap().block_hash(number) + self.state.write().unwrap().block_hash(number) } } @@ -110,11 +135,11 @@ impl DatabaseRef for ArbiterDB { &self, address: revm::primitives::Address, ) -> Result, Self::Error> { - self.0.read().unwrap().basic_ref(address) + self.state.read().unwrap().basic_ref(address) } fn code_by_hash_ref(&self, code_hash: B256) -> Result { - self.0.read().unwrap().code_by_hash_ref(code_hash) + self.state.read().unwrap().code_by_hash_ref(code_hash) } fn storage_ref( @@ -122,11 +147,11 @@ impl DatabaseRef for ArbiterDB { address: revm::primitives::Address, index: U256, ) -> Result { - self.0.read().unwrap().storage_ref(address, index) + self.state.read().unwrap().storage_ref(address, index) } fn block_hash_ref(&self, number: U256) -> Result { - self.0.read().unwrap().block_hash_ref(number) + self.state.read().unwrap().block_hash_ref(number) } } @@ -135,7 +160,7 @@ impl DatabaseCommit for ArbiterDB { &mut self, changes: hashbrown::HashMap, ) { - self.0.write().unwrap().commit(changes) + self.state.write().unwrap().commit(changes) } } diff --git a/core/src/environment/mod.rs b/core/src/environment/mod.rs index a171e706..9289f95a 100644 --- a/core/src/environment/mod.rs +++ b/core/src/environment/mod.rs @@ -77,8 +77,6 @@ pub struct Environment { /// calls and transactions. pub(crate) db: ArbiterDB, - pub(crate) log_storage: Arc>>>, - inspector: Option, /// This gives a means of letting the "outside world" connect to the @@ -147,10 +145,22 @@ impl EnvironmentBuilder { self } - /// Sets the database for the [`Environment`]. This can come from a - /// [`fork::Fork`]. - pub fn with_db(mut self, db: impl Into>) -> Self { - self.db = ArbiterDB(Arc::new(RwLock::new(db.into()))); + /// Sets the state for the [`Environment`]. This can come from a saved state of a simulation + /// or a [`database::fork::Fork`]. + pub fn with_state(mut self, state: impl Into>) -> Self { + self.db.state = Arc::new(RwLock::new(state.into())); + self + } + + /// Sets the logs for the [`Environment`]. This can come from a saved state of a simulation and can be useful for doing analysis. + pub fn with_logs(mut self, logs: impl Into>>) -> Self { + self.db.logs = Arc::new(RwLock::new(logs.into())); + self + } + + /// Sets the entire database for the [`Environment`] including both the state and logs. This can come from the saved state of a simulation and can be useful for doing analysis. + pub fn with_arbiter_db(mut self, db: ArbiterDB) -> Self { + self.db = db; self } @@ -175,7 +185,7 @@ impl Environment { pub fn builder() -> EnvironmentBuilder { EnvironmentBuilder { parameters: EnvironmentParameters::default(), - db: ArbiterDB(Arc::new(RwLock::new(CacheDB::new(EmptyDB::new())))), + db: ArbiterDB::default(), } } @@ -202,7 +212,6 @@ impl Environment { inspector, parameters, db, - log_storage: Arc::new(RwLock::new(BTreeMap::new())), handle: None, } } @@ -215,7 +224,6 @@ impl Environment { // Bring in the EVM db and log storage by cloning the interior Arc (lightweight). let db = self.db.clone(); - let log_storage = self.log_storage.clone(); // Bring in the EVM ENV let mut env = Env::default(); @@ -232,7 +240,7 @@ impl Environment { let handle = thread::spawn(move || { // Create a new EVM builder let mut evm = Evm::builder() - .with_db(db) + .with_db(db.clone()) .with_env(Box::new(env)) .with_external_context(inspector) .append_handler_register(inspector_handle_register) @@ -260,14 +268,7 @@ impl Environment { account_state: AccountState::None, storage: HashMap::new(), }; - let db = &mut evm.context.evm.db; - match db - .0 - .write() - .unwrap() - .accounts - .insert(recast_address, account) - { + match db.state.write()?.accounts.insert(recast_address, account) { None => outcome_sender.send(Ok(Outcome::AddAccountCompleted))?, Some(_) => { outcome_sender.send(Err(ArbiterCoreError::AccountCreationError))?; @@ -307,13 +308,11 @@ impl Environment { key, block: _, } => { - let db = &mut evm.context.evm.db; - let recast_address = Address::from(account.as_fixed_bytes()); let recast_key = B256::from(key.as_fixed_bytes()).into(); // Get the account storage value at the key in the db. - match db.0.write().unwrap().accounts.get_mut(&recast_address) { + match db.state.write().unwrap().accounts.get_mut(&recast_address) { Some(account) => { // Returns zero if the account is missing. let value: U256 = match account.storage.get::(&recast_key) @@ -336,16 +335,13 @@ impl Environment { key, value, } => { - // Get the underlying database - let db = &mut evm.context.evm.db; - let recast_address = Address::from(account.as_fixed_bytes()); let recast_key = B256::from(key.as_fixed_bytes()); let recast_value = B256::from(value.as_fixed_bytes()); // Mutate the db by inserting the new key-value pair into the account's // storage and send the successful CheatcodeCompleted outcome. - match db.0.write().unwrap().accounts.get_mut(&recast_address) { + match db.state.write().unwrap().accounts.get_mut(&recast_address) { Some(account) => { account .storage @@ -362,9 +358,8 @@ impl Environment { }; } Cheatcodes::Deal { address, amount } => { - let db = &mut evm.context.evm.db; let recast_address = Address::from(address.as_fixed_bytes()); - match db.0.write().unwrap().accounts.get_mut(&recast_address) { + match db.state.write().unwrap().accounts.get_mut(&recast_address) { Some(account) => { account.info.balance += U256::from_limbs(amount.0); outcome_sender.send(Ok(Outcome::CheatcodeReturn( @@ -378,10 +373,8 @@ impl Environment { }; } Cheatcodes::Access { address } => { - let db = &mut evm.context.evm.db; let recast_address = Address::from(address.as_fixed_bytes()); - - match db.0.write().unwrap().accounts.get(&recast_address) { + match db.state.write().unwrap().accounts.get(&recast_address) { Some(account) => { let account_state = match account.account_state { AccountState::None => AccountStateSerializable::None, @@ -472,8 +465,7 @@ impl Environment { cumulative_gas_per_block, }; - // TODO: Don't unwrap this. - let mut logs = log_storage.write().unwrap(); + let mut logs = db.logs.write()?; match logs.get_mut(&evm.block().number) { Some(log_vec) => { log_vec.extend(revm_logs_to_ethers_logs( @@ -525,10 +517,8 @@ impl Environment { Ok(Outcome::QueryReturn(evm.tx().gas_price.to_string())) } EnvironmentData::Balance(address) => { - // This unwrap should never fail. - let db = &mut evm.context.evm.db; match db - .0 + .state .read() .unwrap() .accounts @@ -541,9 +531,8 @@ impl Environment { } } EnvironmentData::TransactionCount(address) => { - let db = &mut evm.context.evm.db; match db - .0 + .state .read() .unwrap() .accounts @@ -556,7 +545,7 @@ impl Environment { } } EnvironmentData::Logs { filter } => { - let logs = log_storage.read().unwrap(); + let logs = db.logs.read().unwrap(); let from_block = U256::from( filter .block_option @@ -623,7 +612,6 @@ impl Environment { warn!("Stop signal was not sent to any listeners. Are there any listeners?") } } - let (db, _) = evm.into_db_and_env_with_handler_cfg(); outcome_sender.send(Ok(Outcome::StopCompleted(db)))?; break; } diff --git a/core/src/errors.rs b/core/src/errors.rs index 3a67728a..d9e33299 100644 --- a/core/src/errors.rs +++ b/core/src/errors.rs @@ -1,6 +1,8 @@ //! Errors that can occur when managing or interfacing with Arbiter's sandboxed //! Ethereum environment. +use std::sync::{PoisonError, RwLockWriteGuard}; + // use crossbeam_channel::SendError; use crossbeam_channel::{RecvError, SendError}; use ethers::{ @@ -101,6 +103,9 @@ pub enum ArbiterCoreError { /// Failed to reply to instruction. #[error("{0}")] ReplyError(String), + + #[error("{0}")] + RwLockError(String), } impl From>> for ArbiterCoreError { @@ -109,6 +114,12 @@ impl From>> for ArbiterCoreError { } } +impl From>> for ArbiterCoreError { + fn from(e: PoisonError>) -> Self { + ArbiterCoreError::RwLockError(e.to_string()) + } +} + impl MiddlewareError for ArbiterCoreError { type Inner = ProviderError;