Evm state override (#2623)

* Add support for state override

* Add tests for state override call

lib/ain-evm/src/backend.rs

@@ -64,8 +64,8 @@ struct OverlayData {
     storage: HashMap<H256, H256>,
 }
 
-#[derive(Debug, Clone)]
-struct Overlay {
+#[derive(Debug, Clone, Default)]
+pub struct Overlay {
     state: HashMap<H160, OverlayData>,
     changeset: Vec<HashMap<H160, OverlayData>>,
     deletes: HashSet<H160>,
@@ -145,6 +145,7 @@ impl EVMBackend {
         trie_store: Arc<TrieDBStore>,
         storage: Arc<Storage>,
         vicinity: Vicinity,
+        overlay: Option<Overlay>,
     ) -> Result<Self> {
         let state = trie_store
             .trie_db
@@ -156,7 +157,7 @@ impl EVMBackend {
             trie_store,
             storage,
             vicinity,
-            overlay: Overlay::new(),
+            overlay: overlay.unwrap_or_default(),
         })
     }
 

lib/ain-evm/src/core.rs

@@ -28,7 +28,7 @@ use parking_lot::Mutex;
 use vsdb_core::vsdb_set_base_dir;
 
 use crate::{
-    backend::{BackendError, EVMBackend, Vicinity},
+    backend::{BackendError, EVMBackend, Overlay, Vicinity},
     block::INITIAL_BASE_FEE,
     blocktemplate::BlockTemplate,
     executor::{AinExecutor, ExecutorContext, TxResponse},
@@ -239,7 +239,7 @@ impl EVMCoreService {
         self.trie_store.flush()
     }
 
-    pub fn call(&self, arguments: EthCallArgs) -> Result<TxResponse> {
+    pub fn call(&self, arguments: EthCallArgs, overlay: Option<Overlay>) -> Result<TxResponse> {
         let EthCallArgs {
             caller,
             to,
@@ -276,6 +276,7 @@ impl EVMCoreService {
             Arc::clone(&self.trie_store),
             Arc::clone(&self.storage),
             vicinity,
+            overlay,
         )
         .map_err(|e| format_err!("------ Could not restore backend {}", e))?;
 
@@ -766,6 +767,7 @@ impl EVMCoreService {
             Arc::clone(&self.trie_store),
             Arc::clone(&self.storage),
             vicinity,
+            None,
         )
         .map_err(|e| format_err!("Could not restore backend {}", e))?;
         backend.update_vicinity_from_tx(tx)?;
@@ -832,6 +834,7 @@ impl EVMCoreService {
             Arc::clone(&self.trie_store),
             Arc::clone(&self.storage),
             Vicinity::default(),
+            None,
         )?;
         Ok(backend.get_account(&address))
     }
@@ -901,6 +904,7 @@ impl EVMCoreService {
             Arc::clone(&self.trie_store),
             Arc::clone(&self.storage),
             vicinity,
+            None,
         )
     }
 
@@ -918,6 +922,7 @@ impl EVMCoreService {
                 ),
                 ..Vicinity::default()
             },
+            None,
         )
     }
 

lib/ain-evm/src/evm.rs

@@ -526,6 +526,7 @@ impl EVMServices {
             Arc::clone(&self.core.trie_store),
             Arc::clone(&self.storage),
             vicinity.clone(),
+            None,
         )?;
 
         let template = BlockTemplate::new(vicinity, ctx, backend);

lib/ain-evm/src/lib.rs

@@ -1,6 +1,6 @@
 //! Defichain EVM consensus, runtime and storage implementation
 
-mod backend;
+pub mod backend;
 pub mod block;
 pub mod blocktemplate;
 pub mod bytes;

lib/ain-evm/src/trie.rs

@@ -66,6 +66,7 @@ impl TrieDBStore {
             Arc::clone(trie_store),
             Arc::clone(storage),
             Vicinity::default(),
+            None,
         )
         .expect("Could not restore backend");
 

lib/ain-grpc/src/call_request.rs

@@ -1,6 +1,8 @@
-use ain_evm::bytes::Bytes;
-use ethereum::AccessListItem;
-use ethereum_types::{H160, U256};
+use std::collections::{BTreeMap, HashMap};
+
+use ain_evm::{backend::Overlay, bytes::Bytes};
+use ethereum::{AccessListItem, Account};
+use ethereum_types::{H160, H256, U256};
 use jsonrpsee::core::Error;
 use serde::Deserialize;
 
@@ -127,3 +129,53 @@ impl CallRequest {
         }
     }
 }
+
+// State override
+#[derive(Clone, Debug, Default, Eq, PartialEq, Deserialize, Serialize)]
+#[serde(deny_unknown_fields)]
+#[serde(rename_all = "camelCase")]
+pub struct CallStateOverride {
+    /// Fake balance to set for the account before executing the call.
+    pub balance: Option<U256>,
+    /// Fake nonce to set for the account before executing the call.
+    pub nonce: Option<U256>,
+    /// Fake EVM bytecode to inject into the account before executing the call.
+    pub code: Option<Bytes>,
+    /// Fake key-value mapping to override all slots in the account storage before
+    /// executing the call.
+    pub state: Option<BTreeMap<H256, H256>>,
+    /// Fake key-value mapping to override individual slots in the account storage before
+    /// executing the call.
+    pub state_diff: Option<BTreeMap<H256, H256>>,
+}
+
+pub fn override_to_overlay(r#override: BTreeMap<H160, CallStateOverride>) -> Overlay {
+    let mut overlay = Overlay::default();
+
+    for (address, state_override) in r#override {
+        let code = state_override.code.map(|b| b.into_vec());
+        let mut storage = state_override
+            .state
+            .unwrap_or_default()
+            .into_iter()
+            .collect::<HashMap<_, _>>();
+
+        let account = Account {
+            balance: state_override.balance.unwrap_or_default(),
+            nonce: state_override.nonce.unwrap_or_default(),
+            storage_root: H256::zero(),
+            code_hash: H256::zero(),
+        };
+
+        let reset_storage = storage.is_empty();
+        if let Some(diff) = state_override.state_diff {
+            for (k, v) in diff {
+                storage.insert(k, v);
+            }
+        }
+
+        overlay.apply(address, account, code, storage, reset_storage);
+    }
+
+    overlay
+}

lib/ain-grpc/src/rpc/debug.rs

@@ -130,16 +130,19 @@ impl MetachainDebugRPCServer for MetachainDebugRPCModule {
         let TxResponse { used_gas, .. } = self
             .handler
             .core
-            .call(EthCallArgs {
-                caller,
-                to: call.to,
-                value: call.value.unwrap_or_default(),
-                data,
-                gas_limit,
-                gas_price,
-                access_list: call.access_list.unwrap_or_default(),
-                block_number,
-            })
+            .call(
+                EthCallArgs {
+                    caller,
+                    to: call.to,
+                    value: call.value.unwrap_or_default(),
+                    data,
+                    gas_limit,
+                    gas_price,
+                    access_list: call.access_list.unwrap_or_default(),
+                    block_number,
+                },
+                None,
+            )
             .map_err(RPCError::EvmError)?;
 
         let used_gas = U256::from(used_gas);

lib/ain-grpc/src/rpc/eth.rs

@@ -1,4 +1,4 @@
-use std::{convert::Into, str::FromStr, sync::Arc};
+use std::{collections::BTreeMap, convert::Into, str::FromStr, sync::Arc};
 
 use ain_cpp_imports::get_eth_priv_key;
 use ain_evm::{
@@ -22,7 +22,7 @@ use log::{debug, trace};
 
 use crate::{
     block::{BlockNumber, RpcBlock, RpcFeeHistory},
-    call_request::CallRequest,
+    call_request::{override_to_overlay, CallRequest, CallStateOverride},
     codegen::types::EthTransactionInfo,
     errors::{to_custom_err, RPCError},
     filters::{GetFilterChangesResult, NewFilterRequest},
@@ -42,7 +42,12 @@ pub trait MetachainRPC {
     /// Makes a call to the Ethereum node without creating a transaction on the blockchain.
     /// Returns the output data as a hexadecimal string.
     #[method(name = "call")]
-    fn call(&self, call: CallRequest, block_number: Option<BlockNumber>) -> RpcResult<Bytes>;
+    fn call(
+        &self,
+        input: CallRequest,
+        block_number: Option<BlockNumber>,
+        state_overrides: Option<BTreeMap<H160, CallStateOverride>>,
+    ) -> RpcResult<Bytes>;
 
     /// Retrieves the list of accounts managed by the node.
     /// Returns a vector of Ethereum addresses as hexadecimal strings.
@@ -207,8 +212,12 @@ pub trait MetachainRPC {
 
     /// Estimate gas needed for execution of given contract.
     #[method(name = "estimateGas")]
-    fn estimate_gas(&self, call: CallRequest, block_number: Option<BlockNumber>)
-        -> RpcResult<U256>;
+    fn estimate_gas(
+        &self,
+        input: CallRequest,
+        block_number: Option<BlockNumber>,
+        state_overrides: Option<BTreeMap<H160, CallStateOverride>>,
+    ) -> RpcResult<U256>;
 
     /// Returns current gas_price.
     #[method(name = "gasPrice")]
@@ -309,7 +318,12 @@ impl MetachainRPCModule {
 }
 
 impl MetachainRPCServer for MetachainRPCModule {
-    fn call(&self, call: CallRequest, block_number: Option<BlockNumber>) -> RpcResult<Bytes> {
+    fn call(
+        &self,
+        call: CallRequest,
+        block_number: Option<BlockNumber>,
+        state_overrides: Option<BTreeMap<H160, CallStateOverride>>,
+    ) -> RpcResult<Bytes> {
         debug!(target:"rpc",  "Call, input {:#?}", call);
 
         let caller = call.from.unwrap_or_default();
@@ -330,16 +344,19 @@ impl MetachainRPCServer for MetachainRPCModule {
         } = self
             .handler
             .core
-            .call(EthCallArgs {
-                caller,
-                to: call.to,
-                value: call.value.unwrap_or_default(),
-                data,
-                gas_limit,
-                gas_price,
-                access_list: call.access_list.unwrap_or_default(),
-                block_number: block.header.number,
-            })
+            .call(
+                EthCallArgs {
+                    caller,
+                    to: call.to,
+                    value: call.value.unwrap_or_default(),
+                    data,
+                    gas_limit,
+                    gas_price,
+                    access_list: call.access_list.unwrap_or_default(),
+                    block_number: block.header.number,
+                },
+                state_overrides.map(override_to_overlay),
+            )
             .map_err(RPCError::EvmError)?;
 
         match exit_reason {
@@ -764,6 +781,7 @@ impl MetachainRPCServer for MetachainRPCModule {
         &self,
         call: CallRequest,
         block_number: Option<BlockNumber>,
+        state_overrides: Option<BTreeMap<H160, CallStateOverride>>,
     ) -> RpcResult<U256> {
         debug!(target:"rpc",  "Estimate gas, input {:#?}", call);
 
@@ -776,6 +794,7 @@ impl MetachainRPCServer for MetachainRPCModule {
         let call_gas = u64::try_from(call.gas.unwrap_or(U256::from(block_gas_limit)))
             .map_err(to_custom_err)?;
 
+        let overlay = state_overrides.map(override_to_overlay);
         // Determine the lowest and highest possible gas limits to binary search in between
         let mut lo = Self::CONFIG.gas_transaction_call - 1;
         let mut hi = call_gas;
@@ -828,16 +847,19 @@ impl MetachainRPCServer for MetachainRPCModule {
             let tx_response = self
                 .handler
                 .core
-                .call(EthCallArgs {
-                    caller,
-                    to: call.to,
-                    value: call.value.unwrap_or_default(),
-                    data,
-                    gas_limit,
-                    gas_price: fee_cap,
-                    access_list: call.access_list.clone().unwrap_or_default(),
-                    block_number: block.header.number,
-                })
+                .call(
+                    EthCallArgs {
+                        caller,
+                        to: call.to,
+                        value: call.value.unwrap_or_default(),
+                        data,
+                        gas_limit,
+                        gas_price: fee_cap,
+                        access_list: call.access_list.clone().unwrap_or_default(),
+                        block_number: block.header.number,
+                    },
+                    overlay.clone(),
+                )
                 .map_err(RPCError::EvmError)?;
 
             match tx_response.exit_reason {