feat(contract): support state overrides for Multicall (#2478)

* feat(contract): support state overrides for Multicall

* feat(contract): add tests for state overrides of Multicall

ethers-contract/src/multicall/mod.rs

@@ -2,10 +2,11 @@ use crate::call::{ContractCall, ContractError};
 use ethers_core::{
     abi::{Detokenize, Function, Token, Tokenizable},
     types::{
-        transaction::eip2718::TypedTransaction, Address, BlockNumber, Bytes, NameOrAddress, U256,
+        transaction::eip2718::TypedTransaction, Address, BlockId, BlockNumber, Bytes,
+        NameOrAddress, U256,
     },
 };
-use ethers_providers::{Middleware, PendingTransaction};
+use ethers_providers::{spoof::State, Middleware, PendingTransaction, RawCall};
 use std::{convert::TryFrom, fmt, result::Result as StdResult, sync::Arc};
 
 /// The Multicall contract bindings. Auto-generated with `abigen`.
@@ -205,7 +206,10 @@ pub struct Multicall<M> {
     pub legacy: bool,
 
     /// The `block` field of the Multicall aggregate call.
-    pub block: Option<BlockNumber>,
+    pub block: Option<BlockId>,
+
+    /// The state overrides of the Multicall aggregate
+    pub state: Option<State>,
 
     /// The internal call vector.
     calls: Vec<Call>,
@@ -220,6 +224,7 @@ impl<M> Clone for Multicall<M> {
             legacy: self.legacy,
             block: self.block,
             calls: self.calls.clone(),
+            state: self.state.clone(),
         }
     }
 }
@@ -231,6 +236,7 @@ impl<M> fmt::Debug for Multicall<M> {
             .field("version", &self.version)
             .field("legacy", &self.legacy)
             .field("block", &self.block)
+            .field("state", &self.state)
             .field("calls", &self.calls)
             .finish()
     }
@@ -278,6 +284,7 @@ impl<M: Middleware> Multicall<M> {
             version: MulticallVersion::Multicall3,
             legacy: false,
             block: None,
+            state: None,
             calls: vec![],
             contract,
         })
@@ -327,6 +334,7 @@ impl<M: Middleware> Multicall<M> {
             version: MulticallVersion::Multicall3,
             legacy: false,
             block: None,
+            state: None,
             calls: vec![],
             contract,
         })
@@ -363,7 +371,13 @@ impl<M: Middleware> Multicall<M> {
 
     /// Sets the `block` field of the Multicall aggregate call.
     pub fn block(mut self, block: impl Into<BlockNumber>) -> Self {
-        self.block = Some(block.into());
+        self.block = Some(block.into().into());
+        self
+    }
+
+    /// Sets the `overriding state` of the Multicall aggregate call.
+    pub fn state(mut self, state: State) -> Self {
+        self.state = Some(state);
         self
     }
 
@@ -684,7 +698,11 @@ impl<M: Middleware> Multicall<M> {
             // Wrap the return data with `success: true` since version 1 reverts if any call failed
             MulticallVersion::Multicall => {
                 let call = self.as_aggregate();
-                let (_, bytes) = ContractCall::call(&call).await?;
+                let (_, bytes) = if let Some(state) = &self.state {
+                    ContractCall::call_raw(&call).state(state).await?
+                } else {
+                    ContractCall::call(&call).await?
+                };
                 self.parse_call_result(
                     bytes
                         .into_iter()
@@ -698,7 +716,11 @@ impl<M: Middleware> Multicall<M> {
                 } else {
                     self.as_aggregate_3()
                 };
-                let results = ContractCall::call(&call).await?;
+                let results = if let Some(state) = &self.state {
+                    ContractCall::call_raw(&call).state(state).await?
+                } else {
+                    ContractCall::call(&call).await?
+                };
                 self.parse_call_result(results.into_iter())
             }
         }
@@ -869,7 +891,7 @@ impl<M: Middleware> Multicall<M> {
     /// Sets the block and legacy flags on a [ContractCall] if they were set on Multicall.
     fn set_call_flags<D: Detokenize>(&self, mut call: ContractCall<M, D>) -> ContractCall<M, D> {
         if let Some(block) = self.block {
-            call.block = Some(block.into());
+            call = call.block(block);
         }
 
         if self.legacy {
@@ -879,3 +901,14 @@ impl<M: Middleware> Multicall<M> {
         }
     }
 }
+
+impl<'a, M: Middleware> RawCall<'a> for Multicall<M> {
+    fn block(self, id: ethers_core::types::BlockId) -> Self {
+        let mut call = self;
+        call.block = Some(id);
+        call
+    }
+    fn state(self, state: &'a ethers_providers::spoof::State) -> Self {
+        self.state(state.clone())
+    }
+}

ethers-contract/tests/it/contract.rs

@@ -8,7 +8,7 @@ use ethers_core::{
     types::{Address, BlockId, Bytes, Filter, ValueOrArray, H160, H256, U256},
     utils::{keccak256, Anvil},
 };
-use ethers_providers::{Http, Middleware, MiddlewareError, Provider, StreamExt, Ws};
+use ethers_providers::{spoof, Http, Middleware, MiddlewareError, Provider, StreamExt, Ws};
 use std::{sync::Arc, time::Duration};
 
 #[derive(Debug)]
@@ -777,3 +777,95 @@ async fn multicall_aggregate() {
     assert_eq!(bytes[..4], keccak256("CustomErrorWithData(string)")[..4]);
     assert_eq!(bytes[4..], encode(&[Token::String("Data".to_string())]));
 }
+
+#[tokio::test]
+async fn test_multicall_state_overrides() {
+    // get ABI and bytecode for the Multicall contract
+    let (multicall_abi, multicall_bytecode) = get_contract("Multicall.json");
+
+    // get ABI and bytecode for the NotSoSimpleStorage contract
+    let (slot_storage_abi, slot_storage_bytecode) = get_contract("SlotStorage.json");
+
+    // launch anvil
+    let anvil = Anvil::new().spawn();
+
+    let client = connect(&anvil, 0);
+    let client2 = connect(&anvil, 1);
+
+    // create a factory which will be used to deploy instances of the contract
+    let multicall_factory = ContractFactory::new(multicall_abi, multicall_bytecode, client.clone());
+    let slot_storage_factory =
+        ContractFactory::new(slot_storage_abi, slot_storage_bytecode, client2.clone());
+
+    let multicall_contract = multicall_factory.deploy(()).unwrap().legacy().send().await.unwrap();
+    let multicall_addr = multicall_contract.address();
+
+    let value: H256 =
+        "0x312c22f60e0b666af7fce7332bfbe2a3247e19b8d612289c16b8f2e37516de36".parse().unwrap();
+    let addr = "0x851a842060FC8ae05848d08872653E30FD4c9829".parse().unwrap();
+    let slot: H256 =
+        "0xa35a6bd95953594c6d23a75dc715af91915e970ba4d87f1141e13b915e0201a3".parse().unwrap();
+
+    let slot_storage_contract =
+        slot_storage_factory.deploy(value).unwrap().legacy().send().await.unwrap();
+
+    // initiate the Multicall instance and add calls one by one in builder style
+    let mut multicall =
+        Multicall::<Provider<Http>>::new(client.clone(), Some(multicall_addr)).await.unwrap();
+
+    // test balance override
+    multicall = multicall.version(MulticallVersion::Multicall3);
+
+    let balance = 100.into();
+    let mut state = spoof::state();
+    state.account(addr).balance(balance);
+
+    multicall = multicall.state(state);
+    let (get_balance,): (U256,) =
+        multicall.clear_calls().add_get_eth_balance(addr, true).call().await.unwrap();
+    assert_eq!(get_balance, balance);
+
+    // test code override
+    let deployed_bytecode = client.get_code(slot_storage_contract.address(), None).await.unwrap();
+    state = spoof::state();
+    state.account(addr).code(deployed_bytecode);
+
+    multicall = multicall.state(state);
+    let new_value: H256 =
+        "0x5d2c59f6581053209078988fe8cad8edb594bad62e570e99ad4f5ea38049677b".parse().unwrap();
+    let (get_old_value, get_value): (H256, H256) = multicall
+        .clear_calls()
+        .add_call(
+            slot_storage_contract.at(addr).method::<_, H256>("setValue", new_value).unwrap(),
+            false,
+        )
+        .add_call(slot_storage_contract.at(addr).method::<_, H256>("getValue", ()).unwrap(), false)
+        .call()
+        .await
+        .unwrap();
+    assert_eq!(get_old_value, H256::default());
+    assert_eq!(get_value, new_value);
+
+    // test slot override
+    let deployed_bytecode = client.get_code(slot_storage_contract.address(), None).await.unwrap();
+    let old_value =
+        "0xfce2394e4cb6779bdacc1983fb24636007e9c843211586811e46b52c86d97c34".parse().unwrap();
+    state = spoof::state();
+    state.account(addr).code(deployed_bytecode).store(slot, old_value);
+
+    multicall = multicall.state(state);
+    let new_value: H256 =
+        "0x5d2c59f6581053209078988fe8cad8edb594bad62e570e99ad4f5ea38049677b".parse().unwrap();
+    let (get_old_value, get_value): (H256, H256) = multicall
+        .clear_calls()
+        .add_call(
+            slot_storage_contract.at(addr).method::<_, H256>("setValue", new_value).unwrap(),
+            false,
+        )
+        .add_call(slot_storage_contract.at(addr).method::<_, H256>("getValue", ()).unwrap(), false)
+        .call()
+        .await
+        .unwrap();
+    assert_eq!(get_old_value, old_value);
+    assert_eq!(get_value, new_value);
+}

ethers-contract/tests/solidity-contracts/SlotStorage.json

@@ -0,0 +1 @@
+{"abi":[{"inputs":[{"internalType":"bytes32","name":"value","type":"bytes32"}],"stateMutability":"nonpayable","type":"constructor"},{"anonymous":false,"inputs":[{"indexed":true,"internalType":"address","name":"author","type":"address"},{"indexed":true,"internalType":"address","name":"oldAuthor","type":"address"},{"indexed":false,"internalType":"bytes32","name":"oldValue","type":"bytes32"},{"indexed":false,"internalType":"bytes32","name":"newValue","type":"bytes32"}],"name":"ValueChanged","type":"event"},{"inputs":[],"name":"getValue","outputs":[{"internalType":"bytes32","name":"val","type":"bytes32"}],"stateMutability":"view","type":"function"},{"inputs":[],"name":"lastSender","outputs":[{"internalType":"address","name":"","type":"address"}],"stateMutability":"view","type":"function"},{"inputs":[{"internalType":"bytes32","name":"value","type":"bytes32"}],"name":"setValue","outputs":[{"internalType":"bytes32","name":"val","type":"bytes32"}],"stateMutability":"nonpayable","type":"function"}],"bytecode":"0x608060405234801561001057600080fd5b506040516106453803806106458339818101604052810190610032919061025f565b60006100426100c360201b60201c565b9050600073ffffffffffffffffffffffffffffffffffffffff163373ffffffffffffffffffffffffffffffffffffffff167f3153ff958de22b1b71d4695c0881bb8f12ee75beedd960463fb88cd70a3da3da83856040516100a492919061029b565b60405180910390a36100bb826100fc60201b60201c565b5050506102c4565b60006100f77fa35a6bd95953594c6d23a75dc715af91915e970ba4d87f1141e13b915e0201a360001b61021260201b60201c565b905090565b60008061010d6100c360201b60201c565b905060008054906101000a900473ffffffffffffffffffffffffffffffffffffffff1673ffffffffffffffffffffffffffffffffffffffff163373ffffffffffffffffffffffffffffffffffffffff167f3153ff958de22b1b71d4695c0881bb8f12ee75beedd960463fb88cd70a3da3da838660405161018e92919061029b565b60405180910390a36101c97fa35a6bd95953594c6d23a75dc715af91915e970ba4d87f1141e13b915e0201a360001b8461021d60201b60201c565b336000806101000a81548173ffffffffffffffffffffffffffffffffffffffff021916908373ffffffffffffffffffffffffffffffffffffffff16021790555080915050919050565b600081549050919050565b8082555050565b600080fd5b6000819050919050565b61023c81610229565b811461024757600080fd5b50565b60008151905061025981610233565b92915050565b60006020828403121561027557610274610224565b5b60006102838482850161024a565b91505092915050565b61029581610229565b82525050565b60006040820190506102b0600083018561028c565b6102bd602083018461028c565b9392505050565b610372806102d36000396000f3fe608060405234801561001057600080fd5b50600436106100415760003560e01c80632096525514610046578063256fec881461006457806358825b1014610082575b600080fd5b61004e6100b2565b60405161005b919061023e565b60405180910390f35b61006c6100e5565b604051610079919061029a565b60405180910390f35b61009c600480360381019061009791906102e6565b610109565b6040516100a9919061023e565b60405180910390f35b60006100e07fa35a6bd95953594c6d23a75dc715af91915e970ba4d87f1141e13b915e0201a360001b610213565b905090565b60008054906101000a900473ffffffffffffffffffffffffffffffffffffffff1681565b6000806101146100b2565b905060008054906101000a900473ffffffffffffffffffffffffffffffffffffffff1673ffffffffffffffffffffffffffffffffffffffff163373ffffffffffffffffffffffffffffffffffffffff167f3153ff958de22b1b71d4695c0881bb8f12ee75beedd960463fb88cd70a3da3da8386604051610195929190610313565b60405180910390a36101ca7fa35a6bd95953594c6d23a75dc715af91915e970ba4d87f1141e13b915e0201a360001b8461021e565b336000806101000a81548173ffffffffffffffffffffffffffffffffffffffff021916908373ffffffffffffffffffffffffffffffffffffffff16021790555080915050919050565b600081549050919050565b8082555050565b6000819050919050565b61023881610225565b82525050565b6000602082019050610253600083018461022f565b92915050565b600073ffffffffffffffffffffffffffffffffffffffff82169050919050565b600061028482610259565b9050919050565b61029481610279565b82525050565b60006020820190506102af600083018461028b565b92915050565b600080fd5b6102c381610225565b81146102ce57600080fd5b50565b6000813590506102e0816102ba565b92915050565b6000602082840312156102fc576102fb6102b5565b5b600061030a848285016102d1565b91505092915050565b6000604082019050610328600083018561022f565b610335602083018461022f565b939250505056fea264697066735822122063dccd7daa00e1af40edec15ecda7848016311815f4a960355f938f9173743ad64736f6c63430008110033"}

ethers-contract/tests/solidity-contracts/SlotStorage.sol

@@ -0,0 +1,47 @@
+pragma solidity >=0.4.24;
+
+contract SlotStorage {
+    event ValueChanged(
+        address indexed author,
+        address indexed oldAuthor,
+        bytes32 oldValue,
+        bytes32 newValue
+    );
+
+    bytes32 private constant KEY =
+        bytes32(
+            0xa35a6bd95953594c6d23a75dc715af91915e970ba4d87f1141e13b915e0201a3
+        );
+
+    address public lastSender;
+
+    constructor(bytes32 value) {
+        bytes32 _value = getValue();
+        emit ValueChanged(msg.sender, address(0), _value, value);
+        setValue(value);
+    }
+
+    function getValue() public view returns (bytes32 val) {
+        val = readBytes32(KEY);
+    }
+
+    function setValue(bytes32 value) public returns (bytes32 val) {
+        bytes32 _value = getValue();
+        emit ValueChanged(msg.sender, lastSender, _value, value);
+        writeBytes32(KEY, value);
+        lastSender = msg.sender;
+        val = _value;
+    }
+
+    function writeBytes32(bytes32 _key, bytes32 _val) internal {
+        assembly {
+            sstore(_key, _val)
+        }
+    }
+
+    function readBytes32(bytes32 _key) internal view returns (bytes32 val) {
+        assembly {
+            val := sload(_key)
+        }
+    }
+}