From 20797c512a38a86577da5f541dc08d7bd70de372 Mon Sep 17 00:00:00 2001 From: wildmolasses Date: Thu, 25 Jan 2024 12:29:32 -0500 Subject: [PATCH] wip: Begin invariant test suite Co-authored-by: Alexander Keating --- foundry.toml | 9 +++ lib/forge-std | 2 +- test/UniStaker.invariants.t.sol | 67 ++++++++++++++++ test/helpers/AddressSet.sol | 46 +++++++++++ test/helpers/UniStaker.handler.sol | 120 +++++++++++++++++++++++++++++ 5 files changed, 243 insertions(+), 1 deletion(-) 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 6b1fb8f..3b4ef5d 100644 --- a/foundry.toml +++ b/foundry.toml @@ -34,3 +34,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 = 15 + dictionary_weight = 80 + fail_on_revert = false + include_push_bytes = true + include_storage = true + runs = 100 diff --git a/lib/forge-std b/lib/forge-std index 76e89e5..4513bc2 160000 --- a/lib/forge-std +++ b/lib/forge-std @@ -1 +1 @@ -Subproject commit 76e89e592ac58df4f0d16202ee230bbf0b8098fc +Subproject commit 4513bc2063f23c57bee6558799584b518d387a39 diff --git a/test/UniStaker.invariants.t.sol b/test/UniStaker.invariants.t.sol new file mode 100644 index 0000000..73cfc06 --- /dev/null +++ b/test/UniStaker.invariants.t.sol @@ -0,0 +1,67 @@ +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[](1); + selectors[0] = UniStakerHandler.stake.selector; + + targetSelector(FuzzSelector({addr: address(handler), selectors: selectors})); + + targetContract(address(handler)); + } + + function invariant_Sum_of_beneficiary_earning_power_equals_total_stake() public { + assertEq(uniStaker.totalSupply(), handler.reduceBeneficiaries(0, this.accumulateEarningPower)); + } + + function invariant_Sum_of_surrogate_balance_equals_total_stake() public { + assertEq(uniStaker.totalSupply(), handler.reduceDelegates(0, this.accumulateSurrogateBalance)); + } + + function accumulateEarningPower(uint256 earningPower, address caller) + external + view + returns (uint256) + { + return earningPower + uniStaker.earningPower(caller); + } + + 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); + } + + function invariant_callSummary() public view { + handler.callSummary(); + } +} 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..1079b7e --- /dev/null +++ b/test/helpers/UniStaker.handler.sol @@ -0,0 +1,120 @@ +// 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; + + UniStaker public uniStaker; + IERC20 public stakeToken; + + uint256 public ghost_stakeSum; + + mapping(bytes32 => uint256) public calls; + + AddressSet internal _actors; + address internal currentActor; + + AddressSet internal _delegates; + AddressSet internal _beneficiaries; + + modifier createActor() { + currentActor = msg.sender; + _actors.add(msg.sender); + _; + } + + modifier useActor(uint256 actorIndexSeed) { + currentActor = _actors.rand(actorIndexSeed); + _; + } + + modifier countCall(bytes32 key) { + calls[key]++; + _; + } + + constructor(UniStaker _uniStaker) { + uniStaker = _uniStaker; + stakeToken = IERC20(address(_uniStaker.STAKE_TOKEN())); + } + + function _mintStakeToken(address _to, uint256 _amount) internal { + vm.assume(_to != address(0)); + deal(address(stakeToken), _to, _amount); + } + + function stake(uint256 _amount, address _delegatee, address _beneficiary) + public + createActor + countCall("stake") + { + // 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))); + + _beneficiaries.add(_beneficiary); + _delegates.add(_delegatee); + // todo: adjust upper bound + _amount = bound(_amount, 0, 100_000_000e18); + + // assumes user has stake amount + _mintStakeToken(currentActor, _amount); + + vm.startPrank(currentActor); + stakeToken.approve(address(uniStaker), _amount); + uniStaker.stake(_amount, _delegatee, _beneficiary); + vm.stopPrank(); + + ghost_stakeSum += _amount; + } + + function _getBeneficiaryEarningPower(address _beneficiary) internal view returns (uint256) { + return uniStaker.earningPower(_beneficiary); + } + + function forEachActor(function(address) external func) public { + return _actors.forEach(func); + } + + function reduceActors(uint256 acc, function(uint256,address) external returns (uint256) func) + public + returns (uint256) + { + return _actors.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 actors() external view returns (address[] memory) { + return _actors.addrs; + } + + function callSummary() external view { + console.log("Call summary:"); + console.log("-------------------"); + console.log("stake", calls["stake"]); + console.log("-------------------"); + } + + receive() external payable {} +}