feat(invariants): use SlotIdentifier for identifying complex types and sampling them (#11450)

* feat(`forge`): sample typed storage values

* arc it

* nit

* clippy

* nit

* strip file prefixes

* fmt

* don't add adjacent values to sample

* feat(cheatcodes): add contract identifier to AccountStateDiffs

* forge fmt

* doc nits

* fix tests

* feat(`cheatcodes`): include `SlotInfo` in SlotStateDiff

* cleanup + identify slots of static arrays

* nits

* nit

* nits

* test + nits

* docs

* handle 2d arrays

* use DynSolType

* feat: decode storage values

* doc nit

* skip decoded serialization if none

* nit

* fmt

* fix

* fix

* fix

* feat(cheatcodes): decode structs in state diff output

* fix

* while decode

* fix: show only decoded in plaintext / display output + test

* feat: format slots to only significant bits in vm.getStateDiff output

* encode_prefixed

* nit

* chore: add @onbjerg to `CODEOWNERS` (#11343)

* add @onbjerg

* add @0xrusowsky

* resolve conflicts

* fix: disable tx gas limit cap (#11347)

* chore(deps): bump all dependencies (#11349)

* chore: use get_or_calculate_hash better (#11350)

* resolve more conflicts

* fix(lint): 'unwrapped-modifier-logic' incorrectly marked with `Severity::Gas` (#11358)

fix(lint): 'unwrapped-modifier-logic' incorrectly marked with Severity::Gas

* feat: identify and decode nested structs

* cleanup

* decode structs and members recursively

* cleanup

* doc fix

* feat(cheatcodes): decode mappings in state diffs (#11381)

* feat(cheatcodes): decode mappings in state diffs

* feat: decode nested mappings

* assert vm.getStateDiff output

* feat: add `keys` fields to `SlotInfo` in case of mappings

* remove wrapper

* refactor: moves state diff decoding to common (#11413)

* refactor: storage decoder

* cleanup

* dedup MappingSlots by moving it to common

* move decoding logic into SlotInfo

* rename to SlotIndentifier

* docs

* fix: delegate identification according to encoding types

* clippy + fmt

* docs fix

* fix

* merge match arms

* merge ifs

* recurse handle_struct

* dedup assertContains test util

* fix

* Update crates/common/src/slot_identifier.rs

Co-authored-by: zerosnacks <[email protected]>

* Review changes: simplify get or insert, use common fmt

* alloy-dyn-abi.workspace

* identify slot types using `SlotIdentifier`

* clippy

* feat(`invariants`): record mapping keys and slots to identify their types for sampling

* fix

* only insert if value decodes

* tracing::info logs

* remove logs

* nit

* rm

---------

Co-authored-by: Yash Atreya <[email protected]>
Co-authored-by: zerosnacks <[email protected]>
Co-authored-by: Matthias Seitz <[email protected]>
Co-authored-by: DaniPopes <[email protected]>
Co-authored-by: srdtrk <[email protected]>
Co-authored-by: 0xrusowsky <[email protected]>
Co-authored-by: grandizzy <[email protected]>
Co-authored-by: grandizzy <[email protected]>

Author: 8 people

Cargo.lock

@@ -1,11 +1,7 @@
 use crate::{Cheatcode, Cheatcodes, Result, Vm::*};
-use alloy_primitives::{Address, B256, U256, keccak256, map::AddressHashMap};
+use alloy_primitives::{Address, B256};
 use alloy_sol_types::SolValue;
 use foundry_common::mapping_slots::MappingSlots;
-use revm::{
-    bytecode::opcode,
-    interpreter::{Interpreter, interpreter_types::Jumps},
-};
 
 impl Cheatcode for startMappingRecordingCall {
     fn apply(&self, state: &mut Cheatcodes) -> Result {
@@ -74,29 +70,3 @@ fn slot_child<'a>(
 ) -> Option<&'a Vec<B256>> {
     mapping_slot(state, target)?.children.get(slot)
 }
-
-#[cold]
-pub(crate) fn step(mapping_slots: &mut AddressHashMap<MappingSlots>, interpreter: &Interpreter) {
-    match interpreter.bytecode.opcode() {
-        opcode::KECCAK256 => {
-            if interpreter.stack.peek(1) == Ok(U256::from(0x40)) {
-                let address = interpreter.input.target_address;
-                let offset = interpreter.stack.peek(0).expect("stack size > 1").saturating_to();
-                let data = interpreter.memory.slice_len(offset, 0x40);
-                let low = B256::from_slice(&data[..0x20]);
-                let high = B256::from_slice(&data[0x20..]);
-                let result = keccak256(&*data);
-
-                mapping_slots.entry(address).or_default().seen_sha3.insert(result, (low, high));
-            }
-        }
-        opcode::SSTORE => {
-            if let Some(mapping_slots) = mapping_slots.get_mut(&interpreter.input.target_address)
-                && let Ok(slot) = interpreter.stack.peek(0)
-            {
-                mapping_slots.insert(slot.into());
-            }
-        }
-        _ => {}
-    }
-}

crates/cheatcodes/src/evm/mapping.rs

@@ -4,7 +4,7 @@ use crate::{
     CheatsConfig, CheatsCtxt, DynCheatcode, Error, Result,
     Vm::{self, AccountAccess},
     evm::{
-        DealRecord, GasRecord, RecordAccess, mapping,
+        DealRecord, GasRecord, RecordAccess,
         mock::{MockCallDataContext, MockCallReturnData},
         prank::Prank,
     },
@@ -33,7 +33,9 @@ use alloy_rpc_types::{
 };
 use alloy_sol_types::{SolCall, SolInterface, SolValue};
 use foundry_common::{
-    SELECTOR_LEN, TransactionMaybeSigned, evm::Breakpoints, mapping_slots::MappingSlots,
+    SELECTOR_LEN, TransactionMaybeSigned,
+    evm::Breakpoints,
+    mapping_slots::{MappingSlots, step as mapping_step},
 };
 use foundry_evm_core::{
     InspectorExt,
@@ -1103,7 +1105,7 @@ impl Inspector<EthEvmContext<&mut dyn DatabaseExt>> for Cheatcodes {
 
         // `startMappingRecording`: record SSTORE and KECCAK256.
         if let Some(mapping_slots) = &mut self.mapping_slots {
-            mapping::step(mapping_slots, interpreter);
+            mapping_step(mapping_slots, interpreter);
         }
 
         // `snapshotGas*`: take a snapshot of the current gas.

crates/cheatcodes/src/inspector.rs

@@ -44,6 +44,8 @@ alloy-transport.workspace = true
 alloy-consensus = { workspace = true, features = ["k256"] }
 alloy-network.workspace = true
 
+revm.workspace = true
+
 solar.workspace = true
 
 tower.workspace = true

crates/common/Cargo.toml

@@ -1,4 +1,11 @@
-use alloy_primitives::{B256, map::B256HashMap};
+use alloy_primitives::{
+    B256, U256, keccak256,
+    map::{AddressHashMap, B256HashMap},
+};
+use revm::{
+    bytecode::opcode,
+    interpreter::{Interpreter, interpreter_types::Jumps},
+};
 
 /// Recorded mapping slots.
 #[derive(Clone, Debug, Default)]
@@ -36,3 +43,30 @@ impl MappingSlots {
         }
     }
 }
+
+/// Function to be used in Inspector::step to record mapping slots and keys
+#[cold]
+pub fn step(mapping_slots: &mut AddressHashMap<MappingSlots>, interpreter: &Interpreter) {
+    match interpreter.bytecode.opcode() {
+        opcode::KECCAK256 => {
+            if interpreter.stack.peek(1) == Ok(U256::from(0x40)) {
+                let address = interpreter.input.target_address;
+                let offset = interpreter.stack.peek(0).expect("stack size > 1").saturating_to();
+                let data = interpreter.memory.slice_len(offset, 0x40);
+                let low = B256::from_slice(&data[..0x20]);
+                let high = B256::from_slice(&data[0x20..]);
+                let result = keccak256(&*data);
+
+                mapping_slots.entry(address).or_default().seen_sha3.insert(result, (low, high));
+            }
+        }
+        opcode::SSTORE => {
+            if let Some(mapping_slots) = mapping_slots.get_mut(&interpreter.input.target_address)
+                && let Ok(slot) = interpreter.stack.peek(0)
+            {
+                mapping_slots.insert(slot.into());
+            }
+        }
+        _ => {}
+    }
+}

crates/common/src/mapping_slots.rs

@@ -197,6 +197,7 @@ impl SlotIdentifier {
     ///
     /// It can also identify whether a slot belongs to a mapping if provided with [`MappingSlots`].
     pub fn identify(&self, slot: &B256, mapping_slots: Option<&MappingSlots>) -> Option<SlotInfo> {
+        trace!(?slot, "identifying slot");
         let slot_u256 = U256::from_be_bytes(slot.0);
         let slot_str = slot_u256.to_string();
 
@@ -629,7 +630,7 @@ impl SlotIdentifier {
 
         // Check if the value is another mapping (nested case)
         if let Some(value_storage_type) = self.storage_layout.types.get(value_type_ref) {
-            if value_storage_type.encoding == "mapping" {
+            if value_storage_type.encoding == ENCODING_MAPPING {
                 // Recursively resolve the nested mapping
                 let (nested_keys, final_value, _) = self.resolve_mapping_type(value_type_ref)?;
                 key_types.extend(nested_keys);

crates/common/src/slot_identifier.rs

@@ -2,7 +2,10 @@ use crate::{
     executors::{Executor, RawCallResult},
     inspectors::Fuzzer,
 };
-use alloy_primitives::{Address, Bytes, FixedBytes, Selector, U256, map::HashMap};
+use alloy_primitives::{
+    Address, Bytes, FixedBytes, Selector, U256,
+    map::{AddressMap, HashMap},
+};
 use alloy_sol_types::{SolCall, sol};
 use eyre::{ContextCompat, Result, eyre};
 use foundry_common::contracts::{ContractsByAddress, ContractsByArtifact};
@@ -598,6 +601,15 @@ impl<'a> InvariantExecutor<'a> {
             ));
         }
 
+        // If any of the targeted contracts have the storage layout enabled then we can sample
+        // mapping values. To accomplish, we need to record the mapping storage slots and keys.
+        let fuzz_state =
+            if targeted_contracts.targets.lock().iter().any(|(_, t)| t.storage_layout.is_some()) {
+                fuzz_state.with_mapping_slots(AddressMap::default())
+            } else {
+                fuzz_state
+            };
+
         self.executor.inspector_mut().set_fuzzer(Fuzzer {
             call_generator,
             fuzz_state: fuzz_state.clone(),

crates/evm/evm/src/executors/invariant/mod.rs

@@ -1,4 +1,5 @@
 use crate::{invariant::RandomCallGenerator, strategies::EvmFuzzState};
+use foundry_common::mapping_slots::step as mapping_step;
 use revm::{
     Inspector,
     context::{ContextTr, Transaction},
@@ -26,6 +27,9 @@ where
         // We only collect `stack` and `memory` data before and after calls.
         if self.collect {
             self.collect_data(interp);
+            if let Some(mapping_slots) = &mut self.fuzz_state.mapping_slots {
+                mapping_step(mapping_slots, interp);
+            }
         }
     }