Skip to content

Commit

Permalink
fix(cheatcodes): overflow in randomNumber w/range (#8361)
Browse files Browse the repository at this point in the history
  • Loading branch information
DaniPopes committed Jul 4, 2024
1 parent eff3f43 commit 56dbd20
Show file tree
Hide file tree
Showing 2 changed files with 19 additions and 14 deletions.
9 changes: 7 additions & 2 deletions crates/cheatcodes/src/utils.rs
Original file line number Diff line number Diff line change
Expand Up @@ -162,8 +162,13 @@ impl Cheatcode for randomUint_1Call {
ensure!(min <= max, "min must be less than or equal to max");
// Generate random between range min..=max
let mut rng = rand::thread_rng();
let range = max - min + U256::from(1);
let random_number = rng.gen::<U256>() % range + min;
let exclusive_modulo = max - min;
let mut random_number = rng.gen::<U256>();
if exclusive_modulo != U256::MAX {
let inclusive_modulo = exclusive_modulo + U256::from(1);
random_number %= inclusive_modulo;
}
random_number += min;
Ok(random_number.abi_encode())
}
}
Expand Down
24 changes: 12 additions & 12 deletions testdata/default/cheats/RandomUint.t.sol
Original file line number Diff line number Diff line change
Expand Up @@ -7,27 +7,27 @@ import "cheats/Vm.sol";
contract RandomUint is DSTest {
Vm constant vm = Vm(HEVM_ADDRESS);

// All tests use `>=` and `<=` to verify that ranges are inclusive and that
// a value of zero may be generated.
function testRandomUint() public {
uint256 rand = vm.randomUint();
assertTrue(rand >= 0);
vm.randomUint();
}

function testRandomUint(uint256 min, uint256 max) public {
vm.assume(max >= min);
uint256 rand = vm.randomUint(min, max);
assertTrue(rand >= min, "rand >= min");
assertTrue(rand <= max, "rand <= max");
function testRandomUintRangeOverflow() public {
vm.randomUint(0, uint256(int256(-1)));
}

function testRandomUint(uint256 val) public {
function testRandomUintSame(uint256 val) public {
uint256 rand = vm.randomUint(val, val);
assertTrue(rand == val);
}

function testRandomUintRange(uint256 min, uint256 max) public {
vm.assume(max >= min);
uint256 rand = vm.randomUint(min, max);
assertTrue(rand >= min, "rand >= min");
assertTrue(rand <= max, "rand <= max");
}

function testRandomAddress() public {
address rand = vm.randomAddress();
assertTrue(rand >= address(0));
vm.randomAddress();
}
}

0 comments on commit 56dbd20

Please sign in to comment.