diff --git a/src/Staking.sol b/src/Staking.sol index 9415fc1..94fa337 100644 --- a/src/Staking.sol +++ b/src/Staking.sol @@ -13,7 +13,7 @@ import {IRewardsDistributor} from "./interfaces/IRewardsDistributor.sol"; // TODO should be pausable? // TODO is this vulnerable to first deposit attack? // TODO check calculations -contract StakingV2 is ERC20VotesUpgradeable, Ownable2StepUpgradeable { +contract Staking is ERC20VotesUpgradeable, Ownable2StepUpgradeable { /*////////////////////////////////////////////////////////////// LIBRARIES //////////////////////////////////////////////////////////////*/ diff --git a/test/unit/StakingUnitTestV2.t.sol b/test/Staking.t.sol similarity index 96% rename from test/unit/StakingUnitTestV2.t.sol rename to test/Staking.t.sol index be8d29b..81c2705 100644 --- a/test/unit/StakingUnitTestV2.t.sol +++ b/test/Staking.t.sol @@ -5,11 +5,11 @@ import "@forge-std/Test.sol"; import {FixedPointMathLib} from "@solmate/utils/FixedPointMathLib.sol"; import {TransparentUpgradeableProxy, ITransparentUpgradeableProxy} from "@openzeppelin/contracts/proxy/transparent/TransparentUpgradeableProxy.sol"; - import {Staking} from "src/Staking.sol"; import {RewardsDistributor} from "src/RewardsDistributor.sol"; import {IRewardsDistributor} from "src/interfaces/IRewardsDistributor.sol"; import {MockGovToken} from "test/mocks/MockGovToken.sol"; +import {ProxyUtils} from "test/helper/ProxyUtils.sol"; contract StakingTest is Test { using FixedPointMathLib for uint256; @@ -106,7 +106,13 @@ contract StakingTest is Test { address _keyper, uint256 _amount ) internal returns (uint256 _depositId) { - vm.assume(_keyper != address(0)); + vm.assume( + _keyper != address(0) && + uint160(_keyper) > 0x100 && // ignore precompiled address + _keyper != address(this) && + _keyper != address(staking) && + _keyper != ProxyUtils.getAdminAddress(_keyper) + ); vm.startPrank(_keyper); govToken.approve(address(staking), _amount); @@ -159,7 +165,6 @@ contract Stake is StakingTest { _setKeyper(_depositor, true); _stake(_depositor, _amount); - assertEq(govToken.balanceOf(_depositor), 0, "Wrong balance"); assertEq( govToken.balanceOf(address(staking)), @@ -355,8 +360,6 @@ contract Stake is StakingTest { _mintGovToken(_depositor, _amount); _setKeyper(_depositor, true); - vm.assume(_depositor != address(0)); - uint256 _contractBalanceBefore = govToken.balanceOf(address(staking)); _stake(_depositor, _amount); @@ -435,7 +438,7 @@ contract Stake is StakingTest { _mintGovToken(_depositor, _amount1 + _amount2); _setKeyper(_depositor, true); - vm.assume(_depositor != address(0)); + vm.assume(_depositor != address(0) && _depositor != address(this)); uint256 depositIndex1 = _stake(_depositor, _amount1); uint256 depositIndex2 = _stake(_depositor, _amount2); diff --git a/test/helper/ProxyUtils.sol b/test/helper/ProxyUtils.sol new file mode 100644 index 0000000..c1906f0 --- /dev/null +++ b/test/helper/ProxyUtils.sol @@ -0,0 +1,16 @@ +// SPDX-License-Identifier: MIT +pragma solidity 0.8.25; + +import {Vm} from "@forge-std/Vm.sol"; +import {ERC1967Utils} from "@openzeppelin/contracts/proxy/ERC1967/ERC1967Utils.sol"; + +library ProxyUtils { + address constant CHEATCODE_ADDRESS = + 0x7109709ECfa91a80626fF3989D68f67F5b1DD12D; + Vm constant vm = Vm(CHEATCODE_ADDRESS); + + function getAdminAddress(address proxy) public view returns (address) { + bytes32 adminSlot = vm.load(proxy, ERC1967Utils.ADMIN_SLOT); + return address(uint160(uint256(adminSlot))); + } +}