Invariant test suite

---------

Co-authored-by: Alexander Keating <[email protected]>

foundry.toml

@@ -41,3 +41,12 @@
   # with assume statements because we don't have the surrogate address until it's deployed later in
   # the test.
   include_storage = false
+
+[invariant]
+  call_override = false
+  depth = 50
+  dictionary_weight = 80
+  fail_on_revert = false
+  include_push_bytes = true
+  include_storage = true
+  runs = 256

test/README.md

@@ -0,0 +1,68 @@
+# Invariant Suite
+
+The invariant suite is a collection of tests designed to build confidence around certain properties of the system expected to be true.
+
+## Invariants under test
+
+- The total staked balance should equal the sum of all individual depositors' balances
+- The sum of beneficiary earning power should equal the total staked balance
+- The sum of all surrogate balance should equal the total staked balance
+- Cumulative deposits minus withdrawals should equal the total staked balance
+- The sum of all notified rewards should be greater or equal to all claimed rewards plus the rewards balance in the staking contract (TODO: not strictly equal because of stray transfers in, which are not yet implemented in handler)
+- Sum of unclaimed reward across all beneficiaries should be less than or equal to total rewards
+- `rewardPerTokenAccumulatedCheckpoint` should be greater or equal to the last `rewardPerTokenAccumulatedCheckpoint` value
+
+## Invariant Handler
+
+The handler contract specifies the actions that should be taken in the black box of an invariant run. Here is a list of implemented actions the handler contract can take, as well as ideas for further actions.
+
+### Valid user actions
+
+These actions are typical user actions that can be taken on the system. They are used to test the system's behavior under normal conditions.
+
+- [x] stake: a user deposits some amount of STAKE_TOKEN, specifying a delegatee and optionally a beneficiary.
+  - Action taken by: any user
+- [x] stakeMore: a user increments the balance on an existing deposit that she owns.
+  - Action taken by: existing depositors
+- [x] withdraw: a user withdraws some balance from a deposit that she owns.
+  - Action taken by: existing depositors
+- [x] claimReward: A beneficiary claims the reward that is due to her.
+  - Action taken by: existing beneficiaries
+- [ ] alterDelegatee
+- [ ] alterBeneficiary
+- [ ] permitAndStake
+- [x] enable rewards notifier
+- [x] notifyRewardAmount
+- [ ] all of the `onBehalf` methods
+- [ ] multicall
+
+### Invalid user actions
+
+- [ ] Staking without sufficient ERC20 approval
+- [ ] Stake more on a deposit that does not belong to you
+- [ ] State more on a deposit that does not exist
+- [ ] Alter beneficiary and alter delegatee on a deposit that is not yours or does not exist
+- [ ] withdraw on deposit that's not yours
+- [ ] call notifyRewardsAmount if you are not rewards notifier, or insufficient/incorrect reward balance
+- [ ] setAdmin and setRewardNotifier without being the admin
+- [ ] Invalid signature on the `onBehalf` methods
+- [ ] multicall
+
+### Weird user actions
+
+These are actions that are outside the normal use of the system. They are used to test the system's behavior under abnormal conditions.
+
+- [ ] directly transfer in some amount of STAKE_TOKEN to UniStaker
+- [ ] directly transfer some amount of REWARD_TOKEN to UniStaker
+- [ ] transfer stake directly to surrogate
+- [ ] reentrancy attempts
+- [ ] SELFDESTRUCT to this contract
+- [ ] flash loan?
+- [ ] User uses the staking contract as the from address in a `transferFrom`
+- [ ] A non-beneficiary calls claim reward
+- [x] withdraw with zero amount
+- [ ] multicall
+
+### Utility actions
+
+- [x] `warpAhead`: warp the block timestamp ahead by a specified number of seconds.

test/UniStaker.invariants.t.sol

@@ -0,0 +1,120 @@
+// SPDX-License-Identifier: AGPL-3.0-only
+pragma solidity ^0.8.23;
+
+import {Test} from "forge-std/Test.sol";
+import {IERC20} from "openzeppelin/token/ERC20/IERC20.sol";
+
+import {UniStaker} from "src/UniStaker.sol";
+import {UniStakerHandler} from "test/helpers/UniStaker.handler.sol";
+import {ERC20VotesMock} from "test/mocks/MockERC20Votes.sol";
+import {ERC20Fake} from "test/fakes/ERC20Fake.sol";
+
+contract UniStakerInvariants is Test {
+  UniStakerHandler public handler;
+  UniStaker public uniStaker;
+  ERC20Fake rewardToken;
+  ERC20VotesMock govToken;
+  address rewardsNotifier;
+
+  function setUp() public {
+    rewardToken = new ERC20Fake();
+    vm.label(address(rewardToken), "Rewards Token");
+
+    govToken = new ERC20VotesMock();
+    vm.label(address(govToken), "Governance Token");
+
+    rewardsNotifier = address(0xaffab1ebeef);
+    vm.label(rewardsNotifier, "Rewards Notifier");
+    uniStaker = new UniStaker(rewardToken, govToken, rewardsNotifier);
+    handler = new UniStakerHandler(uniStaker);
+
+    bytes4[] memory selectors = new bytes4[](7);
+    selectors[0] = UniStakerHandler.stake.selector;
+    selectors[1] = UniStakerHandler.validStakeMore.selector;
+    selectors[2] = UniStakerHandler.validWithdraw.selector;
+    selectors[3] = UniStakerHandler.warpAhead.selector;
+    selectors[4] = UniStakerHandler.claimReward.selector;
+    selectors[5] = UniStakerHandler.enableRewardNotifier.selector;
+    selectors[6] = UniStakerHandler.notifyRewardAmount.selector;
+
+    targetSelector(FuzzSelector({addr: address(handler), selectors: selectors}));
+
+    targetContract(address(handler));
+  }
+
+  // Invariants
+
+  function invariant_Sum_of_all_depositor_balances_equals_total_stake() public {
+    assertEq(uniStaker.totalStaked(), handler.reduceDepositors(0, this.accumulateDeposits));
+  }
+
+  function invariant_Sum_of_beneficiary_earning_power_equals_total_stake() public {
+    assertEq(uniStaker.totalStaked(), handler.reduceBeneficiaries(0, this.accumulateEarningPower));
+  }
+
+  function invariant_Sum_of_surrogate_balance_equals_total_stake() public {
+    assertEq(uniStaker.totalStaked(), handler.reduceDelegates(0, this.accumulateSurrogateBalance));
+  }
+
+  function invariant_Cumulative_staked_minus_withdrawals_equals_total_stake() public {
+    assertEq(uniStaker.totalStaked(), handler.ghost_stakeSum() - handler.ghost_stakeWithdrawn());
+  }
+
+  function invariant_Sum_of_notified_rewards_equals_all_claimed_rewards_plus_rewards_left() public {
+    assertEq(
+      handler.ghost_rewardsNotified(),
+      rewardToken.balanceOf(address(uniStaker)) + handler.ghost_rewardsClaimed()
+    );
+  }
+
+  function invariant_Sum_of_unclaimed_reward_should_be_less_than_or_equal_to_total_rewards() public {
+    assertLe(
+      handler.reduceBeneficiaries(0, this.accumulateUnclaimedReward),
+      rewardToken.balanceOf(address(uniStaker))
+    );
+  }
+
+  function invariant_RewardPerTokenAccumulatedCheckpoint_should_be_greater_or_equal_to_the_last_rewardPerTokenAccumulatedCheckpoint(
+  ) public {
+    assertGe(
+      uniStaker.rewardPerTokenAccumulatedCheckpoint(),
+      handler.ghost_prevRewardPerTokenAccumulatedCheckpoint()
+    );
+  }
+
+  // Used to see distribution of non-reverting calls
+  function invariant_callSummary() public view {
+    handler.callSummary();
+  }
+
+  // Helpers
+
+  function accumulateDeposits(uint256 balance, address depositor) external view returns (uint256) {
+    return balance + uniStaker.depositorTotalStaked(depositor);
+  }
+
+  function accumulateEarningPower(uint256 earningPower, address caller)
+    external
+    view
+    returns (uint256)
+  {
+    return earningPower + uniStaker.earningPower(caller);
+  }
+
+  function accumulateUnclaimedReward(uint256 unclaimedReward, address beneficiary)
+    external
+    view
+    returns (uint256)
+  {
+    return unclaimedReward + uniStaker.unclaimedReward(beneficiary);
+  }
+
+  function accumulateSurrogateBalance(uint256 balance, address delegate)
+    external
+    view
+    returns (uint256)
+  {
+    address surrogateAddr = address(uniStaker.surrogates(delegate));
+    return balance + IERC20(address(uniStaker.STAKE_TOKEN())).balanceOf(surrogateAddr);
+  }
+}

test/helpers/AddressSet.sol

@@ -0,0 +1,49 @@
+// SPDX-License-Identifier: AGPL-3.0-or-later
+pragma solidity ^0.8.23;
+
+// AddressSet.sol comes from
+// https://github.com/horsefacts/weth-invariant-testing/blob/973156bc9b6684f0cf62de19e9bb4c5c27a41bb2/test/helpers/AddressSet.sol
+
+struct AddressSet {
+  address[] addrs;
+  mapping(address => bool) saved;
+}
+
+library LibAddressSet {
+  function add(AddressSet storage s, address addr) internal {
+    if (!s.saved[addr]) {
+      s.addrs.push(addr);
+      s.saved[addr] = true;
+    }
+  }
+
+  function contains(AddressSet storage s, address addr) internal view returns (bool) {
+    return s.saved[addr];
+  }
+
+  function count(AddressSet storage s) internal view returns (uint256) {
+    return s.addrs.length;
+  }
+
+  function rand(AddressSet storage s, uint256 seed) internal view returns (address) {
+    if (s.addrs.length > 0) return s.addrs[seed % s.addrs.length];
+    else return address(0);
+  }
+
+  function forEach(AddressSet storage s, function(address) external func) internal {
+    for (uint256 i; i < s.addrs.length; ++i) {
+      func(s.addrs[i]);
+    }
+  }
+
+  function reduce(
+    AddressSet storage s,
+    uint256 acc,
+    function(uint256,address) external returns (uint256) func
+  ) internal returns (uint256) {
+    for (uint256 i; i < s.addrs.length; ++i) {
+      acc = func(acc, s.addrs[i]);
+    }
+    return acc;
+  }
+}