From 71da1208dd923782423784acf32612655db3ccf9 Mon Sep 17 00:00:00 2001 From: wildmolasses Date: Thu, 25 Jan 2024 12:29:32 -0500 Subject: [PATCH 1/6] Invariants Co-authored-by: Alexander Keating --- foundry.toml | 9 ++ test/README.md | 64 +++++++++ test/UniStaker.invariants.t.sol | 113 ++++++++++++++++ test/helpers/AddressSet.sol | 46 +++++++ test/helpers/UniStaker.handler.sol | 209 +++++++++++++++++++++++++++++ 5 files changed, 441 insertions(+) create mode 100644 test/README.md create mode 100644 test/UniStaker.invariants.t.sol create mode 100644 test/helpers/AddressSet.sol create mode 100644 test/helpers/UniStaker.handler.sol diff --git a/foundry.toml b/foundry.toml index 97b7b1d..fd8d835 100644 --- a/foundry.toml +++ b/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 = 100 diff --git a/test/README.md b/test/README.md new file mode 100644 index 0000000..f51fecd --- /dev/null +++ b/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. diff --git a/test/UniStaker.invariants.t.sol b/test/UniStaker.invariants.t.sol new file mode 100644 index 0000000..a09ed34 --- /dev/null +++ b/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); + } +} diff --git a/test/helpers/AddressSet.sol b/test/helpers/AddressSet.sol new file mode 100644 index 0000000..e57f494 --- /dev/null +++ b/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; + } +} diff --git a/test/helpers/UniStaker.handler.sol b/test/helpers/UniStaker.handler.sol new file mode 100644 index 0000000..bc405f5 --- /dev/null +++ b/test/helpers/UniStaker.handler.sol @@ -0,0 +1,209 @@ +// SPDX-License-Identifier: GPL-3.0-or-later +pragma solidity ^0.8.13; + +import {CommonBase} from "forge-std/Base.sol"; +import {StdCheats} from "forge-std/StdCheats.sol"; +import {StdUtils} from "forge-std/StdUtils.sol"; +import {console} from "forge-std/console.sol"; +import {AddressSet, LibAddressSet} from "../helpers/AddressSet.sol"; +import {UniStaker} from "src/UniStaker.sol"; +import {IERC20} from "openzeppelin/token/ERC20/IERC20.sol"; + +contract UniStakerHandler is CommonBase, StdCheats, StdUtils { + using LibAddressSet for AddressSet; + + // system setup + UniStaker public uniStaker; + IERC20 public stakeToken; + IERC20 public rewardToken; + address public admin; + + // actors, deposit state + address internal currentActor; + AddressSet internal _depositors; + AddressSet internal _delegates; + AddressSet internal _beneficiaries; + AddressSet internal _surrogates; + AddressSet internal _rewardNotifiers; + mapping(address => uint256[]) internal _depositIds; + mapping(bytes32 => uint256) public calls; + + // ghost vars + uint256 public ghost_stakeSum; + uint256 public ghost_stakeWithdrawn; + uint256 public ghost_depositCount; + uint256 public ghost_rewardsClaimed; + uint256 public ghost_rewardsNotified; + + modifier countCall(bytes32 key) { + calls[key]++; + _; + } + + constructor(UniStaker _uniStaker) { + uniStaker = _uniStaker; + stakeToken = IERC20(address(_uniStaker.STAKE_TOKEN())); + rewardToken = IERC20(address(_uniStaker.REWARD_TOKEN())); + admin = uniStaker.admin(); + } + + function _mintStakeToken(address _to, uint256 _amount) internal { + vm.assume(_to != address(0)); + deal(address(stakeToken), _to, _amount, true); + } + + function _mintRewardToken(address _to, uint256 _amount) internal { + vm.assume(_to != address(0)); + deal(address(rewardToken), _to, _amount, true); + } + + function enableRewardNotifier(address _notifier) public countCall("enableRewardNotifier") { + vm.assume(_notifier != address(0)); + _rewardNotifiers.add(_notifier); + vm.prank(admin); + uniStaker.setRewardNotifier(_notifier, true); + } + + function notifyRewardAmount(uint256 _amount, uint256 _actorSeed) + public + countCall("notifyRewardAmount") + { + _useActor(_rewardNotifiers, _actorSeed); + vm.assume(currentActor != address(0)); + _amount = bound(_amount, 0, 100_000_000e18); + _mintRewardToken(currentActor, _amount); + vm.startPrank(currentActor); + rewardToken.transfer(address(uniStaker), _amount); + uniStaker.notifyRewardAmount(_amount); + vm.stopPrank(); + ghost_rewardsNotified += _amount; + } + + // TODO: distinguish between valid and invalid stake + function stake(uint256 _amount, address _delegatee, address _beneficiary) + public + countCall("stake") + { + _createDepositor(); + + _beneficiaries.add(_beneficiary); + _delegates.add(_delegatee); + // todo: adjust upper bound + _amount = bound(_amount, 0, 100_000_000e18); + + // assume user has stake amount + _mintStakeToken(currentActor, _amount); + + vm.startPrank(currentActor); + stakeToken.approve(address(uniStaker), _amount); + uniStaker.stake(_amount, _delegatee, _beneficiary); + vm.stopPrank(); + + // update handler state + _depositIds[currentActor].push(ghost_depositCount); + ghost_depositCount++; + _surrogates.add(address(uniStaker.surrogates(_delegatee))); + ghost_stakeSum += _amount; + } + + function validStakeMore(uint256 _amount, uint256 _actorSeed, uint256 _actorDepositSeed) + public + countCall("validStakeMore") + { + _useActor(_depositors, _actorSeed); + vm.assume(currentActor != address(0)); + vm.assume(_depositIds[currentActor].length > 0); + UniStaker.DepositIdentifier _depositId = + UniStaker.DepositIdentifier.wrap(_getActorRandDepositId(_actorDepositSeed)); + (uint256 _balance,,,) = uniStaker.deposits(_depositId); + _amount = bound(_amount, 0, _balance); + vm.startPrank(currentActor); + stakeToken.approve(address(uniStaker), _amount); + uniStaker.stakeMore(_depositId, _amount); + vm.stopPrank(); + ghost_stakeSum += _amount; + } + + // TODO: include invalid withdrawals + function validWithdraw(uint256 _amount, uint256 _actorSeed, uint256 _actorDepositSeed) + public + countCall("validWithdraw") + { + _useActor(_depositors, _actorSeed); + vm.assume(currentActor != address(0)); + vm.assume(_depositIds[currentActor].length > 0); + UniStaker.DepositIdentifier _depositId = + UniStaker.DepositIdentifier.wrap(_getActorRandDepositId(_actorDepositSeed)); + (uint256 _balance,,,) = uniStaker.deposits(_depositId); + _amount = bound(_amount, 0, _balance); + vm.startPrank(currentActor); + uniStaker.withdraw(_depositId, _amount); + vm.stopPrank(); + ghost_stakeWithdrawn += _amount; + } + + function claimReward(uint256 _actorSeed) public countCall("claimReward") { + _useActor(_beneficiaries, _actorSeed); + vm.startPrank(currentActor); + uint256 rewardsClaimed = uniStaker.unclaimedRewardCheckpoint(currentActor); + uniStaker.claimReward(); + vm.stopPrank(); + ghost_rewardsClaimed += rewardsClaimed; + } + + function warpAhead(uint256 _seconds) public countCall("warpAhead") { + _seconds = bound(_seconds, 0, uniStaker.REWARD_DURATION() * 2); + skip(_seconds); + } + + function _getActorRandDepositId(uint256 _randomDepositSeed) internal view returns (uint256) { + return _depositIds[currentActor][_randomDepositSeed % _depositIds[currentActor].length]; + } + + function _createDepositor() internal { + currentActor = msg.sender; + // Surrogates can't stake. We won't include them as potential depositors. + vm.assume(!_surrogates.contains(currentActor)); + _depositors.add(msg.sender); + } + + function _useActor(AddressSet storage _set, uint256 _randomActorSeed) internal { + currentActor = _set.rand(_randomActorSeed); + } + + function reduceDepositors(uint256 acc, function(uint256,address) external returns (uint256) func) + public + returns (uint256) + { + return _depositors.reduce(acc, func); + } + + function reduceBeneficiaries( + uint256 acc, + function(uint256,address) external returns (uint256) func + ) public returns (uint256) { + return _beneficiaries.reduce(acc, func); + } + + function reduceDelegates(uint256 acc, function(uint256,address) external returns (uint256) func) + public + returns (uint256) + { + return _delegates.reduce(acc, func); + } + + function callSummary() external view { + console.log("\nCall summary:"); + console.log("-------------------"); + console.log("stake", calls["stake"]); + console.log("validStakeMore", calls["validStakeMore"]); + console.log("validWithdraw", calls["validWithdraw"]); + console.log("claimReward", calls["claimReward"]); + console.log("enableRewardNotifier", calls["enableRewardNotifier"]); + console.log("notifyRewardAmount", calls["notifyRewardAmount"]); + console.log("warpAhead", calls["warpAhead"]); + console.log("-------------------\n"); + } + + receive() external payable {} +} From c3939a30873ab637cfe39388a0d69d01982d4d6f Mon Sep 17 00:00:00 2001 From: wildmolasses Date: Wed, 21 Feb 2024 12:34:43 -0500 Subject: [PATCH 2/6] Readme work, etc --- test/README.md | 68 ++++++++++++++++-------------- test/UniStaker.invariants.t.sol | 10 ++++- test/helpers/UniStaker.handler.sol | 25 +++++++++-- 3 files changed, 67 insertions(+), 36 deletions(-) diff --git a/test/README.md b/test/README.md index f51fecd..24daad7 100644 --- a/test/README.md +++ b/test/README.md @@ -1,16 +1,20 @@ # Invariant Suite -The invariant suite is a collection of tests that ensure certain properties of the system are always true. +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 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 +- 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. Included in the handler contract are actions such as: +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 @@ -24,41 +28,41 @@ These actions are typical user actions that can be taken on the system. They are - 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 +- [ ] alterDelegatee +- [ ] alterBeneficiary +- [ ] permitAndStake +- [x] enable rewards notifier +- [x] 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 +- [ ] 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. -- 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 +- [ ] 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 -- `warpAhead`: warp the block timestamp ahead by a specified number of seconds. +- [x] `warpAhead`: warp the block timestamp ahead by a specified number of seconds. diff --git a/test/UniStaker.invariants.t.sol b/test/UniStaker.invariants.t.sol index a09ed34..25dfa7c 100644 --- a/test/UniStaker.invariants.t.sol +++ b/test/UniStaker.invariants.t.sol @@ -68,13 +68,21 @@ contract UniStakerInvariants is Test { ); } - function invariant_Unclaimed_reward_LTE_total_rewards() public { + 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(); diff --git a/test/helpers/UniStaker.handler.sol b/test/helpers/UniStaker.handler.sol index bc405f5..307e2a4 100644 --- a/test/helpers/UniStaker.handler.sol +++ b/test/helpers/UniStaker.handler.sol @@ -34,12 +34,18 @@ contract UniStakerHandler is CommonBase, StdCheats, StdUtils { uint256 public ghost_depositCount; uint256 public ghost_rewardsClaimed; uint256 public ghost_rewardsNotified; + uint256 public ghost_prevRewardPerTokenAccumulatedCheckpoint; modifier countCall(bytes32 key) { calls[key]++; _; } + modifier doCheckpoints() { + _checkpoint_ghost_prevRewardPerTokenAccumulatedCheckpoint(); + _; + } + constructor(UniStaker _uniStaker) { uniStaker = _uniStaker; stakeToken = IERC20(address(_uniStaker.STAKE_TOKEN())); @@ -57,7 +63,11 @@ contract UniStakerHandler is CommonBase, StdCheats, StdUtils { deal(address(rewardToken), _to, _amount, true); } - function enableRewardNotifier(address _notifier) public countCall("enableRewardNotifier") { + function enableRewardNotifier(address _notifier) + public + countCall("enableRewardNotifier") + doCheckpoints + { vm.assume(_notifier != address(0)); _rewardNotifiers.add(_notifier); vm.prank(admin); @@ -67,10 +77,12 @@ contract UniStakerHandler is CommonBase, StdCheats, StdUtils { function notifyRewardAmount(uint256 _amount, uint256 _actorSeed) public countCall("notifyRewardAmount") + doCheckpoints { _useActor(_rewardNotifiers, _actorSeed); vm.assume(currentActor != address(0)); _amount = bound(_amount, 0, 100_000_000e18); + ghost_prevRewardPerTokenAccumulatedCheckpoint = uniStaker.rewardPerTokenAccumulatedCheckpoint(); _mintRewardToken(currentActor, _amount); vm.startPrank(currentActor); rewardToken.transfer(address(uniStaker), _amount); @@ -83,6 +95,7 @@ contract UniStakerHandler is CommonBase, StdCheats, StdUtils { function stake(uint256 _amount, address _delegatee, address _beneficiary) public countCall("stake") + doCheckpoints { _createDepositor(); @@ -109,6 +122,7 @@ contract UniStakerHandler is CommonBase, StdCheats, StdUtils { function validStakeMore(uint256 _amount, uint256 _actorSeed, uint256 _actorDepositSeed) public countCall("validStakeMore") + doCheckpoints { _useActor(_depositors, _actorSeed); vm.assume(currentActor != address(0)); @@ -128,6 +142,7 @@ contract UniStakerHandler is CommonBase, StdCheats, StdUtils { function validWithdraw(uint256 _amount, uint256 _actorSeed, uint256 _actorDepositSeed) public countCall("validWithdraw") + doCheckpoints { _useActor(_depositors, _actorSeed); vm.assume(currentActor != address(0)); @@ -142,7 +157,7 @@ contract UniStakerHandler is CommonBase, StdCheats, StdUtils { ghost_stakeWithdrawn += _amount; } - function claimReward(uint256 _actorSeed) public countCall("claimReward") { + function claimReward(uint256 _actorSeed) public countCall("claimReward") doCheckpoints { _useActor(_beneficiaries, _actorSeed); vm.startPrank(currentActor); uint256 rewardsClaimed = uniStaker.unclaimedRewardCheckpoint(currentActor); @@ -151,7 +166,7 @@ contract UniStakerHandler is CommonBase, StdCheats, StdUtils { ghost_rewardsClaimed += rewardsClaimed; } - function warpAhead(uint256 _seconds) public countCall("warpAhead") { + function warpAhead(uint256 _seconds) public countCall("warpAhead") doCheckpoints { _seconds = bound(_seconds, 0, uniStaker.REWARD_DURATION() * 2); skip(_seconds); } @@ -192,6 +207,10 @@ contract UniStakerHandler is CommonBase, StdCheats, StdUtils { return _delegates.reduce(acc, func); } + function _checkpoint_ghost_prevRewardPerTokenAccumulatedCheckpoint() internal { + ghost_prevRewardPerTokenAccumulatedCheckpoint = uniStaker.rewardPerTokenAccumulatedCheckpoint(); + } + function callSummary() external view { console.log("\nCall summary:"); console.log("-------------------"); From 163fedbd9ce97e2b354b9c85d1945358954dc4cb Mon Sep 17 00:00:00 2001 From: wildmolasses Date: Thu, 22 Feb 2024 10:40:46 -0500 Subject: [PATCH 3/6] credit horsefacts --- test/helpers/AddressSet.sol | 3 +++ test/helpers/UniStaker.handler.sol | 38 +++++++++++++++--------------- 2 files changed, 22 insertions(+), 19 deletions(-) diff --git a/test/helpers/AddressSet.sol b/test/helpers/AddressSet.sol index e57f494..83327a7 100644 --- a/test/helpers/AddressSet.sol +++ b/test/helpers/AddressSet.sol @@ -1,6 +1,9 @@ // 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; diff --git a/test/helpers/UniStaker.handler.sol b/test/helpers/UniStaker.handler.sol index 307e2a4..87d33c0 100644 --- a/test/helpers/UniStaker.handler.sol +++ b/test/helpers/UniStaker.handler.sol @@ -19,7 +19,7 @@ contract UniStakerHandler is CommonBase, StdCheats, StdUtils { address public admin; // actors, deposit state - address internal currentActor; + address internal _currentActor; AddressSet internal _depositors; AddressSet internal _delegates; AddressSet internal _beneficiaries; @@ -80,11 +80,11 @@ contract UniStakerHandler is CommonBase, StdCheats, StdUtils { doCheckpoints { _useActor(_rewardNotifiers, _actorSeed); - vm.assume(currentActor != address(0)); + vm.assume(_currentActor != address(0)); _amount = bound(_amount, 0, 100_000_000e18); ghost_prevRewardPerTokenAccumulatedCheckpoint = uniStaker.rewardPerTokenAccumulatedCheckpoint(); - _mintRewardToken(currentActor, _amount); - vm.startPrank(currentActor); + _mintRewardToken(_currentActor, _amount); + vm.startPrank(_currentActor); rewardToken.transfer(address(uniStaker), _amount); uniStaker.notifyRewardAmount(_amount); vm.stopPrank(); @@ -105,15 +105,15 @@ contract UniStakerHandler is CommonBase, StdCheats, StdUtils { _amount = bound(_amount, 0, 100_000_000e18); // assume user has stake amount - _mintStakeToken(currentActor, _amount); + _mintStakeToken(_currentActor, _amount); - vm.startPrank(currentActor); + vm.startPrank(_currentActor); stakeToken.approve(address(uniStaker), _amount); uniStaker.stake(_amount, _delegatee, _beneficiary); vm.stopPrank(); // update handler state - _depositIds[currentActor].push(ghost_depositCount); + _depositIds[_currentActor].push(ghost_depositCount); ghost_depositCount++; _surrogates.add(address(uniStaker.surrogates(_delegatee))); ghost_stakeSum += _amount; @@ -125,13 +125,13 @@ contract UniStakerHandler is CommonBase, StdCheats, StdUtils { doCheckpoints { _useActor(_depositors, _actorSeed); - vm.assume(currentActor != address(0)); - vm.assume(_depositIds[currentActor].length > 0); + vm.assume(_currentActor != address(0)); + vm.assume(_depositIds[_currentActor].length > 0); UniStaker.DepositIdentifier _depositId = UniStaker.DepositIdentifier.wrap(_getActorRandDepositId(_actorDepositSeed)); (uint256 _balance,,,) = uniStaker.deposits(_depositId); _amount = bound(_amount, 0, _balance); - vm.startPrank(currentActor); + vm.startPrank(_currentActor); stakeToken.approve(address(uniStaker), _amount); uniStaker.stakeMore(_depositId, _amount); vm.stopPrank(); @@ -145,13 +145,13 @@ contract UniStakerHandler is CommonBase, StdCheats, StdUtils { doCheckpoints { _useActor(_depositors, _actorSeed); - vm.assume(currentActor != address(0)); - vm.assume(_depositIds[currentActor].length > 0); + vm.assume(_currentActor != address(0)); + vm.assume(_depositIds[_currentActor].length > 0); UniStaker.DepositIdentifier _depositId = UniStaker.DepositIdentifier.wrap(_getActorRandDepositId(_actorDepositSeed)); (uint256 _balance,,,) = uniStaker.deposits(_depositId); _amount = bound(_amount, 0, _balance); - vm.startPrank(currentActor); + vm.startPrank(_currentActor); uniStaker.withdraw(_depositId, _amount); vm.stopPrank(); ghost_stakeWithdrawn += _amount; @@ -159,8 +159,8 @@ contract UniStakerHandler is CommonBase, StdCheats, StdUtils { function claimReward(uint256 _actorSeed) public countCall("claimReward") doCheckpoints { _useActor(_beneficiaries, _actorSeed); - vm.startPrank(currentActor); - uint256 rewardsClaimed = uniStaker.unclaimedRewardCheckpoint(currentActor); + vm.startPrank(_currentActor); + uint256 rewardsClaimed = uniStaker.unclaimedRewardCheckpoint(_currentActor); uniStaker.claimReward(); vm.stopPrank(); ghost_rewardsClaimed += rewardsClaimed; @@ -172,18 +172,18 @@ contract UniStakerHandler is CommonBase, StdCheats, StdUtils { } function _getActorRandDepositId(uint256 _randomDepositSeed) internal view returns (uint256) { - return _depositIds[currentActor][_randomDepositSeed % _depositIds[currentActor].length]; + return _depositIds[_currentActor][_randomDepositSeed % _depositIds[_currentActor].length]; } function _createDepositor() internal { - currentActor = msg.sender; + _currentActor = msg.sender; // Surrogates can't stake. We won't include them as potential depositors. - vm.assume(!_surrogates.contains(currentActor)); + vm.assume(!_surrogates.contains(_currentActor)); _depositors.add(msg.sender); } function _useActor(AddressSet storage _set, uint256 _randomActorSeed) internal { - currentActor = _set.rand(_randomActorSeed); + _currentActor = _set.rand(_randomActorSeed); } function reduceDepositors(uint256 acc, function(uint256,address) external returns (uint256) func) From e6928b5d5367d4ca9285eb6be52f1adcc7f82e8d Mon Sep 17 00:00:00 2001 From: wildmolasses Date: Thu, 22 Feb 2024 10:47:41 -0500 Subject: [PATCH 4/6] feedback --- test/UniStaker.invariants.t.sol | 1 - 1 file changed, 1 deletion(-) diff --git a/test/UniStaker.invariants.t.sol b/test/UniStaker.invariants.t.sol index 25dfa7c..4c80ce1 100644 --- a/test/UniStaker.invariants.t.sol +++ b/test/UniStaker.invariants.t.sol @@ -17,7 +17,6 @@ contract UniStakerInvariants is Test { address rewardsNotifier; function setUp() public { - // deploy UniStaker rewardToken = new ERC20Fake(); vm.label(address(rewardToken), "Rewards Token"); From 38d893a38154c5953a7e09cb096ae0687409ec79 Mon Sep 17 00:00:00 2001 From: wildmolasses Date: Thu, 22 Feb 2024 11:34:24 -0500 Subject: [PATCH 5/6] use foundry default runs amount in foundry.toml --- foundry.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/foundry.toml b/foundry.toml index fd8d835..7292500 100644 --- a/foundry.toml +++ b/foundry.toml @@ -49,4 +49,4 @@ fail_on_revert = false include_push_bytes = true include_storage = true - runs = 100 + runs = 256 From c6e46f41d556329723e9fce3a81a90870b308383 Mon Sep 17 00:00:00 2001 From: wildmolasses Date: Thu, 22 Feb 2024 11:35:22 -0500 Subject: [PATCH 6/6] remove TODOs --- test/helpers/UniStaker.handler.sol | 3 --- 1 file changed, 3 deletions(-) diff --git a/test/helpers/UniStaker.handler.sol b/test/helpers/UniStaker.handler.sol index 87d33c0..f8fe335 100644 --- a/test/helpers/UniStaker.handler.sol +++ b/test/helpers/UniStaker.handler.sol @@ -91,7 +91,6 @@ contract UniStakerHandler is CommonBase, StdCheats, StdUtils { ghost_rewardsNotified += _amount; } - // TODO: distinguish between valid and invalid stake function stake(uint256 _amount, address _delegatee, address _beneficiary) public countCall("stake") @@ -101,7 +100,6 @@ contract UniStakerHandler is CommonBase, StdCheats, StdUtils { _beneficiaries.add(_beneficiary); _delegates.add(_delegatee); - // todo: adjust upper bound _amount = bound(_amount, 0, 100_000_000e18); // assume user has stake amount @@ -138,7 +136,6 @@ contract UniStakerHandler is CommonBase, StdCheats, StdUtils { ghost_stakeSum += _amount; } - // TODO: include invalid withdrawals function validWithdraw(uint256 _amount, uint256 _actorSeed, uint256 _actorDepositSeed) public countCall("validWithdraw")