Add support for state override

lib/ain-evm/src/backend.rs

@@ -49,8 +49,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>,
@@ -130,6 +130,7 @@ impl EVMBackend {
         trie_store: Arc<TrieDBStore>,
         storage: Arc<Storage>,
         vicinity: Vicinity,
+        overlay: Option<Overlay>,
     ) -> Result<Self> {
         let state = trie_store
             .trie_db
@@ -141,7 +142,7 @@ impl EVMBackend {
             trie_store,
             storage,
             vicinity,
-            overlay: Overlay::new(),
+            overlay: overlay.unwrap_or_default(),
         })
     }
 

lib/ain-evm/src/core.rs

@@ -22,7 +22,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},
@@ -215,7 +215,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,
@@ -280,6 +280,7 @@ impl EVMCoreService {
             Arc::clone(&self.trie_store),
             Arc::clone(&self.storage),
             vicinity,
+            overlay,
         )
         .map_err(|e| format_err!("------ Could not restore backend {}", e))?;
 
@@ -763,6 +764,7 @@ impl EVMCoreService {
             Arc::clone(&self.trie_store),
             Arc::clone(&self.storage),
             Vicinity::default(),
+            None, // TODO state override
         )?;
         Ok(backend.get_account(&address))
     }
@@ -843,6 +845,7 @@ impl EVMCoreService {
                 ),
                 ..Vicinity::default()
             },
+            None,
         )
     }
 
@@ -860,6 +863,7 @@ impl EVMCoreService {
                 ),
                 ..Vicinity::default()
             },
+            None,
         )
     }
 

lib/ain-evm/src/evm.rs

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

lib/ain-evm/src/lib.rs

@@ -1,4 +1,4 @@
-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 serde::Deserialize;
 
 /// Call request
@@ -34,3 +36,53 @@ pub struct CallRequest {
     #[serde(rename = "type")]
     pub transaction_type: Option<U256>,
 }
+
+// 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(overide: BTreeMap<H160, CallStateOverride>) -> Overlay {
+    let mut overlay = Overlay::default();
+
+    for (address, state_override) in overide {
+        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

@@ -146,18 +146,21 @@ impl MetachainDebugRPCServer for MetachainDebugRPCModule {
         let TxResponse { used_gas, .. } = self
             .handler
             .core
-            .call(EthCallArgs {
-                caller: from,
-                to,
-                value: value.unwrap_or_default(),
-                data: &data.map(|d| d.0).unwrap_or_default(),
-                gas_limit,
-                gas_price,
-                max_fee_per_gas,
-                access_list: access_list.unwrap_or_default(),
-                block_number,
-                transaction_type,
-            })
+            .call(
+                EthCallArgs {
+                    caller: from,
+                    to,
+                    value: value.unwrap_or_default(),
+                    data: &data.map(|d| d.0).unwrap_or_default(),
+                    gas_limit,
+                    gas_price,
+                    max_fee_per_gas,
+                    access_list: access_list.unwrap_or_default(),
+                    block_number,
+                    transaction_type,
+                },
+                None,
+            )
             .map_err(|e| Error::Custom(format!("Error calling EVM : {e:?}")))?;
 
         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::{
@@ -23,7 +23,7 @@ use log::{debug, trace};
 use super::to_jsonrpsee_custom_error;
 use crate::{
     block::{BlockNumber, RpcBlock, RpcFeeHistory},
-    call_request::CallRequest,
+    call_request::{override_to_overlay, CallRequest, CallStateOverride},
     codegen::types::EthTransactionInfo,
     filters::{GetFilterChangesResult, NewFilterRequest},
     receipt::ReceiptResult,
@@ -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, input: 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.
@@ -211,6 +216,7 @@ pub trait MetachainRPC {
         &self,
         input: CallRequest,
         block_number: Option<BlockNumber>,
+        state_overrides: Option<BTreeMap<H160, CallStateOverride>>,
     ) -> RpcResult<U256>;
 
     /// Returns current gas_price.
@@ -314,7 +320,12 @@ impl MetachainRPCModule {
 }
 
 impl MetachainRPCServer for MetachainRPCModule {
-    fn call(&self, input: CallRequest, block_number: Option<BlockNumber>) -> RpcResult<Bytes> {
+    fn call(
+        &self,
+        input: CallRequest,
+        block_number: Option<BlockNumber>,
+        state_overrides: Option<BTreeMap<H160, CallStateOverride>>,
+    ) -> RpcResult<Bytes> {
         debug!(target:"rpc",  "[RPC] Call input {:#?}", input);
         let CallRequest {
             from,
@@ -343,24 +354,27 @@ impl MetachainRPCServer for MetachainRPCModule {
         let TxResponse { data, .. } = self
             .handler
             .core
-            .call(EthCallArgs {
-                caller: from,
-                to,
-                value: value.unwrap_or_default(),
-                // https://github.com/ethereum/go-ethereum/blob/281e8cd5abaac86ed3f37f98250ff147b3c9fe62/internal/ethapi/transaction_args.go#L67
-                // We accept "data" and "input" for backwards-compatibility reasons.
-                //  "input" is the newer name and should be preferred by clients.
-                // 	Issue detail: https://github.com/ethereum/go-ethereum/issues/15628
-                data: &input
-                    .map(|d| d.0)
-                    .unwrap_or(data.map(|d| d.0).unwrap_or_default()),
-                gas_limit,
-                gas_price,
-                max_fee_per_gas,
-                access_list: access_list.unwrap_or_default(),
-                block_number: self.block_number_to_u256(block_number)?,
-                transaction_type,
-            })
+            .call(
+                EthCallArgs {
+                    caller: from,
+                    to,
+                    value: value.unwrap_or_default(),
+                    // https://github.com/ethereum/go-ethereum/blob/281e8cd5abaac86ed3f37f98250ff147b3c9fe62/internal/ethapi/transaction_args.go#L67
+                    // We accept "data" and "input" for backwards-compatibility reasons.
+                    //  "input" is the newer name and should be preferred by clients.
+                    // 	Issue detail: https://github.com/ethereum/go-ethereum/issues/15628
+                    data: &input
+                        .map(|d| d.0)
+                        .unwrap_or(data.map(|d| d.0).unwrap_or_default()),
+                    gas_limit,
+                    gas_price,
+                    max_fee_per_gas,
+                    access_list: access_list.unwrap_or_default(),
+                    block_number: self.block_number_to_u256(block_number)?,
+                    transaction_type,
+                },
+                state_overrides.map(override_to_overlay),
+            )
             .map_err(|e| {
                 debug!("Error calling EVM : {e:?}");
                 Error::Custom(format!("Error calling EVM : {e:?}"))
@@ -798,6 +812,7 @@ impl MetachainRPCServer for MetachainRPCModule {
         &self,
         input: CallRequest,
         block_number: Option<BlockNumber>,
+        state_overrides: Option<BTreeMap<H160, CallStateOverride>>,
     ) -> RpcResult<U256> {
         let CallRequest {
             from,
@@ -827,18 +842,21 @@ impl MetachainRPCServer for MetachainRPCModule {
         let TxResponse { used_gas, .. } = self
             .handler
             .core
-            .call(EthCallArgs {
-                caller: from,
-                to,
-                value: value.unwrap_or_default(),
-                data: &data.map(|d| d.0).unwrap_or_default(),
-                gas_limit,
-                gas_price,
-                max_fee_per_gas,
-                access_list: access_list.unwrap_or_default(),
-                block_number,
-                transaction_type,
-            })
+            .call(
+                EthCallArgs {
+                    caller: from,
+                    to,
+                    value: value.unwrap_or_default(),
+                    data: &data.map(|d| d.0).unwrap_or_default(),
+                    gas_limit,
+                    gas_price,
+                    max_fee_per_gas,
+                    access_list: access_list.unwrap_or_default(),
+                    block_number,
+                    transaction_type,
+                },
+                state_overrides.map(override_to_overlay),
+            )
             .map_err(|e| Error::Custom(format!("Error calling EVM : {e:?}")))?;
 
         debug!(target:"rpc",  "estimateGas: {:#?} at block {:#x}", used_gas, block_number);