disallow surrogates as actors

uniswapfoundation · Jan 26, 2024 · 7a6f026 · 7a6f026
1 parent fcb8cca
commit 7a6f026
Showing 1 changed file with 8 additions and 5 deletions.
diff --git a/test/helpers/UniStaker.handler.sol b/test/helpers/UniStaker.handler.sol
@@ -24,11 +24,13 @@ contract UniStakerHandler is CommonBase, StdCheats, StdUtils {
 
   AddressSet internal _delegates;
   AddressSet internal _beneficiaries;
+  AddressSet internal _surrogates;
 
-  modifier createActor() {
+  function _createActor() internal {
     currentActor = msg.sender;
+    // Surrogates can't stake. We won't include them as potential actors.
+    vm.assume(!_surrogates.contains(currentActor));
     _actors.add(msg.sender);
-    _;
   }
 
   modifier useActor(uint256 actorIndexSeed) {
@@ -48,14 +50,15 @@ contract UniStakerHandler is CommonBase, StdCheats, StdUtils {
 
   function _mintStakeToken(address _to, uint256 _amount) internal {
     vm.assume(_to != address(0));
-    deal(address(stakeToken), _to, _amount);
+    deal(address(stakeToken), _to, _amount, true);
   }
 
   function stake(uint256 _amount, address _delegatee, address _beneficiary)
     public
-    createActor
     countCall("stake")
   {
+    _createActor();
+
     // TODO: decide if we want reverts in stake
     //_beneficiary = address(uint160(bound(uint160(_beneficiary), 1, type(uint160).max)));
     //_delegatee = address(uint160(bound(uint160(_delegatee), 1, type(uint160).max)));
@@ -72,7 +75,7 @@ contract UniStakerHandler is CommonBase, StdCheats, StdUtils {
     stakeToken.approve(address(uniStaker), _amount);
     uniStaker.stake(_amount, _delegatee, _beneficiary);
     vm.stopPrank();
-
+    _surrogates.add(address(uniStaker.surrogates(_delegatee)));
     ghost_stakeSum += _amount;
   }