From 0a20bd72cdf2e72fea74efd0bf6b76e3fcbe256c Mon Sep 17 00:00:00 2001 From: Dariusz Doktorski Date: Tue, 28 Jan 2025 17:04:10 +0100 Subject: [PATCH] Add `generate_arg` cheatcode --- .../forge_runtime_extension/fuzzer.rs | 38 +++++++++++++++++++ .../forge_runtime_extension/mod.rs | 12 ++++++ crates/forge-runner/src/running.rs | 1 + snforge_std/src/cheatcodes.cairo | 1 + snforge_std/src/cheatcodes/generate_arg.cairo | 8 ++++ 5 files changed, 60 insertions(+) create mode 100644 crates/cheatnet/src/runtime_extensions/forge_runtime_extension/fuzzer.rs create mode 100644 snforge_std/src/cheatcodes/generate_arg.cairo diff --git a/crates/cheatnet/src/runtime_extensions/forge_runtime_extension/fuzzer.rs b/crates/cheatnet/src/runtime_extensions/forge_runtime_extension/fuzzer.rs new file mode 100644 index 0000000000..03c4be7f39 --- /dev/null +++ b/crates/cheatnet/src/runtime_extensions/forge_runtime_extension/fuzzer.rs @@ -0,0 +1,38 @@ +use anyhow::ensure; +use num_bigint::RandBigInt; +use rand::prelude::StdRng; +use starknet_types_core::felt::Felt; +use std::sync::{Arc, Mutex}; + +pub(crate) fn generate_arg( + fuzzer_rng: Option>>, + min_value: Felt, + max_value: Felt, +) -> anyhow::Result { + let min_big_int = if min_value > (Felt::MAX + Felt::from(i128::MIN)) && min_value > max_value { + // Negative value x is serialized as P + x, where P is the STARK prime number + // hence to deserialize and get the actual x we need to subtract P (== Felt::MAX + 1) + min_value.to_bigint() - Felt::MAX.to_bigint() - 1 + } else { + min_value.to_bigint() + }; + + let max_big_int = max_value.to_bigint(); + + ensure!( + min_big_int <= max_big_int, + format!("`generate_arg` cheatcode: `min_value` must be <= `max_value`, provided values after deserialization: {min_big_int} and {max_big_int}") + ); + + let value = if let Some(fuzzer_rng) = fuzzer_rng { + fuzzer_rng + .lock() + .unwrap() + .gen_bigint_range(&min_big_int, &(max_big_int + 1)) + } else { + // `generate_arg` cheatcode can be also used outside the fuzzer context + rand::thread_rng().gen_bigint_range(&min_big_int, &(max_big_int + 1)) + }; + + Ok(Felt::from(value)) +} diff --git a/crates/cheatnet/src/runtime_extensions/forge_runtime_extension/mod.rs b/crates/cheatnet/src/runtime_extensions/forge_runtime_extension/mod.rs index 2938c7bf97..489d221034 100644 --- a/crates/cheatnet/src/runtime_extensions/forge_runtime_extension/mod.rs +++ b/crates/cheatnet/src/runtime_extensions/forge_runtime_extension/mod.rs @@ -40,6 +40,7 @@ use conversions::felt::TryInferFormat; use conversions::serde::deserialize::BufferReader; use conversions::serde::serialize::CairoSerialize; use data_transformer::cairo_types::CairoU256; +use rand::prelude::StdRng; use runtime::{ CheatcodeHandlingResult, EnhancedHintError, ExtendedRuntime, ExtensionLogic, SyscallHandlingResult, @@ -48,16 +49,19 @@ use starknet::signers::SigningKey; use starknet_api::{core::ClassHash, deprecated_contract_class::EntryPointType::L1Handler}; use starknet_types_core::felt::Felt; use std::collections::HashMap; +use std::sync::{Arc, Mutex}; pub mod cheatcodes; pub mod contracts_data; mod file_operations; +mod fuzzer; pub type ForgeRuntime<'a> = ExtendedRuntime>; pub struct ForgeExtension<'a> { pub environment_variables: &'a HashMap, pub contracts_data: &'a ContractsData, + pub fuzzer_rng: Option>>, } // This runtime extension provides an implementation logic for functions from snforge_std library. @@ -476,6 +480,14 @@ impl<'a> ExtensionLogic for ForgeExtension<'a> { "generate_random_felt" => Ok(CheatcodeHandlingResult::from_serializable( generate_random_felt(), )), + "generate_arg" => { + let min_value: Felt = input_reader.read()?; + let max_value: Felt = input_reader.read()?; + + Ok(CheatcodeHandlingResult::from_serializable( + fuzzer::generate_arg(self.fuzzer_rng.clone(), min_value, max_value)?, + )) + } _ => Ok(CheatcodeHandlingResult::Forwarded), } } diff --git a/crates/forge-runner/src/running.rs b/crates/forge-runner/src/running.rs index ea913c2caf..2df1af58b0 100644 --- a/crates/forge-runner/src/running.rs +++ b/crates/forge-runner/src/running.rs @@ -202,6 +202,7 @@ pub fn run_test_case( let forge_extension = ForgeExtension { environment_variables: runtime_config.environment_variables, contracts_data: runtime_config.contracts_data, + fuzzer_rng: None, }; let mut forge_runtime = ExtendedRuntime { diff --git a/snforge_std/src/cheatcodes.cairo b/snforge_std/src/cheatcodes.cairo index 1b3665bc9d..db2af2e0b9 100644 --- a/snforge_std/src/cheatcodes.cairo +++ b/snforge_std/src/cheatcodes.cairo @@ -8,6 +8,7 @@ pub mod storage; pub mod execution_info; pub mod message_to_l1; pub mod generate_random_felt; +pub mod generate_arg; /// Enum used to specify how long the target should be cheated for. #[derive(Copy, Drop, Serde, PartialEq, Clone, Debug)] diff --git a/snforge_std/src/cheatcodes/generate_arg.cairo b/snforge_std/src/cheatcodes/generate_arg.cairo new file mode 100644 index 0000000000..dd7c046351 --- /dev/null +++ b/snforge_std/src/cheatcodes/generate_arg.cairo @@ -0,0 +1,8 @@ +use super::super::_cheatcode::execute_cheatcode_and_deserialize; + +// Generates a random number that is used for creating data for fuzz tests +pub fn generate_arg, +Drop, +Into>(min_value: T, max_value: T) -> T { + execute_cheatcode_and_deserialize::< + 'generate_arg' + >(array![min_value.into(), max_value.into()].span()) +}