From b5963f72c19e6f3e0a671bc74bc7fb7d3f2d4ac1 Mon Sep 17 00:00:00 2001
From: r4bbit <445106+0x-r4bbit@users.noreply.github.com>
Date: Thu, 11 Jan 2024 14:54:39 +0100
Subject: [PATCH] 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, or rather `safeTransfer`
using SafeERC20 utilities.
---
contracts/StakeVault.sol | 7 +++++--
test/StakeManager.t.sol | 18 ++++++++++++++++++
2 files changed, 23 insertions(+), 2 deletions(-)
diff --git a/contracts/StakeVault.sol b/contracts/StakeVault.sol
index dea4af3..84ce1ad 100644
--- a/contracts/StakeVault.sol
+++ b/contracts/StakeVault.sol
@@ -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";
/**
@@ -13,6 +14,8 @@ import { StakeManager } from "./StakeManager.sol";
*/
contract StakeVault is Ownable {
+ using SafeERC20 for ERC20;
+
error StakeVault__MigrationNotAvailable();
StakeManager private stakeManager;
@@ -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);
@@ -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 {
diff --git a/test/StakeManager.t.sol b/test/StakeManager.t.sol
index d5745b9..97322b4 100644
--- a/test/StakeManager.t.sol
+++ b/test/StakeManager.t.sol
@@ -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 {