Skip to content

Commit

Permalink
Add support for bytes and string fixtures, fixture strategy macro. So…
Browse files Browse the repository at this point in the history
…lidity test
  • Loading branch information
grandizzy committed Mar 25, 2024
1 parent f93254a commit 69e8c51
Show file tree
Hide file tree
Showing 7 changed files with 313 additions and 40 deletions.
38 changes: 14 additions & 24 deletions crates/evm/fuzz/src/strategies/address.rs
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
use crate::strategies::fixture_strategy;
use alloy_dyn_abi::DynSolValue;
use alloy_primitives::Address;
use proptest::{
Expand All @@ -24,30 +25,19 @@ pub struct AddressStrategy {}
impl AddressStrategy {
/// Create a new address strategy.
pub fn init(fixtures: Option<&[DynSolValue]>) -> BoxedStrategy<DynSolValue> {
if let Some(fixtures) = fixtures {
let address_fixtures: Vec<DynSolValue> =
fixtures.iter().enumerate().map(|(_, value)| value.to_owned()).collect();
let address_fixtures_len = address_fixtures.len();
any::<prop::sample::Index>()
.prop_map(move |index| {
// Generate value tree from fixture.
// If fixture is not a valid address, raise error and generate random value.
let index = index.index(address_fixtures_len);
if let Some(addr_fixture) = address_fixtures.get(index) {
if let Some(addr_fixture) = addr_fixture.as_address() {
return DynSolValue::Address(addr_fixture);
}
}
error!(
"{:?} is not a valid address fixture, generate random value",
address_fixtures.get(index)
);
DynSolValue::Address(Address::random())
})
.boxed()
} else {
// If no config for addresses dictionary then create unbounded addresses strategy.
let value_from_fixture = |fixture: Option<&DynSolValue>| {
if let Some(fixture) = fixture {
if let Some(fixture) = fixture.as_address() {
return DynSolValue::Address(fixture);
}
}
error!("{:?} is not a valid address fixture, generate random value", fixture);
DynSolValue::Address(Address::random())
};
fixture_strategy!(
fixtures,
value_from_fixture,
any::<Address>().prop_map(DynSolValue::Address).boxed()
}
)
}
}
86 changes: 86 additions & 0 deletions crates/evm/fuzz/src/strategies/bytes.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,86 @@
use crate::strategies::fixture_strategy;
use alloy_dyn_abi::{DynSolType, DynSolValue};
use alloy_primitives::B256;
use proptest::{
arbitrary::any,
prelude::{prop, BoxedStrategy},
strategy::Strategy,
};

/// The bytes strategy combines 2 different strategies:
/// 1. A random bytes strategy if no fixture defined for current parameter.
/// 2. A fixture based strategy if configured values for current parameters.
/// If fixture is not a valid type then an error is raised and test suite will continue to execute
/// with random values.
///
///
/// For example:
/// To define fixture for `backup` fuzzed parameter, return an array of possible values from
/// `function fixture_backup() external pure returns (bytes[] memory)`.
/// Use `backup` named parameter in fuzzed test in order to create a custom strategy
/// `function testFuzz_backupValue(bytes memory backup)`.
#[derive(Debug, Default)]
pub struct BytesStrategy {}

impl BytesStrategy {
/// Create a new bytes strategy.
pub fn init(fixtures: Option<&[DynSolValue]>) -> BoxedStrategy<DynSolValue> {
let value_from_fixture = |fixture: Option<&DynSolValue>| {
if let Some(fixture) = fixture {
if let Some(fixture) = fixture.as_bytes() {
return DynSolValue::Bytes(fixture.to_vec());
}
}
error!("{:?} is not a valid bytes fixture, generate random value", fixture);
let random: [u8; 32] = rand::random();
DynSolValue::Bytes(random.to_vec())
};
fixture_strategy!(
fixtures,
value_from_fixture,
DynSolValue::type_strategy(&DynSolType::Bytes).boxed()
)
}
}

/// The fixed bytes strategy combines 2 different strategies:
/// 1. A random fixed bytes strategy if no fixture defined for current parameter.
/// 2. A fixture based strategy if configured values for current parameters.
/// If fixture is not a valid type then an error is raised and test suite will continue to execute
/// with random values.
///
///
/// For example:
/// To define fixture for `key` fuzzed parameter, return an array of possible values from
/// `function fixture_key() external pure returns (bytes32[] memory)`.
/// Use `key` named parameter in fuzzed test in order to create a custom strategy
/// `function testFuzz_keyValue(bytes32 key)`.
#[derive(Debug, Default)]
pub struct FixedBytesStrategy {}

impl FixedBytesStrategy {
/// Create a new fixed bytes strategy.
pub fn init(size: usize, fixtures: Option<&[DynSolValue]>) -> BoxedStrategy<DynSolValue> {
let value_from_fixture = move |fixture: Option<&DynSolValue>| {
if let Some(fixture) = fixture {
if let Some(fixture) = fixture.as_fixed_bytes() {
if fixture.1 == size {
return DynSolValue::FixedBytes(B256::from_slice(fixture.0), fixture.1);
}
}
}
error!("{:?} is not a valid fixed bytes fixture, generate random value", fixture);
DynSolValue::FixedBytes(B256::random(), size)
};
fixture_strategy!(
fixtures,
value_from_fixture,
any::<B256>()
.prop_map(move |mut v| {
v[size..].fill(0);
DynSolValue::FixedBytes(v, size)
})
.boxed()
)
}
}
33 changes: 33 additions & 0 deletions crates/evm/fuzz/src/strategies/mod.rs
Original file line number Diff line number Diff line change
@@ -1,6 +1,9 @@
mod address;
pub use address::AddressStrategy;

mod bytes;
pub use bytes::{BytesStrategy, FixedBytesStrategy};

mod int;
pub use int::IntStrategy;

Expand All @@ -18,5 +21,35 @@ pub use state::{
build_initial_state, collect_created_contracts, collect_state_from_call, EvmFuzzState,
};

mod string;
pub use string::StringStrategy;

mod invariants;
pub use invariants::{fuzz_contract_with_calldata, invariant_strat, override_call_strat};

/// Macro to create strategy with fixtures.
/// 1. A default strategy if no fixture defined for current parameter.
/// 2. A fixture based strategy if configured values for current parameter.
/// If fixture is not a valid type then an error is raised and test suite will continue to execute
/// with random values.
macro_rules! fixture_strategy {
($fixtures:ident, $value_from_fixture:expr, $default_strategy:expr) => {
if let Some(fixtures) = $fixtures {
let custom_fixtures: Vec<DynSolValue> =
fixtures.iter().enumerate().map(|(_, value)| value.to_owned()).collect();
let custom_fixtures_len = custom_fixtures.len();
any::<prop::sample::Index>()
.prop_map(move |index| {
// Generate value tree from fixture.
// If fixture is not a valid type, raise error and generate random value.
let index = index.index(custom_fixtures_len);
$value_from_fixture(custom_fixtures.get(index))
})
.boxed()
} else {
return $default_strategy
}
};
}

pub(crate) use fixture_strategy;
21 changes: 5 additions & 16 deletions crates/evm/fuzz/src/strategies/param.rs
Original file line number Diff line number Diff line change
Expand Up @@ -33,23 +33,12 @@ pub fn fuzz_param(
DynSolType::Uint(n @ 8..=256) => super::UintStrategy::new(n, fuzz_fixtures)
.prop_map(move |x| DynSolValue::Uint(x, n))
.boxed(),
DynSolType::Function | DynSolType::Bool | DynSolType::Bytes => {
DynSolValue::type_strategy(param).boxed()
DynSolType::Function | DynSolType::Bool => DynSolValue::type_strategy(param).boxed(),
DynSolType::Bytes => super::BytesStrategy::init(fuzz_fixtures),
DynSolType::FixedBytes(size @ 1..=32) => {
super::FixedBytesStrategy::init(size, fuzz_fixtures)
}
DynSolType::FixedBytes(size @ 1..=32) => any::<B256>()
.prop_map(move |mut v| {
v[size..].fill(0);
DynSolValue::FixedBytes(v, size)
})
.boxed(),
DynSolType::String => DynSolValue::type_strategy(param)
.prop_map(move |value| {
DynSolValue::String(
value.as_str().unwrap().trim().trim_end_matches('\0').to_string(),
)
})
.boxed(),

DynSolType::String => super::StringStrategy::init(fuzz_fixtures),
DynSolType::Tuple(ref params) => params
.iter()
.map(|p| fuzz_param(p, None))
Expand Down
53 changes: 53 additions & 0 deletions crates/evm/fuzz/src/strategies/string.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,53 @@
use crate::strategies::fixture_strategy;
use alloy_dyn_abi::{DynSolType, DynSolValue};
use proptest::{
arbitrary::any,
prelude::{prop, BoxedStrategy},
strategy::Strategy,
};
use rand::{distributions::Alphanumeric, thread_rng, Rng};

/// The address strategy combines 2 different strategies:
/// 1. A random addresses strategy if no fixture defined for current parameter.
/// 2. A fixture based strategy if configured values for current parameters.
/// If fixture is not a valid type then an error is raised and test suite will continue to execute
// with random strings.
///
///
/// For example:
/// To define fixture for `person` fuzzed parameter, return an array of possible values from
/// `function fixture_person() public returns (string[] memory)`.
/// Use `person` named parameter in fuzzed test in order to create a custom strategy
/// `function testFuzz_personValue(string memory person)`.
#[derive(Debug, Default)]
pub struct StringStrategy {}

impl StringStrategy {
/// Create a new string strategy.
pub fn init(fixtures: Option<&[DynSolValue]>) -> BoxedStrategy<DynSolValue> {
let value_from_fixture = |fixture: Option<&DynSolValue>| {
if let Some(fixture) = fixture {
if let Some(fixture) = fixture.as_str() {
return DynSolValue::String(fixture.to_string());
}
}
error!("{:?} is not a valid string fixture, generate random value", fixture);
let mut rng = thread_rng();
let string_len = rng.gen_range(0..128);
let random: String =
(&mut rng).sample_iter(Alphanumeric).map(char::from).take(string_len).collect();
DynSolValue::String(random)
};
fixture_strategy!(
fixtures,
value_from_fixture,
DynSolValue::type_strategy(&DynSolType::String)
.prop_map(move |value| {
DynSolValue::String(
value.as_str().unwrap().trim().trim_end_matches('\0').to_string(),
)
})
.boxed()
)
}
}
30 changes: 30 additions & 0 deletions crates/forge/tests/it/invariant.rs
Original file line number Diff line number Diff line change
Expand Up @@ -151,6 +151,16 @@ async fn test_invariant() {
"default/fuzz/invariant/common/InvariantAssume.t.sol:InvariantAssume",
vec![("invariant_dummy()", true, None, None, None)],
),
(
"default/fuzz/invariant/common/InvariantFixtures.t.sol:InvariantFixtures",
vec![(
"invariant_target_not_compromised()",
false,
Some("<empty revert data>".into()),
None,
None,
)],
),
]),
);
}
Expand Down Expand Up @@ -400,3 +410,23 @@ async fn test_invariant_assume_respects_restrictions() {
)]),
);
}

#[tokio::test(flavor = "multi_thread")]
async fn test_invariant_fixtures() {
let filter = Filter::new(".*", ".*", ".*fuzz/invariant/common/InvariantFixtures.t.sol");
let mut runner = TEST_DATA_DEFAULT.runner();
let results = runner.test_collect(&filter);
assert_multiple(
&results,
BTreeMap::from([(
"default/fuzz/invariant/common/InvariantFixtures.t.sol:InvariantFixtures",
vec![(
"invariant_target_not_compromised()",
false,
Some("<empty revert data>".into()),
None,
None,
)],
)]),
);
}
92 changes: 92 additions & 0 deletions testdata/default/fuzz/invariant/common/InvariantFixtures.t.sol
Original file line number Diff line number Diff line change
@@ -0,0 +1,92 @@
// SPDX-License-Identifier: MIT OR Apache-2.0
pragma solidity ^0.8.0;

import "ds-test/test.sol";

contract Target {
bool ownerFound;
bool amountFound;
bool magicFound;
bool keyFound;
bool backupFound;
bool extraStringFound;

function fuzzWithFixtures(
address owner,
uint256 amount,
int32 magic,
bytes32 key,
bytes memory backup,
string memory extra
) external {
if (owner == 0x6B175474E89094C44Da98b954EedeAC495271d0F)
ownerFound = true;
if (amount == 1122334455) amountFound = true;
if (magic == -777) magicFound = true;
if (key == "abcd1234") keyFound = true;
if (keccak256(backup) == keccak256("qwerty1234")) backupFound = true;
if (
keccak256(abi.encodePacked(extra)) ==
keccak256(abi.encodePacked("112233aabbccdd"))
) extraStringFound = true;
}

function isCompromised() public view returns (bool) {
return
ownerFound &&
amountFound &&
magicFound &&
keyFound &&
backupFound &&
extraStringFound;
}
}

/// Try to compromise target contract by finding all accepted values using fixtures.
contract InvariantFixtures is DSTest {
Target target;

function setUp() public {
target = new Target();
}

function fixture_owner() external pure returns (address[] memory) {
address[] memory addressFixture = new address[](1);
addressFixture[0] = 0x6B175474E89094C44Da98b954EedeAC495271d0F;
return addressFixture;
}

function fixture_amount() external pure returns (uint256[] memory) {
uint256[] memory amountFixture = new uint256[](1);
amountFixture[0] = 1122334455;
return amountFixture;
}

function fixture_magic() external pure returns (int32[] memory) {
int32[] memory magicFixture = new int32[](1);
magicFixture[0] = -777;
return magicFixture;
}

function fixture_key() external pure returns (bytes32[] memory) {
bytes32[] memory keyFixture = new bytes32[](1);
keyFixture[0] = "abcd1234";
return keyFixture;
}

function fixture_backup() external pure returns (bytes[] memory) {
bytes[] memory backupFixture = new bytes[](1);
backupFixture[0] = "qwerty1234";
return backupFixture;
}

function fixture_extra() external pure returns (string[] memory) {
string[] memory extraFixture = new string[](1);
extraFixture[0] = "112233aabbccdd";
return extraFixture;
}

function invariant_target_not_compromised() public {
assertEq(target.isCompromised(), false);
}
}

0 comments on commit 69e8c51

Please sign in to comment.