diff --git a/test/UniStaker.t.sol b/test/UniStaker.t.sol index bc286c9..00ccfd6 100644 --- a/test/UniStaker.t.sol +++ b/test/UniStaker.t.sol @@ -778,16 +778,28 @@ contract StakeOnBehalf is UniStakerTest { bytes32 _messageHash = keccak256(abi.encodePacked("\x19\x01", EIP712_DOMAIN_SEPARATOR, _message)); - bytes memory _signature = _sign(_privateKey, _messageHash); - if (_amount % 3 == 0) _amount++; - else if (_amount % 3 == 1) _delegatee = address(uint160(bytes20(_delegatee)) + 1); - else _depositor = address(uint160(bytes20(_depositor)) + 1); + // Here we use `_amount` as an arbitrary source of randomness to replace a legit parameter + // with an attack-like one. + if (_amount % 4 == 0) _amount++; + else if (_amount % 4 == 1) _delegatee = vm.addr(uint256(uint160(bytes20(_delegatee)))); + else if (_amount % 4 == 2) _depositor = vm.addr(uint256(uint160(bytes20(_depositor)))); + else if (_amount % 4 == 3) _messageHash = _modifyMessage(_messageHash, _amount % _privateKey); + bytes memory _signature = _sign(_privateKey, _messageHash); vm.prank(_sender); vm.expectRevert(UniStaker.UniStaker__InvalidSignature.selector); uniStaker.stakeOnBehalf(_amount, _delegatee, _depositor, _signature); } + + function _modifyMessage(bytes32 _message, uint256 _index) internal pure returns (bytes32) { + _index = bound(_index, 0, 31); + bytes memory _messageBytes = abi.encodePacked(_message); + // zero out the byte at the given index, or set it to 1 if it's already zero + if (_messageBytes[_index] == 0) _messageBytes[_index] = bytes1(uint8(1)); + else _messageBytes[_index] = bytes1(uint8(0)); + return bytes32(_messageBytes); + } } contract StakeMore is UniStakerTest {