Skip to content

Commit

Permalink
fix(StakeVault): make unstaking actually work
Browse files Browse the repository at this point in the history
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, or rather `safeTransfer`
using SafeERC20 utilities.
  • Loading branch information
0x-r4bbit committed Jan 11, 2024
1 parent cf7a8b6 commit b5963f7
Show file tree
Hide file tree
Showing 2 changed files with 23 additions and 2 deletions.
7 changes: 5 additions & 2 deletions contracts/StakeVault.sol
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@ pragma solidity ^0.8.18;

import { Ownable } from "@openzeppelin/contracts/access/Ownable.sol";
import { ERC20 } from "@openzeppelin/contracts/token/ERC20/ERC20.sol";
import { SafeERC20 } from "@openzeppelin/contracts/token/ERC20/utils/SafeERC20.sol";
import { StakeManager } from "./StakeManager.sol";

/**
Expand All @@ -13,6 +14,8 @@ import { StakeManager } from "./StakeManager.sol";
*/

contract StakeVault is Ownable {
using SafeERC20 for ERC20;

error StakeVault__MigrationNotAvailable();

StakeManager private stakeManager;
Expand All @@ -28,7 +31,7 @@ contract StakeVault is Ownable {
}

function stake(uint256 _amount, uint256 _time) external onlyOwner {
STAKED_TOKEN.transferFrom(msg.sender, address(this), _amount);
STAKED_TOKEN.safeTransferFrom(msg.sender, address(this), _amount);
stakeManager.stake(_amount, _time);

emit Staked(msg.sender, address(this), _amount, _time);
Expand All @@ -40,7 +43,7 @@ contract StakeVault is Ownable {

function unstake(uint256 _amount) external onlyOwner {
stakeManager.unstake(_amount);
STAKED_TOKEN.transferFrom(address(this), msg.sender, _amount);
STAKED_TOKEN.safeTransfer(msg.sender, _amount);
}

function leave() external onlyOwner {
Expand Down
18 changes: 18 additions & 0 deletions test/StakeManager.t.sol
Original file line number Diff line number Diff line change
Expand Up @@ -97,6 +97,24 @@ contract UnstakeTest is StakeManagerTest {
vm.expectRevert(StakeManager.StakeManager__FundsLocked.selector);
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 {
Expand Down

0 comments on commit b5963f7

Please sign in to comment.