fix(StakeVault): make unstaking actually work

Unstaking didn't actually work because it was using `transferFrom()` on the
`StakeVault` with the `from` address being the vault itself.
This would result in an approval error because the vault isn't creating
any approvals to spend its own funds.

The solution is to use `transfer` instead and ensuring the return value
is checked.
This commit is contained in:
r4bbit 2024-01-11 14:54:39 +01:00
parent edc44e0c6a
commit 74ff357142
5 changed files with 93 additions and 2 deletions

View File

@ -14,6 +14,10 @@ import { StakeManager } from "./StakeManager.sol";
contract StakeVault is Ownable { contract StakeVault is Ownable {
error StakeVault__MigrationNotAvailable(); error StakeVault__MigrationNotAvailable();
error StakeVault__StakingFailed();
error StakeVault__UnstakingFailed();
StakeManager private stakeManager; StakeManager private stakeManager;
ERC20 private immutable STAKED_TOKEN; ERC20 private immutable STAKED_TOKEN;
@ -27,7 +31,10 @@ contract StakeVault is Ownable {
} }
function stake(uint256 _amount, uint256 _time) external onlyOwner { function stake(uint256 _amount, uint256 _time) external onlyOwner {
STAKED_TOKEN.transferFrom(msg.sender, address(this), _amount); bool success = STAKED_TOKEN.transferFrom(msg.sender, address(this), _amount);
if (!success) {
revert StakeVault__StakingFailed();
}
stakeManager.stake(_amount, _time); stakeManager.stake(_amount, _time);
emit Staked(msg.sender, address(this), _amount, _time); emit Staked(msg.sender, address(this), _amount, _time);
@ -39,7 +46,10 @@ contract StakeVault is Ownable {
function unstake(uint256 _amount) external onlyOwner { function unstake(uint256 _amount) external onlyOwner {
stakeManager.unstake(_amount); stakeManager.unstake(_amount);
STAKED_TOKEN.transferFrom(address(this), msg.sender, _amount); bool success = STAKED_TOKEN.transfer(msg.sender, _amount);
if (!success) {
revert StakeVault__UnstakingFailed();
}
} }
function leave() external onlyOwner { function leave() external onlyOwner {

View File

@ -97,6 +97,24 @@ contract UnstakeTest is StakeManagerTest {
vm.expectRevert(StakeManager.StakeManager__FundsLocked.selector); vm.expectRevert(StakeManager.StakeManager__FundsLocked.selector);
userVault.unstake(100); userVault.unstake(100);
} }
function test_UnstakeShouldReturnFunds() public {
// ensure user has funds
deal(stakeToken, testUser, 1000);
StakeVault userVault = _createTestVault(testUser);
vm.startPrank(testUser);
ERC20(stakeToken).approve(address(userVault), 100);
userVault.stake(100, 0);
assertEq(ERC20(stakeToken).balanceOf(testUser), 900);
userVault.unstake(100);
assertEq(stakeManager.stakeSupply(), 0);
assertEq(ERC20(stakeToken).balanceOf(address(userVault)), 0);
assertEq(ERC20(stakeToken).balanceOf(testUser), 1000);
}
} }
contract LockTest is StakeManagerTest { contract LockTest is StakeManagerTest {

View File

@ -1,8 +1,11 @@
// SPDX-License-Identifier: UNLICENSED // SPDX-License-Identifier: UNLICENSED
pragma solidity ^0.8.19; pragma solidity ^0.8.19;
import { ERC20 } from "@openzeppelin/contracts/token/ERC20/ERC20.sol";
import { Test } from "forge-std/Test.sol"; import { Test } from "forge-std/Test.sol";
import { Deploy } from "../script/Deploy.s.sol"; import { Deploy } from "../script/Deploy.s.sol";
import { DeployBroken } from "./script/DeployBroken.s.sol";
import { DeploymentConfig } from "../script/DeploymentConfig.s.sol"; import { DeploymentConfig } from "../script/DeploymentConfig.s.sol";
import { StakeManager } from "../contracts/StakeManager.sol"; import { StakeManager } from "../contracts/StakeManager.sol";
import { StakeVault } from "../contracts/StakeVault.sol"; import { StakeVault } from "../contracts/StakeVault.sol";
@ -42,3 +45,23 @@ contract StakedTokenTest is StakeVaultTest {
assertEq(address(stakeVault.stakedToken()), stakeToken); assertEq(address(stakeVault.stakedToken()), stakeToken);
} }
} }
contract StakeTest is StakeVaultTest {
function setUp() public override {
DeployBroken deployment = new DeployBroken();
(vaultFactory, stakeManager, stakeToken) = deployment.run();
vm.prank(testUser);
stakeVault = vaultFactory.createVault();
}
function test_RevertWhen_StakeTokenTransferFails() public {
// ensure user has funds
deal(stakeToken, testUser, 1000);
vm.startPrank(address(testUser));
ERC20(stakeToken).approve(address(stakeVault), 100);
vm.expectRevert(StakeVault.StakeVault__StakingFailed.selector);
stakeVault.stake(100, 0);
}
}

View File

@ -0,0 +1,20 @@
// SPDX-License-Identifier: UNLICENSED
pragma solidity ^0.8.19;
import { ERC20 } from "@openzeppelin/contracts/token/ERC20/ERC20.sol";
contract BrokenERC20 is ERC20 {
constructor() ERC20("Mock SNT", "SNT") {
_mint(msg.sender, 1_000_000_000_000_000_000);
}
// solhint-disable-next-line no-unused-vars
function transferFrom(address sender, address recipient, uint256 amount) public override returns (bool) {
return false;
}
// solhint-disable-next-line no-unused-vars
function transfer(address recipient, uint256 amount) public override returns (bool) {
return false;
}
}

View File

@ -0,0 +1,20 @@
// SPDX-License-Identifier: UNLICENSED
pragma solidity ^0.8.19;
import { BaseScript } from "../../script/Base.s.sol";
import { StakeManager } from "../../contracts/StakeManager.sol";
import { VaultFactory } from "../../contracts/VaultFactory.sol";
import { BrokenERC20 } from "../mocks/BrokenERC20.s.sol";
contract DeployBroken is BaseScript {
function run() public returns (VaultFactory, StakeManager, address) {
BrokenERC20 token = new BrokenERC20();
vm.startBroadcast(broadcaster);
StakeManager stakeManager = new StakeManager(address(token), address(0));
VaultFactory vaultFactory = new VaultFactory(address(stakeManager));
vm.stopBroadcast();
return (vaultFactory, stakeManager, address(token));
}
}