-
Notifications
You must be signed in to change notification settings - Fork 1.7k
/
state.rs
275 lines (250 loc) · 9.9 KB
/
state.rs
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
use super::fuzz_param_from_state;
use crate::invariant::{ArtifactFilters, FuzzRunIdentifiedContracts};
use alloy_dyn_abi::JsonAbiExt;
use alloy_json_abi::Function;
use alloy_primitives::{Address, Bytes, Log, B256, U256};
use foundry_common::contracts::{ContractsByAddress, ContractsByArtifact};
use foundry_config::FuzzDictionaryConfig;
use foundry_evm_core::utils::StateChangeset;
use parking_lot::RwLock;
use proptest::prelude::{BoxedStrategy, Strategy};
use revm::{
db::{CacheDB, DatabaseRef},
interpreter::opcode::{self, spec_opcode_gas},
primitives::SpecId,
};
use std::{fmt, sync::Arc};
// We're using `IndexSet` to have a stable element order when restoring persisted state, as well as
// for performance when iterating over the sets.
type FxIndexSet<T> =
indexmap::set::IndexSet<T, std::hash::BuildHasherDefault<rustc_hash::FxHasher>>;
/// A set of arbitrary 32 byte data from the VM used to generate values for the strategy.
///
/// Wrapped in a shareable container.
pub type EvmFuzzState = Arc<RwLock<FuzzDictionary>>;
#[derive(Default)]
pub struct FuzzDictionary {
/// Collected state values.
state_values: FxIndexSet<[u8; 32]>,
/// Addresses that already had their PUSH bytes collected.
addresses: FxIndexSet<Address>,
}
impl fmt::Debug for FuzzDictionary {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("FuzzDictionary")
.field("state_values", &self.state_values.len())
.field("addresses", &self.addresses)
.finish()
}
}
impl FuzzDictionary {
#[inline]
pub fn values(&self) -> &FxIndexSet<[u8; 32]> {
&self.state_values
}
#[inline]
pub fn values_mut(&mut self) -> &mut FxIndexSet<[u8; 32]> {
&mut self.state_values
}
#[inline]
pub fn addresses(&self) -> &FxIndexSet<Address> {
&self.addresses
}
#[inline]
pub fn addresses_mut(&mut self) -> &mut FxIndexSet<Address> {
&mut self.addresses
}
}
/// Given a function and some state, it returns a strategy which generated valid calldata for the
/// given function's input types, based on state taken from the EVM.
pub fn fuzz_calldata_from_state(func: Function, state: &EvmFuzzState) -> BoxedStrategy<Bytes> {
let strats = func
.inputs
.iter()
.map(|input| fuzz_param_from_state(&input.selector_type().parse().unwrap(), state))
.collect::<Vec<_>>();
strats
.prop_map(move |values| {
func.abi_encode_input(&values)
.unwrap_or_else(|_| {
panic!(
"Fuzzer generated invalid arguments for function `{}` with inputs {:?}: {:?}",
func.name, func.inputs, values
)
})
.into()
})
.no_shrink()
.boxed()
}
/// Builds the initial [EvmFuzzState] from a database.
pub fn build_initial_state<DB: DatabaseRef>(
db: &CacheDB<DB>,
config: &FuzzDictionaryConfig,
) -> EvmFuzzState {
let mut state = FuzzDictionary::default();
for (address, account) in db.accounts.iter() {
let address: Address = *address;
// Insert basic account information
state.values_mut().insert(address.into_word().into());
// Insert push bytes
if config.include_push_bytes {
if let Some(code) = &account.info.code {
if state.addresses_mut().insert(address) {
for push_byte in collect_push_bytes(code.bytes()) {
state.values_mut().insert(push_byte);
}
}
}
}
if config.include_storage {
// Insert storage
for (slot, value) in &account.storage {
state.values_mut().insert(B256::from(*slot).0);
state.values_mut().insert(B256::from(*value).0);
// also add the value below and above the storage value to the dictionary.
if *value != U256::ZERO {
let below_value = value - U256::from(1);
state.values_mut().insert(B256::from(below_value).0);
}
if *value != U256::MAX {
let above_value = value + U256::from(1);
state.values_mut().insert(B256::from(above_value).0);
}
}
}
}
// need at least some state data if db is empty otherwise we can't select random data for state
// fuzzing
if state.values().is_empty() {
// prefill with a random addresses
state.values_mut().insert(Address::random().into_word().into());
}
Arc::new(RwLock::new(state))
}
/// Collects state changes from a [StateChangeset] and logs into an [EvmFuzzState] according to the
/// given [FuzzDictionaryConfig].
pub fn collect_state_from_call(
logs: &[Log],
state_changeset: &StateChangeset,
state: &EvmFuzzState,
config: &FuzzDictionaryConfig,
) {
let mut state = state.write();
// Insert log topics and data.
for log in logs {
for topic in log.topics() {
state.values_mut().insert(topic.0);
}
let chunks = log.data.data.chunks_exact(32);
let rem = chunks.remainder();
for chunk in chunks {
state.values_mut().insert(chunk.try_into().unwrap());
}
if !rem.is_empty() {
state.values_mut().insert(B256::right_padding_from(rem).0);
}
}
for (address, account) in state_changeset {
// Insert basic account information
state.values_mut().insert(address.into_word().into());
if config.include_push_bytes && state.addresses.len() < config.max_fuzz_dictionary_addresses
{
// Insert push bytes
if let Some(code) = &account.info.code {
if state.addresses_mut().insert(*address) {
for push_byte in collect_push_bytes(code.bytes()) {
state.values_mut().insert(push_byte);
}
}
}
}
if config.include_storage && state.state_values.len() < config.max_fuzz_dictionary_values {
// Insert storage
for (slot, value) in &account.storage {
let value = value.present_value;
state.values_mut().insert(B256::from(*slot).0);
state.values_mut().insert(B256::from(value).0);
// also add the value below and above the storage value to the dictionary.
if value != U256::ZERO {
let below_value = value - U256::from(1);
state.values_mut().insert(B256::from(below_value).0);
}
if value != U256::MAX {
let above_value = value + U256::from(1);
state.values_mut().insert(B256::from(above_value).0);
}
}
}
}
}
/// The maximum number of bytes we will look at in bytecodes to find push bytes (24 KiB).
///
/// This is to limit the performance impact of fuzz tests that might deploy arbitrarily sized
/// bytecode (as is the case with Solmate).
const PUSH_BYTE_ANALYSIS_LIMIT: usize = 24 * 1024;
/// Collects all push bytes from the given bytecode.
fn collect_push_bytes(code: &[u8]) -> Vec<[u8; 32]> {
let mut bytes: Vec<[u8; 32]> = Vec::new();
// We use [SpecId::LATEST] since we do not really care what spec it is - we are not interested
// in gas costs.
let opcode_infos = spec_opcode_gas(SpecId::LATEST);
let mut i = 0;
while i < code.len().min(PUSH_BYTE_ANALYSIS_LIMIT) {
let op = code[i];
if opcode_infos[op as usize].is_push() {
let push_size = (op - opcode::PUSH1 + 1) as usize;
let push_start = i + 1;
let push_end = push_start + push_size;
// As a precaution, if a fuzz test deploys malformed bytecode (such as using `CREATE2`)
// this will terminate the loop early.
if push_start > code.len() || push_end > code.len() {
return bytes;
}
let push_value = U256::try_from_be_slice(&code[push_start..push_end]).unwrap();
bytes.push(push_value.to_be_bytes());
// also add the value below and above the push value to the dictionary.
if push_value != U256::ZERO {
bytes.push((push_value - U256::from(1)).to_be_bytes());
}
if push_value != U256::MAX {
bytes.push((push_value + U256::from(1)).to_be_bytes());
}
i += push_size;
}
i += 1;
}
bytes
}
/// Collects all created contracts from a StateChangeset which haven't been discovered yet. Stores
/// them at `targeted_contracts` and `created_contracts`.
pub fn collect_created_contracts(
state_changeset: &StateChangeset,
project_contracts: &ContractsByArtifact,
setup_contracts: &ContractsByAddress,
artifact_filters: &ArtifactFilters,
targeted_contracts: FuzzRunIdentifiedContracts,
created_contracts: &mut Vec<Address>,
) -> eyre::Result<()> {
let mut writable_targeted = targeted_contracts.lock();
for (address, account) in state_changeset {
if !setup_contracts.contains_key(address) {
if let (true, Some(code)) = (&account.is_touched(), &account.info.code) {
if !code.is_empty() {
if let Some((artifact, (abi, _))) =
project_contracts.find_by_code(&code.original_bytes())
{
if let Some(functions) =
artifact_filters.get_targeted_functions(artifact, abi)?
{
created_contracts.push(*address);
writable_targeted
.insert(*address, (artifact.name.clone(), abi.clone(), functions));
}
}
}
}
}
}
Ok(())
}