added native withdrawal logic to vaults

contracts/linkStaking/OperatorVault.sol

@@ -98,6 +98,16 @@ contract OperatorVault is Vault {
         IERC677(address(token)).transferAndCall(address(stakeController), _amount, "");
     }
 
+    /**
+     * @notice withdraws tokens from the Chainlink staking contract
+     * @param _amount amount to withdraw
+     */
+    function withdraw(uint256 _amount) external override onlyVaultController {
+        trackedTotalDeposits -= SafeCast.toUint128(_amount);
+        stakeController.unstake(_amount, false);
+        token.safeTransfer(vaultController, _amount);
+    }
+
     /**
      * @notice returns the principal balance of this contract in the Chainlink staking contract
      * @return principal balance

contracts/linkStaking/base/Vault.sol

@@ -65,10 +65,19 @@ abstract contract Vault is Initializable, UUPSUpgradeable, OwnableUpgradeable {
     }
 
     /**
-     * @notice withdrawals are not yet implemented
+     * @notice withdraws tokens from the Chainlink staking contract
+     * @param _amount amount to withdraw
      */
-    function withdraw(uint256) external view onlyVaultController {
-        revert("withdrawals not yet implemented");
+    function withdraw(uint256 _amount) external virtual onlyVaultController {
+        stakeController.unstake(_amount, false);
+        token.safeTransfer(vaultController, _amount);
+    }
+
+    /**
+     * @notice unbonds tokens in the Chainlink staking contract
+     */
+    function unbond() external onlyVaultController {
+        stakeController.unbond();
     }
 
     /**
@@ -96,6 +105,14 @@ abstract contract Vault is Initializable, UUPSUpgradeable, OwnableUpgradeable {
         return rewardsController.getReward(address(this));
     }
 
+    /**
+     * @notice returns whether the unbonding or claim period is active for this contract in the Chainlink staking contract
+     * @return rewards balance
+     */
+    function unbondingActive() external view returns (bool) {
+        return block.timestamp < stakeController.getClaimPeriodEndsAt(address(this));
+    }
+
     /**
      * @dev Checks authorization for contract upgrades
      */

contracts/linkStaking/interfaces/IStaking.sol

@@ -16,5 +16,11 @@ interface IStaking {
 
     function getMerkleRoot() external view returns (bytes32);
 
+    function getClaimPeriodEndsAt(address _staker) external view returns (uint256);
+
     function migrate(bytes calldata data) external;
+
+    function unbond() external;
+
+    function unstake(uint256 _amount, bool _shouldClaimReward) external;
 }

contracts/linkStaking/interfaces/IVault.sol

@@ -4,14 +4,18 @@ pragma solidity 0.8.15;
 interface IVault {
     function deposit(uint256 _amount) external;
 
-    function withdraw(uint256 _amount) external view;
+    function withdraw(uint256 _amount) external;
+
+    function unbond() external;
 
     function getTotalDeposits() external view returns (uint256);
 
     function getPrincipalDeposits() external view returns (uint256);
 
     function getRewards() external view returns (uint256);
 
+    function unbondingActive() external view returns (bool);
+
     function migrate(bytes calldata _data) external;
 
     function upgradeToAndCall(address _newImplementation, bytes memory _data) external;

contracts/linkStaking/test/OperatorVCSMock.sol

@@ -35,6 +35,15 @@ contract OperatorVCSMock {
         vault.deposit(_amount);
     }
 
+    function withdraw(uint256 _amount) external {
+        vault.withdraw(_amount);
+        token.safeTransfer(msg.sender, _amount);
+    }
+
+    function unbond() external {
+        vault.unbond();
+    }
+
     function withdrawOperatorRewards(address _receiver, uint256 _amount) external returns (uint256) {
         uint256 withdrawalAmount = (_amount * withdrawalPercentage) / 10000;
         return withdrawalAmount;

contracts/linkStaking/test/StakingMock.sol

@@ -1,6 +1,8 @@
 // SPDX-License-Identifier: GPL-3.0
 pragma solidity 0.8.15;
 
+import "@openzeppelin/contracts/token/ERC20/utils/SafeERC20.sol";
+
 import "../../core/interfaces/IERC677.sol";
 import "../../core/interfaces/IERC677Receiver.sol";
 
@@ -9,31 +11,52 @@ import "../../core/interfaces/IERC677Receiver.sol";
  * @dev Mocks contract for testing
  */
 contract StakingMock is IERC677Receiver {
+    using SafeERC20 for IERC677;
+
+    struct Staker {
+        uint256 unbondingPeriodEndsAt;
+        uint256 claimPeriodEndsAt;
+        uint256 principal;
+        uint256 removedPrincipal;
+    }
+
     IERC677 public token;
     address public rewardVault;
 
-    mapping(address => uint256) public principalBalances;
-    mapping(address => uint256) public removedPrincipal;
-
     uint256 public depositMin;
     uint256 public depositMax;
     uint256 public maxPoolSize;
 
+    uint256 public unbondingPeriod;
+    uint256 public claimPeriod;
+
     bool public active;
 
+    mapping(address => Staker) public stakers;
+
+    error UnbondingPeriodActive();
+    error NotInClaimPeriod();
+    error UnstakeZeroAmount();
+    error UnstakeExceedsPrincipal();
+    error UnstakePrincipalBelowMinAmount();
+
     constructor(
         address _token,
         address _rewardVault,
         uint256 _depositMin,
         uint256 _depositMax,
-        uint256 _maxPoolSize
+        uint256 _maxPoolSize,
+        uint256 _unbondingPeriod,
+        uint256 _claimPeriod
     ) {
         token = IERC677(_token);
         rewardVault = _rewardVault;
         active = true;
         depositMin = _depositMin;
         depositMax = _depositMax;
         maxPoolSize = _maxPoolSize;
+        unbondingPeriod = _unbondingPeriod;
+        claimPeriod = _claimPeriod;
     }
 
     function onTokenTransfer(
@@ -44,10 +67,48 @@ contract StakingMock is IERC677Receiver {
         require(msg.sender == address(token), "has to be token");
         if (_data.length != 0) {
             address sender = abi.decode(_data, (address));
-            principalBalances[sender] += _value;
+            stakers[sender].principal += _value;
         } else {
-            principalBalances[_sender] += _value;
+            stakers[_sender].principal += _value;
+        }
+
+        delete stakers[_sender].unbondingPeriodEndsAt;
+        delete stakers[_sender].claimPeriodEndsAt;
+    }
+
+    function unbond() external {
+        Staker memory staker = stakers[msg.sender];
+
+        if (staker.unbondingPeriodEndsAt != 0 && block.timestamp <= staker.claimPeriodEndsAt) {
+            revert UnbondingPeriodActive();
+        }
+
+        staker.unbondingPeriodEndsAt = block.timestamp + unbondingPeriod;
+        staker.claimPeriodEndsAt = staker.unbondingPeriodEndsAt + claimPeriod;
+        stakers[msg.sender] = staker;
+    }
+
+    function unstake(uint256 _amount, bool) external {
+        Staker memory staker = stakers[msg.sender];
+
+        if (
+            staker.unbondingPeriodEndsAt == 0 ||
+            block.timestamp < staker.unbondingPeriodEndsAt ||
+            block.timestamp > staker.claimPeriodEndsAt
+        ) {
+            revert NotInClaimPeriod();
         }
+        if (_amount == 0) revert UnstakeZeroAmount();
+
+        if (_amount > staker.principal) revert UnstakeExceedsPrincipal();
+
+        uint256 updatedPrincipal = staker.principal - _amount;
+        if (_amount < staker.principal && updatedPrincipal < depositMin) {
+            revert UnstakePrincipalBelowMinAmount();
+        }
+
+        stakers[msg.sender].principal -= _amount;
+        token.safeTransfer(msg.sender, _amount);
     }
 
     function getStakerLimits() external view returns (uint256, uint256) {
@@ -63,20 +124,24 @@ contract StakingMock is IERC677Receiver {
     }
 
     function getStakerPrincipal(address _staker) external view returns (uint256) {
-        return principalBalances[_staker];
+        return stakers[_staker].principal;
     }
 
     function getRemovedPrincipal(address _staker) external view returns (uint256) {
-        return removedPrincipal[_staker];
+        return stakers[_staker].removedPrincipal;
+    }
+
+    function getClaimPeriodEndsAt(address _staker) external view returns (uint256) {
+        return stakers[_staker].claimPeriodEndsAt;
     }
 
     function getRewardVault() external view returns (address) {
         return rewardVault;
     }
 
     function removePrincipal(address _staker, uint256 _amount) external {
-        principalBalances[_staker] -= _amount;
-        removedPrincipal[_staker] += _amount;
+        stakers[_staker].principal -= _amount;
+        stakers[_staker].removedPrincipal += _amount;
     }
 
     function getMerkleRoot() external view returns (bytes32) {

test/linkStaking/community-vault.test.ts

@@ -30,6 +30,8 @@ describe('CommunityVault', () => {
       toEther(10),
       toEther(100),
       toEther(10000),
+      28 * 86400,
+      7 * 86400,
     ])) as StakingMock
 
     vault = (await deployUpgradeable('CommunityVault', [

test/linkStaking/operator-vault.test.ts

@@ -10,6 +10,10 @@ import {
   StakingMock,
   StakingRewardsMock,
 } from '../../typechain-types'
+import { time } from '@nomicfoundation/hardhat-network-helpers'
+
+const unbondingPeriod = 28 * 86400
+const claimPeriod = 7 * 86400
 
 describe('OperatorVault', () => {
   let token: ERC677
@@ -39,6 +43,8 @@ describe('OperatorVault', () => {
       toEther(10),
       toEther(100),
       toEther(10000),
+      unbondingPeriod,
+      claimPeriod,
     ])) as StakingMock
     pfAlertsController = (await deploy('PFAlertsControllerMock', [
       token.address,
@@ -71,6 +77,21 @@ describe('OperatorVault', () => {
     assert.equal(fromEther(await vault.trackedTotalDeposits()), 200)
   })
 
+  it('withdraw should work correctly', async () => {
+    await strategy.unbond()
+
+    await expect(strategy.withdraw(toEther(30))).to.be.revertedWith('NotInClaimPeriod()')
+
+    await time.increase(unbondingPeriod + 1)
+
+    await strategy.withdraw(toEther(30))
+    assert.equal(fromEther(await token.balanceOf(stakingController.address)), 70)
+    assert.equal(fromEther(await stakingController.getStakerPrincipal(vault.address)), 70)
+    assert.equal(fromEther(await vault.getTotalDeposits()), 70)
+    assert.equal(fromEther(await vault.getUnclaimedRewards()), 0)
+    assert.equal(fromEther(await vault.trackedTotalDeposits()), 70)
+  })
+
   it('raiseAlert should work correctly', async () => {
     await vault.connect(signers[1]).raiseAlert(accounts[5])
     assert.equal(fromEther(await token.balanceOf(strategy.address)), 11.7)

test/linkStaking/vault.test.ts

@@ -9,6 +9,10 @@ import {
   fromEther,
 } from '../utils/helpers'
 import { ERC677, CommunityVault, StakingMock, StakingRewardsMock } from '../../typechain-types'
+import { time } from '@nomicfoundation/hardhat-network-helpers'
+
+const unbondingPeriod = 28 * 86400
+const claimPeriod = 7 * 86400
 
 describe('Vault', () => {
   let token: ERC677
@@ -36,6 +40,8 @@ describe('Vault', () => {
       toEther(10),
       toEther(100),
       toEther(10000),
+      unbondingPeriod,
+      claimPeriod,
     ])) as StakingMock
 
     vault = (await deployUpgradeable('CommunityVault', [
@@ -66,8 +72,27 @@ describe('Vault', () => {
     )
   })
 
-  it('should not be able to withdraw', async () => {
-    await expect(vault.withdraw(toEther(10))).to.be.revertedWith('withdrawals not yet implemented')
+  it('should be able to unbond', async () => {
+    await vault.deposit(toEther(100))
+    await vault.unbond()
+    let ts = (await ethers.provider.getBlock(await ethers.provider.getBlockNumber())).timestamp
+    assert.equal(
+      (await stakingController.getClaimPeriodEndsAt(vault.address)).toNumber(),
+      ts + unbondingPeriod + claimPeriod
+    )
+  })
+
+  it('should be able to withdraw', async () => {
+    await vault.deposit(toEther(100))
+    await vault.unbond()
+
+    await expect(vault.withdraw(toEther(30))).to.be.revertedWith('NotInClaimPeriod()')
+
+    await time.increase(unbondingPeriod + 1)
+
+    await vault.withdraw(toEther(30))
+    assert.equal(fromEther(await vault.getPrincipalDeposits()), 70)
+    assert.equal(fromEther(await token.balanceOf(stakingController.address)), 70)
   })
 
   it('getPrincipalDeposits should work correctly', async () => {
@@ -96,4 +121,20 @@ describe('Vault', () => {
     await rewardsController.setReward(vault.address, toEther(40))
     assert.equal(fromEther(await vault.getTotalDeposits()), 290)
   })
+
+  it('unbondingActive should work correctly', async () => {
+    assert.equal(await vault.unbondingActive(), false)
+
+    await vault.deposit(toEther(100))
+    assert.equal(await vault.unbondingActive(), false)
+
+    await vault.unbond()
+    assert.equal(await vault.unbondingActive(), true)
+
+    await time.increase(unbondingPeriod + 1)
+    assert.equal(await vault.unbondingActive(), true)
+
+    await time.increase(claimPeriod)
+    assert.equal(await vault.unbondingActive(), false)
+  })
 })