Token Staking - unstaking

core/contracts/staking/TokenStaking.sol

@@ -13,11 +13,17 @@ import "../shared/IReceiveApproval.sol";
 contract TokenStaking is IReceiveApproval {
     using SafeERC20 for IERC20;
 
+    struct Staker {
+        uint256 balance;
+        uint256 startStakingTimestamp;
+    }
+
     IERC20 internal immutable token;
 
-    mapping(address => uint256) public balanceOf;
+    mapping(address => Staker) public stakers;
 
     event Staked(address indexed staker, uint256 amount);
+    event Unstaked(address indexed staker, uint256 amount);
 
     constructor(IERC20 _token) {
         require(
@@ -51,6 +57,23 @@ contract TokenStaking is IReceiveApproval {
         _stake(msg.sender, amount);
     }
 
+    /// @notice Reduces stake amount by the provided amount and
+    ///         withdraws tokens to the owner.
+    /// @param amount Amount to unstake and withdraw.
+    function unstake(uint256 amount) external {
+        require((amount > 0), "Amount can not be zero");
+
+        Staker storage staker = stakers[msg.sender];
+
+        require(staker.balance > 0, "Nothing to unstake");
+        require(staker.balance >= amount, "Insufficient funds");
+
+        staker.balance -= amount;
+
+        emit Unstaked(msg.sender, amount);
+        token.safeTransfer(msg.sender, amount);
+    }
+
     /// @notice Returns minimum amount of staking tokens to participate in
     ///         protocol.
     function minimumStake() public pure returns (uint256) {
@@ -71,7 +94,10 @@ contract TokenStaking is IReceiveApproval {
         require(amount <= maximumStake(), "Amount is greater than maxium");
         require(staker != address(0), "Can not be the zero address");
 
-        balanceOf[staker] += amount;
+        Staker storage stakerStruct = stakers[staker];
+        stakerStruct.balance += amount;
+        /* solhint-disable-next-line not-rely-on-time */
+        stakerStruct.startStakingTimestamp = block.timestamp;
 
         // TODO: Mint stBTC token.
         emit Staked(staker, amount);

core/contracts/test/TestToken.sol

@@ -1,6 +1,6 @@
 // SPDX-License-Identifier: GPL-3.0-or-later
 
-pragma solidity 0.8.20;
+pragma solidity ^0.8.20;
 
 import "@openzeppelin/contracts/token/ERC20/ERC20.sol";
 import "../shared/IReceiveApproval.sol";

core/test/staking/TokenStaking.test.ts

@@ -49,9 +49,9 @@ describe("TokenStaking", () => {
         await expect(tokenStaking.connect(tokenHolder).stake(amountToStake))
           .to.emit(tokenStaking, "Staked")
           .withArgs(tokenHolderAddress, amountToStake)
-        expect(await tokenStaking.balanceOf(tokenHolderAddress)).to.be.eq(
-          amountToStake,
-        )
+        expect(
+          (await tokenStaking.stakers(tokenHolderAddress)).balance,
+        ).to.be.eq(amountToStake)
         expect(await token.balanceOf(tokenHolderAddress)).to.be.eq(
           tokenBalanceBeforeStake - amountToStake,
         )
@@ -86,13 +86,54 @@ describe("TokenStaking", () => {
         )
           .to.emit(tokenStaking, "Staked")
           .withArgs(tokenHolderAddress, amountToStake)
-        expect(await tokenStaking.balanceOf(tokenHolderAddress)).to.be.eq(
-          amountToStake,
-        )
+        expect(
+          (await tokenStaking.stakers(tokenHolderAddress)).balance,
+        ).to.be.eq(amountToStake)
         expect(await token.balanceOf(tokenHolderAddress)).to.be.eq(
           tokenBalanceBeforeStake - amountToStake,
         )
       })
     })
   })
+
+  describe("unstaking", () => {
+    const amountToStake = WeiPerEther * 10n
+
+    beforeEach(async () => {
+      // Stake tokens.
+      await token
+        .connect(tokenHolder)
+        .approveAndCall(await tokenStaking.getAddress(), amountToStake, "0x")
+    })
+
+    it("should unstake tokens", async () => {
+      const staker = await tokenHolder.getAddress()
+      const stakingBalance = (await tokenStaking.stakers(staker)).balance
+      const balanceBeforeUnstaking = await token.balanceOf(staker)
+
+      await expect(tokenStaking.connect(tokenHolder).unstake(stakingBalance))
+        .to.emit(tokenStaking, "Unstaked")
+        .withArgs(staker, stakingBalance)
+
+      expect(await token.balanceOf(staker)).to.be.equal(
+        balanceBeforeUnstaking + stakingBalance,
+      )
+      expect((await tokenStaking.stakers(staker)).balance).to.be.eq(0)
+    })
+
+    it("should revert if the unstaked amount is equal 0", async () => {
+      await expect(
+        tokenStaking.connect(tokenHolder).unstake(0),
+      ).to.be.revertedWith("Amount can not be zero")
+    })
+
+    it("should revert if the user wants to unstake more tokens than currently staked", async () => {
+      const staker = await tokenHolder.getAddress()
+      const stakingBalance = (await tokenStaking.stakers(staker)).balance
+
+      await expect(
+        tokenStaking.connect(tokenHolder).unstake(stakingBalance + 10n),
+      ).to.be.revertedWith("Insufficient funds")
+    })
+  })
 })