From 12c776ed862775c8f0ecba9175a793a325ab42e0 Mon Sep 17 00:00:00 2001
From: BkChoy <christian.agnew@outlook.com>
Date: Mon, 21 Oct 2024 09:47:36 -0400
Subject: [PATCH] staking proxy contract

---
 contracts/core/StakingProxy.sol               | 293 +++++++++++
 contracts/core/interfaces/IPriorityPool.sol   |  30 ++
 contracts/core/interfaces/ISDLPool.sol        |   2 +
 contracts/core/interfaces/IWithdrawalPool.sol |  10 +
 .../core/priorityPool/WithdrawalPool.sol      |  21 +
 .../core/priorityPool/withdrawal-pool.test.ts |  11 +
 test/core/staking-proxy.test.ts               | 454 ++++++++++++++++++
 7 files changed, 821 insertions(+)
 create mode 100644 contracts/core/StakingProxy.sol
 create mode 100644 test/core/staking-proxy.test.ts

diff --git a/contracts/core/StakingProxy.sol b/contracts/core/StakingProxy.sol
new file mode 100644
index 00000000..4b975995
--- /dev/null
+++ b/contracts/core/StakingProxy.sol
@@ -0,0 +1,293 @@
+// SPDX-License-Identifier: GPL-3.0
+pragma solidity 0.8.15;
+
+import "@openzeppelin/contracts-upgradeable/proxy/utils/UUPSUpgradeable.sol";
+import "@openzeppelin/contracts-upgradeable/access/OwnableUpgradeable.sol";
+import "@openzeppelin/contracts-upgradeable/token/ERC20/utils/SafeERC20Upgradeable.sol";
+import "@openzeppelin/contracts/token/ERC721/IERC721.sol";
+
+import "./interfaces/IStakingPool.sol";
+import "./interfaces/IPriorityPool.sol";
+import "./interfaces/IWithdrawalPool.sol";
+import "./interfaces/IERC677.sol";
+import "./interfaces/ISDLPool.sol";
+
+/**
+ * @title Staking Proxy
+ * @notice Enables a staker to deposit tokens and earn rewards without ever directly interacting with the LST contracts
+ * @dev When tokens are queued for deposit, the corresponding liquid staking tokens will be distributed using a merkle
+ * tree. The tree is updated once a certain threshold of LSTs is reached at which point LSTs can be claimed. In order
+ * for this contract to claim its LSTs (and execute some other funcion calls), merkle data must be passed as an argument.
+ * This data is stored on IPFS at the hash which can be queried from this contract. For some functions, this data can be
+ * used as is but for ones that require a merkle proof, a merkle tree must be generated using the IPFS data, then a merkle
+ * proof generated using the tree.
+ * @dev Data may or may not need to be passed when depositing and/or withdrawing depending on the underlying LST implementation.
+ * If data does need to be passed, it will need to be fetched from an external source specific to that LST.
+ */
+contract StakingProxy is UUPSUpgradeable, OwnableUpgradeable {
+    using SafeERC20Upgradeable for IERC20Upgradeable;
+
+    // address of asset token
+    IERC20Upgradeable public token;
+    // address of liquid staking token
+    IStakingPool public lst;
+    // address of priority pool
+    IPriorityPool public priorityPool;
+    // address of withdrawal pool
+    IWithdrawalPool public withdrawalPool;
+    // address of SDL pool
+    ISDLPool public sdlPool;
+
+    // address authorized to deposit/withdraw asset tokens
+    address public staker;
+
+    error SenderNotAuthorized();
+    error InvalidToken();
+    error InvalidValue();
+    error InvalidTokenId();
+
+    /// @custom:oz-upgrades-unsafe-allow constructor
+    constructor() {
+        _disableInitializers();
+    }
+
+    /**
+     * @notice Initializes the contract
+     * @param _token address of asset token
+     * @param _lst address of liquid staking token
+     * @param _priorityPool address of priority pool
+     * @param _withdrawalPool address of withdrawal pool
+     * @param _sdlPool address of SDL pool
+     * @param _staker address authorized to deposit/withdraw asset tokens
+     */
+    function initialize(
+        address _token,
+        address _lst,
+        address _priorityPool,
+        address _withdrawalPool,
+        address _sdlPool,
+        address _staker
+    ) public initializer {
+        __Ownable_init();
+        __UUPSUpgradeable_init();
+
+        token = IERC20Upgradeable(_token);
+        lst = IStakingPool(_lst);
+        priorityPool = IPriorityPool(_priorityPool);
+        token.safeApprove(_priorityPool, type(uint256).max);
+        IERC20Upgradeable(_lst).safeApprove(_priorityPool, type(uint256).max);
+        withdrawalPool = IWithdrawalPool(_withdrawalPool);
+        sdlPool = ISDLPool(_sdlPool);
+        staker = _staker;
+    }
+
+    /**
+     * @notice Reverts if sender is not staker
+     */
+    modifier onlyStaker() {
+        if (msg.sender != staker) revert SenderNotAuthorized();
+        _;
+    }
+
+    /**
+     * @notice Returns the IPFS hash of the current merkle distribution tree
+     * @dev returns CIDv0 with no prefix - prefix must be added and hash must be properly encoded offchain
+     * @return IPFS hash
+     */
+    function getMerkleIPFSHash() external view returns (bytes32) {
+        return priorityPool.ipfsHash();
+    }
+
+    /**
+     * @notice Returns the total amount of liquid staking tokens held by this contract
+     * @dev excludes unclaims LSTs sitting in the priority pool
+     * @return total LSTs held by contract
+     */
+    function getTotalDeposits() external view returns (uint256) {
+        return lst.balanceOf(address(this));
+    }
+
+    /**
+     * @notice Returns the total amount of tokens queued for deposit into the staking pool by this contract
+     * @param _distributionAmount amount as recorded in this contract's merkle tree entry (stored on IPFS)
+     * @return total tokens queued for deposit
+     */
+    function getTotalQueuedForDeposit(uint256 _distributionAmount) external view returns (uint256) {
+        return priorityPool.getQueuedTokens(address(this), _distributionAmount);
+    }
+
+    /**
+     * @notice Returns the total amount of tokens queued for withdrawal from the staking pool by this contract
+     * @return total tokens queued for withdrawal
+     */
+    function getTotalQueuedForWithdrawal() external view returns (uint256) {
+        return withdrawalPool.getAccountTotalQueuedWithdrawals(address(this));
+    }
+
+    /**
+     * @notice Returns the total amount of withdrawable tokens for this contract
+     * @param _distributionAmount amount as recorded in this contract's merkle tree entry (stored on IPFS)
+     * @return total amount withdrawable from the priority pool
+     * @return total amount withdrawable from the withdrawal pool
+     * @return withdrawal ids for withdrawal from withdrawal pool
+     * @return batch ids for withdrawal from withdrawal pool
+     */
+    function getTotalWithdrawable(
+        uint256 _distributionAmount
+    ) external view returns (uint256, uint256, uint256[] memory, uint256[] memory) {
+        (uint256[] memory withdrawalIds, uint256 withdrawable) = withdrawalPool
+            .getFinalizedWithdrawalIdsByOwner(address(this));
+        uint256[] memory batchIds = withdrawalPool.getBatchIds(withdrawalIds);
+
+        uint256 priorityPoolCanWithdraw = priorityPool.canWithdraw(
+            address(this),
+            _distributionAmount
+        );
+
+        return (priorityPoolCanWithdraw, withdrawable, withdrawalIds, batchIds);
+    }
+
+    /**
+     * @notice Returns the total amount of claimable liquid staking tokens for this contract
+     * @param _distributionSharesAmount shares amount as recorded in this contract's merkle tree entry (stored on IPFS)
+     * @return total claimable LSTs
+     */
+    function getTotalClaimableLSTs(
+        uint256 _distributionSharesAmount
+    ) external view returns (uint256) {
+        return priorityPool.getLSDTokens(address(this), _distributionSharesAmount);
+    }
+
+    /**
+     * @notice ERC677 implementation to receive deposits
+     * @param _sender address of sender
+     * @param _value value of transfer
+     * @param _calldata encoded deposit data
+     */
+    function onTokenTransfer(address _sender, uint256 _value, bytes calldata _calldata) external {
+        if (msg.sender != address(token)) revert InvalidToken();
+        if (_sender != staker) revert SenderNotAuthorized();
+        if (_value == 0) revert InvalidValue();
+
+        bytes[] memory data = abi.decode(_calldata, (bytes[]));
+        IERC677(address(token)).transferAndCall(
+            address(priorityPool),
+            _value,
+            abi.encode(true, data)
+        );
+    }
+
+    /**
+     * @notice Deposits tokens and/or queues tokens for deposit into the staking pool
+     * @param _amount amount of tokens to deposit
+     * @param _data encoded deposit data
+     */
+    function deposit(uint256 _amount, bytes[] calldata _data) external onlyStaker {
+        token.safeTransferFrom(msg.sender, address(this), _amount);
+        priorityPool.deposit(_amount, true, _data);
+    }
+
+    /**
+     * @notice Withdraws tokens and/or queues tokens for withdrawal from the staking pool
+     * @dev if there is any amount withdrawable from the withdrawal pool, the entire amount will be withdrawn
+     * even if it exceeds _amountToWithdraw
+     * @param _amountToWithdraw amount of tokens to withdraw
+     * @param _distributionAmount amount as recorded in this contract's merkle tree entry (stored on IPFS)
+     * @param _distributionSharesAmount shares amount as recorded in this contract's merkle tree entry (stored on IPFS)
+     * @param _merkleProof merkle proof for this contract's merkle tree entry (generated using IPFS data)
+     * @param _withdrawalIds list of withdrawal ids required if finalizing queued withdrawals
+     * @param _batchIds list of batch ids required if finalizing queued withdrawals
+     * @param _data encoded withdrawal data
+     */
+    function withdraw(
+        uint256 _amountToWithdraw,
+        uint256 _distributionAmount,
+        uint256 _distributionSharesAmount,
+        bytes32[] calldata _merkleProof,
+        uint256[] calldata _withdrawalIds,
+        uint256[] calldata _batchIds,
+        bytes[] calldata _data
+    ) external onlyStaker {
+        uint256 availableTokens;
+
+        if (_withdrawalIds.length != 0) {
+            withdrawalPool.withdraw(_withdrawalIds, _batchIds);
+            availableTokens = token.balanceOf(address(this));
+        }
+
+        if (availableTokens < _amountToWithdraw) {
+            priorityPool.withdraw(
+                _amountToWithdraw - availableTokens,
+                _distributionAmount,
+                _distributionSharesAmount,
+                _merkleProof,
+                true,
+                true,
+                _data
+            );
+            availableTokens = token.balanceOf(address(this));
+        }
+
+        token.safeTransfer(msg.sender, availableTokens);
+    }
+
+    /**
+     * @notice Claims liquid staking tokens from the priority pool
+     * @param _amount amount as recorded in this contract's merkle tree entry (stored on IPFS)
+     * @param _sharesAmount shares amount as recorded in this contract's merkle tree entry (stored on IPFS)
+     * @param _merkleProof merkle proof for this contract's merkle tree entry (generated from IPFS data)
+     */
+    function claimLSTs(
+        uint256 _amount,
+        uint256 _sharesAmount,
+        bytes32[] calldata _merkleProof
+    ) external {
+        priorityPool.claimLSDTokens(_amount, _sharesAmount, _merkleProof);
+    }
+
+    /**
+     * @notice Returns a list of reSDL token ids held by this contract
+     * @return list of token ids
+     */
+    function getRESDLTokenIds() external view returns (uint256[] memory) {
+        return sdlPool.getLockIdsByOwner(address(this));
+    }
+
+    /**
+     * @notice Called when an reSDL token is transferred to this contract using safeTransfer
+     */
+    function onERC721Received(address, address, uint256, bytes calldata) external returns (bytes4) {
+        return this.onERC721Received.selector;
+    }
+
+    /**
+     * @notice Withdraws an reSDL token
+     * @param _tokenId id of token
+     * @param _receiver address to receive token
+     */
+    function withdrawRESDLToken(uint256 _tokenId, address _receiver) external onlyOwner {
+        if (sdlPool.ownerOf(_tokenId) != address(this)) revert InvalidTokenId();
+        IERC721(address(sdlPool)).safeTransferFrom(address(this), _receiver, _tokenId);
+    }
+
+    /**
+     * @notice Claims rewards from the SDL Pool
+     * @dev rewards will be redistributed to the SDL Pool
+     * @param _tokens list of tokens to claim rewards for
+     */
+    function claimRESDLRewards(address[] calldata _tokens) external {
+        sdlPool.withdrawRewards(_tokens);
+        for (uint256 i = 0; i < _tokens.length; ++i) {
+            uint256 balance = IERC20Upgradeable(_tokens[i]).balanceOf(address(this));
+            if (balance != 0) {
+                IERC20Upgradeable(_tokens[i]).safeTransfer(address(sdlPool), balance);
+            }
+        }
+        sdlPool.distributeTokens(_tokens);
+    }
+
+    /**
+     * @dev Checks authorization for contract upgrades
+     */
+    function _authorizeUpgrade(address) internal override onlyOwner {}
+}
diff --git a/contracts/core/interfaces/IPriorityPool.sol b/contracts/core/interfaces/IPriorityPool.sol
index ded96c82..a897a92a 100644
--- a/contracts/core/interfaces/IPriorityPool.sol
+++ b/contracts/core/interfaces/IPriorityPool.sol
@@ -14,11 +14,41 @@ interface IPriorityPool {
 
     function poolStatus() external view returns (PoolStatus);
 
+    function ipfsHash() external view returns (bytes32);
+
     function canWithdraw(
         address _account,
         uint256 _distributionAmount
     ) external view returns (uint256);
 
+    function getQueuedTokens(
+        address _account,
+        uint256 _distributionAmount
+    ) external view returns (uint256);
+
+    function getLSDTokens(
+        address _account,
+        uint256 _distributionShareAmount
+    ) external view returns (uint256);
+
+    function deposit(uint256 _amount, bool _shouldQueue, bytes[] calldata _data) external;
+
+    function withdraw(
+        uint256 _amountToWithdraw,
+        uint256 _amount,
+        uint256 _sharesAmount,
+        bytes32[] calldata _merkleProof,
+        bool _shouldUnqueue,
+        bool _shouldQueueWithdrawal,
+        bytes[] calldata _data
+    ) external;
+
+    function claimLSDTokens(
+        uint256 _amount,
+        uint256 _sharesAmount,
+        bytes32[] calldata _merkleProof
+    ) external;
+
     function pauseForUpdate() external;
 
     function setPoolStatus(PoolStatus _status) external;
diff --git a/contracts/core/interfaces/ISDLPool.sol b/contracts/core/interfaces/ISDLPool.sol
index a5910945..140f7ae0 100644
--- a/contracts/core/interfaces/ISDLPool.sol
+++ b/contracts/core/interfaces/ISDLPool.sol
@@ -16,6 +16,8 @@ interface ISDLPool is IRewardsPoolController {
 
     function ownerOf(uint256 _lockId) external view returns (address);
 
+    function getLockIdsByOwner(address _owner) external view returns (uint256[] memory);
+
     function supportedTokens() external view returns (address[] memory);
 
     function handleOutgoingRESDL(
diff --git a/contracts/core/interfaces/IWithdrawalPool.sol b/contracts/core/interfaces/IWithdrawalPool.sol
index 24f3979b..511a927d 100644
--- a/contracts/core/interfaces/IWithdrawalPool.sol
+++ b/contracts/core/interfaces/IWithdrawalPool.sol
@@ -4,8 +4,18 @@ pragma solidity 0.8.15;
 interface IWithdrawalPool {
     function getTotalQueuedWithdrawals() external view returns (uint256);
 
+    function getAccountTotalQueuedWithdrawals(address _account) external view returns (uint256);
+
+    function getFinalizedWithdrawalIdsByOwner(
+        address _account
+    ) external view returns (uint256[] memory, uint256);
+
+    function getBatchIds(uint256[] memory _withdrawalIds) external view returns (uint256[] memory);
+
     function deposit(uint256 _amount) external;
 
+    function withdraw(uint256[] calldata _withdrawalIds, uint256[] calldata _batchIds) external;
+
     function queueWithdrawal(address _account, uint256 _amount) external;
 
     function performUpkeep(bytes calldata _performData) external;
diff --git a/contracts/core/priorityPool/WithdrawalPool.sol b/contracts/core/priorityPool/WithdrawalPool.sol
index dd6e807f..a4293dbe 100644
--- a/contracts/core/priorityPool/WithdrawalPool.sol
+++ b/contracts/core/priorityPool/WithdrawalPool.sol
@@ -127,6 +127,27 @@ contract WithdrawalPool is UUPSUpgradeable, OwnableUpgradeable {
         return _getStakeByShares(totalQueuedShareWithdrawals);
     }
 
+    /**
+     * @notice Returns the total amount of liquid staking tokens queued for withdrawal by an account
+     * @param _account address of account
+     * @return total amount queued across all account's withdrawals
+     */
+    function getAccountTotalQueuedWithdrawals(address _account) external view returns (uint256) {
+        uint256[] memory withdrawalIds = getWithdrawalIdsByOwner(_account);
+        uint256[] memory batchIds = getBatchIds(withdrawalIds);
+
+        uint256 totalUnfinalizedWithdrawals;
+        for (uint256 i = 0; i < batchIds.length; ++i) {
+            Withdrawal memory withdrawal = queuedWithdrawals[withdrawalIds[i]];
+
+            if (batchIds[i] == 0) {
+                totalUnfinalizedWithdrawals += uint256(withdrawal.sharesRemaining);
+            }
+        }
+
+        return _getStakeByShares(totalUnfinalizedWithdrawals);
+    }
+
     /**
      * @notice Returns a list of withdrawals
      * @param _withdrawalIds list of withdrawal ids
diff --git a/test/core/priorityPool/withdrawal-pool.test.ts b/test/core/priorityPool/withdrawal-pool.test.ts
index 963e2408..cfb35c16 100644
--- a/test/core/priorityPool/withdrawal-pool.test.ts
+++ b/test/core/priorityPool/withdrawal-pool.test.ts
@@ -74,6 +74,12 @@ describe('WithdrawalPool', () => {
 
     assert.equal(fromEther(await stakingPool.balanceOf(withdrawalPool.target)), 1750)
     assert.equal(fromEther(await withdrawalPool.getTotalQueuedWithdrawals()), 1750)
+    assert.equal(
+      fromEther(await withdrawalPool.getAccountTotalQueuedWithdrawals(accounts[0])),
+      1500
+    )
+    assert.equal(fromEther(await withdrawalPool.getAccountTotalQueuedWithdrawals(accounts[1])), 250)
+
     assert.deepEqual(
       (await withdrawalPool.getWithdrawalIdsByOwner(accounts[0])).map((id) => Number(id)),
       [1, 3]
@@ -104,6 +110,11 @@ describe('WithdrawalPool', () => {
     assert.equal(fromEther(await token.balanceOf(withdrawalPool.target)), 400)
     assert.equal(fromEther(await stakingPool.balanceOf(withdrawalPool.target)), 1350)
     assert.equal(fromEther(await withdrawalPool.getTotalQueuedWithdrawals()), 1350)
+    assert.equal(
+      fromEther(await withdrawalPool.getAccountTotalQueuedWithdrawals(accounts[0])),
+      1100
+    )
+    assert.equal(fromEther(await withdrawalPool.getAccountTotalQueuedWithdrawals(accounts[1])), 250)
     assert.equal(Number(await withdrawalPool.indexOfNextWithdrawal()), 1)
     assert.deepEqual(
       (await withdrawalPool.getWithdrawals([1, 2, 3])).map((d: any) => [
diff --git a/test/core/staking-proxy.test.ts b/test/core/staking-proxy.test.ts
new file mode 100644
index 00000000..bb0b883c
--- /dev/null
+++ b/test/core/staking-proxy.test.ts
@@ -0,0 +1,454 @@
+import { assert, expect } from 'chai'
+import {
+  toEther,
+  deploy,
+  fromEther,
+  deployUpgradeable,
+  getAccounts,
+  setupToken,
+} from '../utils/helpers'
+import {
+  ERC677,
+  SDLPoolMock,
+  StakingPool,
+  PriorityPool,
+  StrategyMock,
+  WithdrawalPool,
+  StakingAllowance,
+  RewardsPoolWSD,
+  StakingProxy,
+} from '../../typechain-types'
+import { ethers } from 'hardhat'
+import { StandardMerkleTree } from '@openzeppelin/merkle-tree'
+import { loadFixture } from '@nomicfoundation/hardhat-network-helpers'
+import { SDLPoolPrimary } from '../../typechain-types'
+
+describe('StakingProxy', () => {
+  async function deployFixture() {
+    const { accounts, signers } = await getAccounts()
+
+    const sdlToken = (await deploy('StakingAllowance', ['stake.link', 'SDL'])) as StakingAllowance
+    await sdlToken.mint(accounts[0], toEther(10000000))
+
+    const token = (await deploy('contracts/core/tokens/base/ERC677.sol:ERC677', [
+      'Chainlink',
+      'LINK',
+      1000000000,
+    ])) as ERC677
+    await setupToken(token, accounts, true)
+
+    const stakingPool = (await deployUpgradeable('StakingPool', [
+      token.target,
+      'Staked LINK',
+      'stLINK',
+      [],
+      toEther(10000),
+    ])) as StakingPool
+
+    const wsdToken = await deploy('WrappedSDToken', [
+      stakingPool.target,
+      'Wrapped stLINK',
+      'wstLINK',
+    ])
+
+    const strategy = (await deployUpgradeable('StrategyMock', [
+      token.target,
+      stakingPool.target,
+      toEther(1000),
+      toEther(100),
+    ])) as StrategyMock
+
+    const boostController = await deploy('LinearBoostController', [10, 4 * 365 * 86400, 4])
+
+    const sdlPool = await deployUpgradeable('SDLPoolPrimary', [
+      'Reward Escrowed SDL',
+      'reSDL',
+      sdlToken.target,
+      boostController.target,
+    ])
+
+    const rewardsPool = (await deploy('RewardsPoolWSD', [
+      sdlPool.target,
+      stakingPool.target,
+      wsdToken.target,
+    ])) as RewardsPoolWSD
+
+    const priorityPool = (await deployUpgradeable('PriorityPool', [
+      token.target,
+      stakingPool.target,
+      sdlPool.target,
+      toEther(100),
+      toEther(1000),
+      false,
+    ])) as PriorityPool
+
+    const withdrawalPool = (await deployUpgradeable('WithdrawalPool', [
+      token.target,
+      stakingPool.target,
+      priorityPool.target,
+      toEther(10),
+      0,
+    ])) as WithdrawalPool
+
+    const stakingProxy = (await deployUpgradeable('StakingProxy', [
+      token.target,
+      stakingPool.target,
+      priorityPool.target,
+      withdrawalPool.target,
+      sdlPool.target,
+      accounts[0],
+    ])) as StakingProxy
+
+    await stakingPool.addStrategy(strategy.target)
+    await stakingPool.setPriorityPool(priorityPool.target)
+    await stakingPool.setRebaseController(accounts[0])
+    await priorityPool.setDistributionOracle(accounts[0])
+    await priorityPool.setWithdrawalPool(withdrawalPool.target)
+    await sdlPool.addToken(stakingPool.target, rewardsPool.target)
+
+    await token.connect(signers[1]).approve(priorityPool.target, ethers.MaxUint256)
+    await token.approve(priorityPool.target, ethers.MaxUint256)
+    await token.approve(stakingProxy.target, ethers.MaxUint256)
+    await priorityPool.deposit(1000, false, ['0x'])
+
+    return {
+      signers,
+      accounts,
+      token,
+      stakingPool,
+      strategy,
+      sdlPool,
+      priorityPool,
+      withdrawalPool,
+      stakingProxy,
+      sdlToken,
+    }
+  }
+
+  it('deposit should work correctly', async () => {
+    const { stakingProxy, stakingPool, priorityPool, signers } = await loadFixture(deployFixture)
+
+    await expect(
+      stakingProxy.connect(signers[1]).deposit(toEther(1500), ['0x'])
+    ).to.be.revertedWithCustomError(stakingProxy, 'SenderNotAuthorized()')
+
+    await stakingProxy.deposit(toEther(1500), ['0x'])
+    assert.equal(fromEther(await stakingPool.balanceOf(stakingProxy.target)), 1000)
+    assert.equal(fromEther(await priorityPool.getQueuedTokens(stakingProxy.target, 0)), 500)
+    assert.equal(fromEther(await stakingProxy.getTotalDeposits()), 1000)
+    assert.equal(fromEther(await stakingProxy.getTotalQueuedForDeposit(0)), 500)
+  })
+
+  it('deposit should work correctly using onTokenTransfer', async () => {
+    const { stakingProxy, stakingPool, priorityPool, token, signers, accounts } = await loadFixture(
+      deployFixture
+    )
+
+    await expect(
+      stakingProxy.onTokenTransfer(
+        accounts[0],
+        toEther(1500),
+        ethers.AbiCoder.defaultAbiCoder().encode(['bytes[]'], [['0x']])
+      )
+    ).to.be.revertedWithCustomError(stakingProxy, 'InvalidToken()')
+    await expect(
+      token
+        .connect(signers[1])
+        .transferAndCall(
+          stakingProxy.target,
+          toEther(1500),
+          ethers.AbiCoder.defaultAbiCoder().encode(['bytes[]'], [['0x']])
+        )
+    ).to.be.revertedWithCustomError(stakingProxy, 'SenderNotAuthorized()')
+    await expect(
+      token.transferAndCall(
+        stakingProxy.target,
+        0,
+        ethers.AbiCoder.defaultAbiCoder().encode(['bytes[]'], [['0x']])
+      )
+    ).to.be.revertedWithCustomError(stakingProxy, 'InvalidValue()')
+
+    await token.transferAndCall(
+      stakingProxy.target,
+      toEther(1500),
+      ethers.AbiCoder.defaultAbiCoder().encode(['bytes[]'], [['0x']])
+    )
+    assert.equal(fromEther(await stakingPool.balanceOf(stakingProxy.target)), 1000)
+    assert.equal(fromEther(await priorityPool.getQueuedTokens(stakingProxy.target, 0)), 500)
+    assert.equal(fromEther(await stakingProxy.getTotalDeposits()), 1000)
+    assert.equal(fromEther(await stakingProxy.getTotalQueuedForDeposit(0)), 500)
+  })
+
+  it('claimLSTs should work correctly', async () => {
+    const { stakingProxy, stakingPool, priorityPool, strategy } = await loadFixture(deployFixture)
+
+    await stakingProxy.deposit(toEther(1500), ['0x'])
+    await strategy.setMaxDeposits(toEther(1200))
+    await priorityPool.depositQueuedTokens(0, toEther(500), ['0x'])
+
+    let data = [
+      [ethers.ZeroAddress, toEther(0), toEther(0)],
+      [stakingProxy.target, toEther(50), toEther(50)],
+    ]
+    let tree = StandardMerkleTree.of(data, ['address', 'uint256', 'uint256'])
+
+    await priorityPool.pauseForUpdate()
+    await priorityPool.updateDistribution(
+      tree.root,
+      ethers.encodeBytes32String('ipfs'),
+      toEther(50),
+      toEther(50)
+    )
+
+    assert.equal(fromEther(await stakingProxy.getTotalClaimableLSTs(toEther(50))), 50)
+
+    await stakingProxy.claimLSTs(toEther(50), toEther(50), tree.getProof(1))
+    assert.equal(fromEther(await stakingProxy.getTotalClaimableLSTs(toEther(50))), 0)
+    assert.equal(fromEther(await stakingPool.balanceOf(stakingProxy.target)), 1050)
+    assert.equal(fromEther(await stakingProxy.getTotalDeposits()), 1050)
+    assert.equal(fromEther(await stakingProxy.getTotalQueuedForDeposit(toEther(50))), 450)
+  })
+
+  it('claimLSTs should work correctly', async () => {
+    const { stakingProxy, stakingPool, priorityPool, strategy } = await loadFixture(deployFixture)
+
+    await stakingProxy.deposit(toEther(1500), ['0x'])
+    await strategy.setMaxDeposits(toEther(1200))
+    await priorityPool.depositQueuedTokens(0, toEther(500), ['0x'])
+
+    let data = [
+      [ethers.ZeroAddress, toEther(0), toEther(0)],
+      [stakingProxy.target, toEther(50), toEther(50)],
+    ]
+    let tree = StandardMerkleTree.of(data, ['address', 'uint256', 'uint256'])
+
+    await priorityPool.pauseForUpdate()
+    await priorityPool.updateDistribution(
+      tree.root,
+      ethers.encodeBytes32String('ipfs'),
+      toEther(50),
+      toEther(50)
+    )
+
+    assert.equal(fromEther(await stakingProxy.getTotalClaimableLSTs(toEther(50))), 50)
+
+    await stakingProxy.claimLSTs(toEther(50), toEther(50), tree.getProof(1))
+    assert.equal(fromEther(await stakingProxy.getTotalClaimableLSTs(toEther(50))), 0)
+    assert.equal(fromEther(await stakingPool.balanceOf(stakingProxy.target)), 1050)
+    assert.equal(fromEther(await stakingProxy.getTotalDeposits()), 1050)
+    assert.equal(fromEther(await stakingProxy.getTotalQueuedForDeposit(toEther(50))), 450)
+  })
+
+  it('withdraw should work correctly', async () => {
+    const { stakingProxy, stakingPool, priorityPool, signers, strategy, accounts, token } =
+      await loadFixture(deployFixture)
+
+    await stakingProxy.deposit(toEther(1500), ['0x'])
+    await strategy.setMaxDeposits(toEther(1050))
+    await priorityPool.depositQueuedTokens(0, toEther(50), ['0x'])
+    await priorityPool.connect(signers[1]).deposit(toEther(200), true, ['0x'])
+
+    let data = [
+      [ethers.ZeroAddress, toEther(0), toEther(0)],
+      [stakingProxy.target, toEther(50), toEther(50)],
+    ]
+    let tree = StandardMerkleTree.of(data, ['address', 'uint256', 'uint256'])
+
+    await priorityPool.pauseForUpdate()
+    await priorityPool.updateDistribution(
+      tree.root,
+      ethers.encodeBytes32String('ipfs'),
+      toEther(50),
+      toEther(50)
+    )
+
+    assert.deepEqual(
+      (await stakingProxy.getTotalWithdrawable(toEther(50))).map((d: any, i) => {
+        if (i < 2) return fromEther(d)
+        return d
+      }),
+      [650, 0, [], []]
+    )
+
+    let balance = fromEther(await token.balanceOf(accounts[0]))
+    await stakingProxy.withdraw(
+      toEther(800),
+      toEther(50),
+      toEther(50),
+      tree.getProof(1),
+      [],
+      [],
+      ['0x']
+    )
+    assert.equal(fromEther(await stakingPool.balanceOf(stakingProxy.target)), 650)
+    assert.equal(fromEther(await token.balanceOf(accounts[0])), balance + 650)
+    assert.equal(fromEther(await stakingProxy.getTotalDeposits()), 650)
+    assert.equal(fromEther(await stakingProxy.getTotalQueuedForDeposit(toEther(50))), 0)
+    assert.equal(fromEther(await stakingProxy.getTotalQueuedForWithdrawal()), 150)
+    assert.deepEqual(
+      (await stakingProxy.getTotalWithdrawable(toEther(50))).map((d: any, i) => {
+        if (i < 2) return fromEther(d)
+        return d
+      }),
+      [0, 0, [], []]
+    )
+
+    await priorityPool.deposit(toEther(100), true, ['0x'])
+
+    assert.deepEqual(
+      (await stakingProxy.getTotalWithdrawable(toEther(50))).map((d: any, i) => {
+        if (i < 2) return fromEther(d)
+        return d.map((v: any) => Number(v))
+      }),
+      [0, 100, [1], [0]]
+    )
+
+    balance = fromEther(await token.balanceOf(accounts[0]))
+    await stakingProxy.withdraw(
+      toEther(100),
+      toEther(50),
+      toEther(50),
+      tree.getProof(1),
+      [1],
+      [0],
+      ['0x']
+    )
+    assert.equal(fromEther(await stakingPool.balanceOf(stakingProxy.target)), 650)
+    assert.equal(fromEther(await token.balanceOf(accounts[0])), balance + 100)
+    assert.equal(fromEther(await stakingProxy.getTotalDeposits()), 650)
+    assert.equal(fromEther(await stakingProxy.getTotalQueuedForDeposit(toEther(50))), 0)
+    assert.equal(fromEther(await stakingProxy.getTotalQueuedForWithdrawal()), 50)
+    assert.deepEqual(
+      (await stakingProxy.getTotalWithdrawable(toEther(50))).map((d: any, i) => {
+        if (i < 2) return fromEther(d)
+        return d
+      }),
+      [0, 0, [], []]
+    )
+
+    await priorityPool.deposit(toEther(125), true, ['0x'])
+
+    assert.deepEqual(
+      (await stakingProxy.getTotalWithdrawable(toEther(50))).map((d: any, i) => {
+        if (i < 2) return fromEther(d)
+        return d.map((v: any) => Number(v))
+      }),
+      [75, 50, [1], [2]]
+    )
+
+    balance = fromEther(await token.balanceOf(accounts[0]))
+    await stakingProxy.withdraw(
+      toEther(120),
+      toEther(50),
+      toEther(50),
+      tree.getProof(1),
+      [1],
+      [2],
+      ['0x']
+    )
+    assert.equal(fromEther(await stakingPool.balanceOf(stakingProxy.target)), 580)
+    assert.equal(fromEther(await token.balanceOf(accounts[0])), balance + 120)
+    assert.equal(fromEther(await stakingProxy.getTotalDeposits()), 580)
+    assert.equal(fromEther(await stakingProxy.getTotalQueuedForDeposit(toEther(50))), 0)
+    assert.equal(fromEther(await stakingProxy.getTotalQueuedForWithdrawal()), 0)
+    assert.deepEqual(
+      (await stakingProxy.getTotalWithdrawable(toEther(50))).map((d: any, i) => {
+        if (i < 2) return fromEther(d)
+        return d
+      }),
+      [5, 0, [], []]
+    )
+  })
+
+  it('should be able to deposit reSDL tokens', async () => {
+    const { stakingProxy, sdlPool, sdlToken, accounts } = await loadFixture(deployFixture)
+
+    await sdlToken.transferAndCall(
+      sdlPool.target,
+      toEther(1000),
+      ethers.AbiCoder.defaultAbiCoder().encode(['uint256', 'uint64'], [0, 0])
+    )
+    await sdlToken.transferAndCall(
+      sdlPool.target,
+      toEther(1000),
+      ethers.AbiCoder.defaultAbiCoder().encode(['uint256', 'uint64'], [0, 0])
+    )
+    await sdlToken.transferAndCall(
+      sdlPool.target,
+      toEther(1000),
+      ethers.AbiCoder.defaultAbiCoder().encode(['uint256', 'uint64'], [0, 0])
+    )
+
+    await sdlPool.transferFrom(accounts[0], stakingProxy.target, 1)
+    await sdlPool.safeTransferFrom(accounts[0], stakingProxy.target, 3)
+
+    assert.deepEqual(await stakingProxy.getRESDLTokenIds(), [1n, 3n])
+  })
+
+  it('should be able to claim reSDL rewards', async () => {
+    const { stakingProxy, sdlPool, sdlToken, accounts, stakingPool, priorityPool, token } =
+      await loadFixture(deployFixture)
+
+    await sdlToken.transferAndCall(
+      sdlPool.target,
+      toEther(500),
+      ethers.AbiCoder.defaultAbiCoder().encode(['uint256', 'uint64'], [0, 0])
+    )
+    await sdlToken.transferAndCall(
+      sdlPool.target,
+      toEther(500),
+      ethers.AbiCoder.defaultAbiCoder().encode(['uint256', 'uint64'], [0, 0])
+    )
+    await sdlToken.transferAndCall(
+      sdlPool.target,
+      toEther(1000),
+      ethers.AbiCoder.defaultAbiCoder().encode(['uint256', 'uint64'], [0, 0])
+    )
+    await priorityPool.deposit(toEther(1000), true, ['0x'])
+    await sdlPool.transferFrom(accounts[0], stakingProxy.target, 1)
+    await sdlPool.transferFrom(accounts[0], stakingProxy.target, 2)
+    await stakingPool.transferAndCall(sdlPool.target, toEther(100), '0x')
+    await stakingProxy.claimRESDLRewards([stakingPool.target])
+
+    assert.deepEqual(
+      (await sdlPool.withdrawableRewards(stakingProxy.target)).map((d: any) => fromEther(d)),
+      [25]
+    )
+  })
+
+  it('should be able to withdraw reSDL tokens', async () => {
+    const { stakingProxy, sdlPool, sdlToken, accounts } = await loadFixture(deployFixture)
+
+    await sdlToken.transferAndCall(
+      sdlPool.target,
+      toEther(1000),
+      ethers.AbiCoder.defaultAbiCoder().encode(['uint256', 'uint64'], [0, 0])
+    )
+    await sdlToken.transferAndCall(
+      sdlPool.target,
+      toEther(1000),
+      ethers.AbiCoder.defaultAbiCoder().encode(['uint256', 'uint64'], [0, 0])
+    )
+    await sdlToken.transferAndCall(
+      sdlPool.target,
+      toEther(1000),
+      ethers.AbiCoder.defaultAbiCoder().encode(['uint256', 'uint64'], [0, 0])
+    )
+
+    await sdlPool.transferFrom(accounts[0], stakingProxy.target, 1)
+    await sdlPool.safeTransferFrom(accounts[0], stakingProxy.target, 3)
+
+    await expect(stakingProxy.withdrawRESDLToken(2, accounts[2])).to.be.revertedWithCustomError(
+      stakingProxy,
+      'InvalidTokenId()'
+    )
+
+    await stakingProxy.withdrawRESDLToken(1, accounts[2])
+    assert.deepEqual(await stakingProxy.getRESDLTokenIds(), [3n])
+    assert.equal(await sdlPool.ownerOf(1), accounts[2])
+
+    await stakingProxy.withdrawRESDLToken(3, accounts[3])
+    assert.deepEqual(await stakingProxy.getRESDLTokenIds(), [])
+    assert.equal(await sdlPool.ownerOf(3), accounts[3])
+  })
+})