Invariants

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

foundry.toml

@@ -37,3 +37,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 = 100

test/README.md

@@ -0,0 +1,64 @@
+# Invariant Suite
+
+The invariant suite is a collection of tests that ensure certain properties of the system are always true.
+
+## Invariants under test
+
+- [ ] the earned fees should equal the percentage of durations that have passed on the rewards
+- [ ] Withdrawals for individual depositors + stake should equal the total staked throughout the tests
+- [ ] The sum of all of the notified rewards should equal the rewards balance in the staking contract plus all claimed rewards
+
+## Invariant Handler
+
+The handler contract specifies the actions that should be taken in the black box of an invariant run. Included in the handler contract are actions such as:
+
+### 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
+- enable rewards notifier
+- notifyRewardAmount
+- all of the `onBehalf` methods
+- multicall
+
+### Invalid user actions
+
+- Staking without sufficient EC20 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.
+
+- transfer in arbitrary amount of STAKE_TOKEN
+- transfer in arbitrary amount of REWARD_TOKEN
+- transfer direct 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
+- withdraw with zero amount
+- multicall
+
+### Utility actions
+
+- `warpAhead`: warp the block timestamp ahead by a specified number of seconds.

test/UniStaker.invariants.t.sol

@@ -0,0 +1,113 @@
+// 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 {
+    // deploy UniStaker
+    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_Unclaimed_reward_LTE_total_rewards() public {
+    assertLe(
+      handler.reduceBeneficiaries(0, this.accumulateUnclaimedReward),
+      rewardToken.balanceOf(address(uniStaker))
+    );
+  }
+
+  // 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,46 @@
+// SPDX-License-Identifier: AGPL-3.0-or-later
+pragma solidity ^0.8.23;
+
+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;
+  }
+}