Support copies from arbitrary storage, docs

crates/cheatcodes/src/evm.rs

@@ -15,7 +15,7 @@ use foundry_evm_core::{
 };
 use rand::Rng;
 use revm::{
-    primitives::{Account, Bytecode, EvmStorageSlot, SpecId, KECCAK_EMPTY},
+    primitives::{Account, Bytecode, SpecId, KECCAK_EMPTY},
     InnerEvmContext,
 };
 use std::{
@@ -91,12 +91,23 @@ impl Cheatcode for loadCall {
         ensure_not_precompile!(&target, ccx);
         ccx.ecx.load_account(target)?;
         let mut val = ccx.ecx.sload(target, slot.into())?;
-        // Generate random value if target should have arbitrary storage and storage slot untouched.
-        if ccx.state.arbitrary_storage.contains(&target) && val.is_cold && val.data.is_zero() {
-            val.data = ccx.state.rng().gen();
-            let mut account = ccx.ecx.load_account(target)?;
-            account.storage.insert(slot.into(), EvmStorageSlot::new(val.data));
+
+        if val.is_cold && val.data.is_zero() {
+            let rand_value = ccx.state.rng().gen();
+            let arbitrary_storage = &mut ccx.state.arbitrary_storage;
+            if arbitrary_storage.is_arbitrary(&target) {
+                // If storage slot is untouched and load from a target with arbitrary storage,
+                // then set random value for current slot.
+                arbitrary_storage.save(ccx.ecx, target, slot.into(), rand_value);
+                val.data = rand_value;
+            } else if arbitrary_storage.is_copy(&target) {
+                // If storage slot is untouched and load from a target that copies storage from
+                // a source address with arbitrary storage, then copy existing arbitrary value.
+                // If no arbitrary value generated yet, then the random one is saved and set.
+                val.data = arbitrary_storage.copy(ccx.ecx, target, slot.into(), rand_value);
+            }
         }
+
         Ok(val.abi_encode())
     }
 }

crates/cheatcodes/src/inspector.rs

@@ -254,6 +254,92 @@ impl GasMetering {
     }
 }
 
+/// Holds data about arbitrary storage.
+#[derive(Clone, Debug, Default)]
+pub struct ArbitraryStorage {
+    /// Mapping of arbitrary storage addresses to generated values (slot, arbitrary value).
+    /// (SLOADs return random value if storage slot wasn't accessed).
+    /// Changed values are recorded and used to copy storage to different addresses.
+    pub values: HashMap<Address, HashMap<U256, U256>>,
+    /// Mapping of address with storage copied to arbitrary storage address source.
+    pub copies: HashMap<Address, Address>,
+}
+
+impl ArbitraryStorage {
+    /// Whether the given address has arbitrary storage.
+    pub fn is_arbitrary(&self, address: &Address) -> bool {
+        self.values.contains_key(address)
+    }
+
+    /// Whether the given address is a copy of an address with arbitrary storage.
+    pub fn is_copy(&self, address: &Address) -> bool {
+        self.copies.contains_key(address)
+    }
+
+    /// Marks an address with arbitrary storage.
+    pub fn mark_arbitrary(&mut self, address: &Address) {
+        self.values.insert(*address, HashMap::default());
+    }
+
+    /// Maps an address that copies storage with the arbitrary storage address.
+    pub fn mark_copy(&mut self, from: &Address, to: &Address) {
+        if self.is_arbitrary(from) {
+            self.copies.insert(*to, *from);
+        }
+    }
+
+    /// Saves arbitrary storage value for a given address:
+    /// - store value in changed values cache.
+    /// - update account's storage with given value.
+    pub fn save<DB: DatabaseExt>(
+        &mut self,
+        ecx: &mut InnerEvmContext<DB>,
+        address: Address,
+        slot: U256,
+        data: U256,
+    ) {
+        if let Ok(mut account) = ecx.load_account(address) {
+            self.values
+                .get_mut(&address)
+                .expect("missing arbitrary address entry")
+                .insert(slot, data);
+            account.storage.insert(slot, EvmStorageSlot::new(data));
+        }
+    }
+
+    /// Copies arbitrary storage value from source address to the given target address:
+    /// - if a value is present in arbitrary values cache, then update target storage and return
+    ///   existing value.
+    /// - if no value was yet generated for given slot, then save new value in cache and update both
+    ///   source and target storages.
+    pub fn copy<DB: DatabaseExt>(
+        &mut self,
+        ecx: &mut InnerEvmContext<DB>,
+        target: Address,
+        slot: U256,
+        new_value: U256,
+    ) -> U256 {
+        let source = self.copies.get(&target).expect("missing arbitrary copy target entry");
+        let storage_cache = self.values.get_mut(source).expect("missing arbitrary source storage");
+        let value = match storage_cache.get(&slot) {
+            Some(value) => *value,
+            None => {
+                storage_cache.insert(slot, new_value);
+                // Update source storage with new value.
+                if let Ok(mut source_account) = ecx.load_account(*source) {
+                    source_account.storage.insert(slot, EvmStorageSlot::new(new_value));
+                }
+                new_value
+            }
+        };
+        // Update target storage with new value.
+        if let Ok(mut target_account) = ecx.load_account(target) {
+            target_account.storage.insert(slot, EvmStorageSlot::new(value));
+        }
+        value
+    }
+}
+
 /// List of transactions that can be broadcasted.
 pub type BroadcastableTransactions = VecDeque<BroadcastableTransaction>;
 
@@ -372,9 +458,8 @@ pub struct Cheatcodes {
     /// Ignored traces.
     pub ignored_traces: IgnoredTraces,
 
-    /// Addresses that should have arbitrary storage generated (SLOADs return random value if
-    /// storage slot wasn't accessed).
-    pub arbitrary_storage: Vec<Address>,
+    /// Addresses with arbitrary storage.
+    pub arbitrary_storage: ArbitraryStorage,
 }
 
 // This is not derived because calling this in `fn new` with `..Default::default()` creates a second
@@ -1063,9 +1148,13 @@ impl<DB: DatabaseExt> Inspector<DB> for Cheatcodes {
             self.meter_gas_check(interpreter);
         }
 
-        if self.arbitrary_storage.contains(&interpreter.contract().target_address) {
+        if self.arbitrary_storage.is_arbitrary(&interpreter.contract().target_address) {
             self.ensure_arbitrary_storage(interpreter, ecx);
         }
+
+        if self.arbitrary_storage.is_copy(&interpreter.contract().target_address) {
+            self.copy_arbitrary_storage(interpreter, ecx);
+        }
     }
 
     fn log(&mut self, interpreter: &mut Interpreter, _ecx: &mut EvmContext<DB>, log: &Log) {
@@ -1479,6 +1568,10 @@ impl Cheatcodes {
     }
 
     /// Generates arbitrary values for storage slots.
+    /// Invoked in inspector `step_end`, when the current opcode is not executed.
+    /// If current opcode to execute is `SLOAD` and storage slot is cold, then an arbitrary value
+    /// is generated and saved in target address storage (therefore when `SLOAD` opcode is executed,
+    /// the arbitrary value will be returned).
     #[cold]
     fn ensure_arbitrary_storage<DB: DatabaseExt>(
         &mut self,
@@ -1490,9 +1583,42 @@ impl Cheatcodes {
             let target_address = interpreter.contract().target_address;
             if let Ok(value) = ecx.sload(target_address, key) {
                 if value.is_cold && value.data.is_zero() {
-                    if let Ok(mut target_account) = ecx.load_account(target_address) {
-                        target_account.storage.insert(key, EvmStorageSlot::new(self.rng().gen()));
-                    }
+                    let arbitrary_value = self.rng().gen();
+                    self.arbitrary_storage.save(
+                        &mut ecx.inner,
+                        target_address,
+                        key,
+                        arbitrary_value,
+                    );
+                }
+            }
+        }
+    }
+
+    /// Copies arbitrary values for storage slots.
+    /// Invoked in inspector `step_end`, when the current opcode is not executed.
+    /// If current opcode to execute is `SLOAD` and storage slot is cold, it copies the existing
+    /// arbitrary storage value (or the new generated one if no value in cache) from mapped source
+    /// address to the target address (therefore when `SLOAD` opcode is executed, the arbitrary
+    /// value will be returned).
+    #[cold]
+    fn copy_arbitrary_storage<DB: DatabaseExt>(
+        &mut self,
+        interpreter: &mut Interpreter,
+        ecx: &mut EvmContext<DB>,
+    ) {
+        if interpreter.current_opcode() == op::SLOAD {
+            let key = try_or_return!(interpreter.stack().peek(0));
+            let target_address = interpreter.contract().target_address;
+            if let Ok(value) = ecx.sload(target_address, key) {
+                if value.is_cold && value.data.is_zero() {
+                    let arbitrary_value = self.rng().gen();
+                    self.arbitrary_storage.copy(
+                        &mut ecx.inner,
+                        target_address,
+                        key,
+                        arbitrary_value,
+                    );
                 }
             }
         }

crates/cheatcodes/src/utils.rs

@@ -152,7 +152,7 @@ impl Cheatcode for resumeTracingCall {
 impl Cheatcode for setArbitraryStorageCall {
     fn apply_stateful<DB: DatabaseExt>(&self, ccx: &mut CheatsCtxt<DB>) -> Result {
         let Self { target } = self;
-        ccx.state.arbitrary_storage.push(*target);
+        ccx.state.arbitrary_storage.mark_arbitrary(target);
 
         Ok(Default::default())
     }
@@ -161,10 +161,15 @@ impl Cheatcode for setArbitraryStorageCall {
 impl Cheatcode for copyStorageCall {
     fn apply_stateful<DB: DatabaseExt>(&self, ccx: &mut CheatsCtxt<DB>) -> Result {
         let Self { from, to } = self;
+        ensure!(
+            !ccx.state.arbitrary_storage.is_arbitrary(to),
+            "target address cannot have arbitrary storage"
+        );
         if let Ok(from_account) = ccx.load_account(*from) {
             let from_storage = from_account.storage.clone();
             if let Ok(mut to_account) = ccx.load_account(*to) {
                 to_account.storage = from_storage;
+                ccx.state.arbitrary_storage.mark_copy(from, to);
             }
         }
 

testdata/default/cheats/CopyStorage.t.sol

@@ -7,6 +7,7 @@ import "cheats/Vm.sol";
 contract Counter {
     uint256 public a;
     address public b;
+    int256[] public c;
 
     function setA(uint256 _a) public {
         a = _a;
@@ -22,16 +23,14 @@ contract CounterTest is DSTest {
     Counter public counter1;
     Vm vm = Vm(HEVM_ADDRESS);
 
-    function setUp() public {
+    function test_copy_storage() public {
         counter = new Counter();
         counter.setA(1000);
         counter.setB(address(27));
         counter1 = new Counter();
         counter1.setA(11);
         counter1.setB(address(50));
-    }
 
-    function test_copy_storage() public {
         assertEq(counter.a(), 1000);
         assertEq(counter.b(), address(27));
         assertEq(counter1.a(), 11);
@@ -42,6 +41,26 @@ contract CounterTest is DSTest {
         assertEq(counter1.a(), 1000);
         assertEq(counter1.b(), address(27));
     }
+
+    function test_copy_storage_from_arbitrary() public {
+        counter = new Counter();
+        counter1 = new Counter();
+        vm.setArbitraryStorage(address(counter));
+        vm.copyStorage(address(counter), address(counter1));
+
+        // Make sure untouched storage has same values.
+        assertEq(counter.a(), counter1.a());
+        assertEq(counter.b(), counter1.b());
+        assertEq(counter.c(33), counter1.c(33));
+
+        // Change storage in source storage contract and make sure copy is not changed.
+        counter.setA(1000);
+        counter1.setB(address(50));
+        assertEq(counter.a(), 1000);
+        assertEq(counter1.a(), 40426841063417815470953489044557166618267862781491517122018165313568904172524);
+        assertEq(counter.b(), 0x485E9Cc0ef187E54A3AB45b50c3DcE43f2C223B1);
+        assertEq(counter1.b(), address(50));
+    }
 }
 
 contract CopyStorageContract {
@@ -51,6 +70,7 @@ contract CopyStorageContract {
 contract CopyStorageTest is DSTest {
     CopyStorageContract csc_1;
     CopyStorageContract csc_2;
+    CopyStorageContract csc_3;
     Vm vm = Vm(HEVM_ADDRESS);
 
     function _storeUInt256(address contractAddress, uint256 slot, uint256 value) internal {
@@ -60,6 +80,7 @@ contract CopyStorageTest is DSTest {
     function setUp() public {
         csc_1 = new CopyStorageContract();
         csc_2 = new CopyStorageContract();
+        csc_3 = new CopyStorageContract();
     }
 
     function test_copy_storage() public {
@@ -75,4 +96,63 @@ contract CopyStorageTest is DSTest {
         // `x` of second contract is now the `x` of the first
         assert(csc_2.x() == x_1);
     }
+
+    function test_copy_storage_same_values_on_load() public {
+        // Make the storage of first contract symbolic
+        vm.setArbitraryStorage(address(csc_1));
+        vm.copyStorage(address(csc_1), address(csc_2));
+        uint256 slot1 = vm.randomUint(0, 100);
+        uint256 slot2 = vm.randomUint(0, 100);
+        bytes32 value1 = vm.load(address(csc_1), bytes32(slot1));
+        bytes32 value2 = vm.load(address(csc_1), bytes32(slot2));
+
+        bytes32 value3 = vm.load(address(csc_2), bytes32(slot1));
+        bytes32 value4 = vm.load(address(csc_2), bytes32(slot2));
+
+        // Check storage values are the same for both source and target contracts.
+        assertEq(value1, value3);
+        assertEq(value2, value4);
+    }
+
+    function test_copy_storage_consistent_values() public {
+        // Make the storage of first contract symbolic.
+        vm.setArbitraryStorage(address(csc_1));
+        // Copy arbitrary storage to 2 contracts.
+        vm.copyStorage(address(csc_1), address(csc_2));
+        vm.copyStorage(address(csc_1), address(csc_3));
+        uint256 slot1 = vm.randomUint(0, 100);
+        uint256 slot2 = vm.randomUint(0, 100);
+
+        // Load slot 1 from 1st copied contract and slot2 from symbolic contract.
+        bytes32 value3 = vm.load(address(csc_2), bytes32(slot1));
+        bytes32 value2 = vm.load(address(csc_1), bytes32(slot2));
+
+        bytes32 value1 = vm.load(address(csc_1), bytes32(slot1));
+        bytes32 value4 = vm.load(address(csc_2), bytes32(slot2));
+
+        // Make sure same values for both copied and symbolic contract.
+        assertEq(value3, value1);
+        assertEq(value2, value4);
+
+        uint256 x_1 = vm.randomUint();
+        // Change slot1 of 1st copied contract.
+        _storeUInt256(address(csc_2), slot1, x_1);
+        value3 = vm.load(address(csc_2), bytes32(slot1));
+        bytes32 value5 = vm.load(address(csc_3), bytes32(slot1));
+        // Make sure value for 1st contract copied is different than symbolic contract value.
+        assert(value3 != value1);
+        // Make sure same values for 2nd contract copied and symbolic contract.
+        assertEq(value5, value1);
+
+        uint256 x_2 = vm.randomUint();
+        // Change slot2 of symbolic contract.
+        _storeUInt256(address(csc_1), slot2, x_2);
+        value2 = vm.load(address(csc_1), bytes32(slot2));
+        bytes32 value6 = vm.load(address(csc_3), bytes32(slot2));
+        // Make sure value for symbolic contract value is different than 1st contract copied.
+        assert(value2 != value4);
+        // Make sure value for symbolic contract value is different than 2nd contract copied.
+        assert(value2 != value6);
+        assertEq(value4, value6);
+    }
 }