improve invalid signature test

uniswapfoundation · Feb 5, 2024 · 18fc9bb · 18fc9bb
1 parent e50aca5
commit 18fc9bb
Showing 1 changed file with 16 additions and 4 deletions.
diff --git 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 {