diff --git a/.gas-snapshot b/.gas-snapshot index 763d688..edd796b 100644 --- a/.gas-snapshot +++ b/.gas-snapshot @@ -1,7 +1,11 @@ +ExecuteAccountTest:testDeployment() (gas: 26400) +ExecuteAccountTest:test_RevertWhen_InvalidLimitEpoch() (gas: 974502) LeaveTest:testDeployment() (gas: 26172) +LeaveTest:test_RevertWhen_NoPendingMigration() (gas: 662520) LeaveTest:test_RevertWhen_SenderIsNotVault() (gas: 10562) -LockTest:testDeployment() (gas: 26172) -LockTest:test_RevertWhen_SenderIsNotVault() (gas: 10573) +LockTest:testDeployment() (gas: 26400) +LockTest:test_RevertWhen_DecreasingLockTime() (gas: 976972) +LockTest:test_RevertWhen_SenderIsNotVault() (gas: 10607) MigrateTest:testDeployment() (gas: 26172) MigrateTest:test_RevertWhen_SenderIsNotVault() (gas: 10629) StakeManagerTest:testDeployment() (gas: 26172) diff --git a/contracts/StakeManager.sol b/contracts/StakeManager.sol index 22e88ce..cfdd8df 100644 --- a/contracts/StakeManager.sol +++ b/contracts/StakeManager.sol @@ -9,6 +9,11 @@ import { StakeVault } from "./StakeVault.sol"; contract StakeManager is Ownable { error StakeManager__SenderIsNotVault(); error StakeManager__FundsLocked(); + error StakeManager__DecreasingLockTime(); + error StakeManager__NoPendingMigration(); + error StakeManager__PendingMigration(); + error StakeManager__SenderIsNotPreviousStakeManager(); + error StakeManager__InvalidLimitEpoch(); struct Account { uint256 lockUntil; @@ -93,7 +98,9 @@ contract StakeManager is Ownable { function lock(uint256 _time) external onlyVault { Account storage account = accounts[msg.sender]; processAccount(account, currentEpoch); - require(block.timestamp + _time > account.lockUntil, "Cannot decrease lock time"); + if (block.timestamp + _time < account.lockUntil) { + revert StakeManager__DecreasingLockTime(); + } mintIntialMultiplier(account, _time, account.balance, 0); } @@ -101,7 +108,9 @@ contract StakeManager is Ownable { * @notice leave without processing account */ function leave() external onlyVault { - require(address(migration) != address(0), "Leave only during migration"); + if (address(migration) == address(0)) { + revert StakeManager__NoPendingMigration(); + } Account memory account = accounts[msg.sender]; delete accounts[msg.sender]; multiplierSupply -= account.multiplier; @@ -136,7 +145,9 @@ contract StakeManager is Ownable { */ function migrate() external onlyVault returns (StakeManager newManager) { - require(address(migration) != address(0), "Migration not available"); + if (address(migration) == address(0)) { + revert StakeManager__NoPendingMigration(); + } Account storage account = accounts[msg.sender]; stakedToken.approve(address(migration), account.balance); migration.migrate(msg.sender, account); @@ -151,7 +162,9 @@ contract StakeManager is Ownable { * @param _account Account data */ function migrate(address _vault, Account memory _account) external { - require(msg.sender == address(oldManager), "Unauthorized"); + if (msg.sender != address(oldManager)) { + revert StakeManager__SenderIsNotPreviousStakeManager(); + } stakedToken.transferFrom(address(oldManager), address(this), _account.balance); accounts[_vault] = _account; } @@ -190,8 +203,12 @@ contract StakeManager is Ownable { function processAccount(Account storage account, uint256 _limitEpoch) private { processEpoch(); - require(address(migration) == address(0), "Contract ended, please migrate"); - require(_limitEpoch <= currentEpoch, "Non-sense call"); + if (address(migration) != address(0)) { + revert StakeManager__PendingMigration(); + } + if (_limitEpoch > currentEpoch) { + revert StakeManager__InvalidLimitEpoch(); + } uint256 userReward; uint256 userEpoch = account.epoch; for (Epoch memory iEpoch = epochs[userEpoch]; userEpoch < _limitEpoch; userEpoch++) { diff --git a/test/StakeManager.t.sol b/test/StakeManager.t.sol index 5280872..0ce2768 100644 --- a/test/StakeManager.t.sol +++ b/test/StakeManager.t.sol @@ -89,6 +89,22 @@ contract LockTest is StakeManagerTest { vm.expectRevert(StakeManager.StakeManager__SenderIsNotVault.selector); stakeManager.lock(100); } + + function test_RevertWhen_DecreasingLockTime() public { + // ensure user has funds + deal(stakeToken, testUser, 1000); + StakeVault userVault = _createTestVault(testUser); + + vm.startPrank(testUser); + // ensure user vault can spend user tokens + ERC20(stakeToken).approve(address(userVault), 100); + + uint256 lockTime = 1 days; + userVault.stake(100, lockTime); + + vm.expectRevert(StakeManager.StakeManager__DecreasingLockTime.selector); + userVault.lock(1); + } } contract LeaveTest is StakeManagerTest { @@ -100,6 +116,13 @@ contract LeaveTest is StakeManagerTest { vm.expectRevert(StakeManager.StakeManager__SenderIsNotVault.selector); stakeManager.leave(); } + + function test_RevertWhen_NoPendingMigration() public { + StakeVault userVault = _createTestVault(testUser); + vm.prank(testUser); + vm.expectRevert(StakeManager.StakeManager__NoPendingMigration.selector); + userVault.leave(); + } } contract MigrateTest is StakeManagerTest { @@ -111,4 +134,33 @@ contract MigrateTest is StakeManagerTest { vm.expectRevert(StakeManager.StakeManager__SenderIsNotVault.selector); stakeManager.migrate(); } + + function test_RevertWhen_NoPendingMigration() public { + StakeVault userVault = _createTestVault(testUser); + vm.prank(testUser); + vm.expectRevert(StakeManager.StakeManager__NoPendingMigration.selector); + userVault.updateManager(); + } +} + +contract ExecuteAccountTest is StakeManagerTest { + function setUp() public override { + StakeManagerTest.setUp(); + } + + function test_RevertWhen_InvalidLimitEpoch() public { + deal(stakeToken, testUser, 1000); + StakeVault userVault = _createTestVault(testUser); + + vm.startPrank(testUser); + ERC20(stakeToken).approve(address(userVault), 100); + + uint256 lockTime = 1 days; + userVault.stake(100, lockTime); + + uint256 currentEpoch = stakeManager.currentEpoch(); + + vm.expectRevert(StakeManager.StakeManager__InvalidLimitEpoch.selector); + stakeManager.executeAccount(address(userVault), currentEpoch + 1); + } }