contracts-imtbl/contracts/staking/RewardRouter.sol

// SPDX-License-Identifier: MIT
pragma solidity 0.8.19;

import {IERC20} from "@openzeppelin/contracts/token/ERC20/IERC20.sol";
import {ReentrancyGuard} from "@openzeppelin/contracts/security/ReentrancyGuard.sol";
import {SafeERC20} from "@openzeppelin/contracts/token/ERC20/utils/SafeERC20.sol";

import {IRewardTracker} from "./interfaces/IRewardTracker.sol";
import {IVester} from "./interfaces/IVester.sol";
import {Governable} from "../core/Governable.sol";

contract RewardRouter is ReentrancyGuard, Governable {
  using SafeERC20 for IERC20;

  address public cec;
  address public esCec;

  address public stakedCecTracker;
  address public cecVester;

  mapping(address sender => address receiver) public pendingReceivers;

  event StakeCec(address account, address token, uint256 amount);
  event UnstakeCec(address account, address token, uint256 amount);

  constructor(address _cec, address _esCec, address _stakedCecTracker, address _cecVester) {
    cec = _cec;
    esCec = _esCec;
    stakedCecTracker = _stakedCecTracker;
    cecVester = _cecVester;
  }

  // to help users who accidentally send their tokens to this contract
  function withdrawToken(address _token, address _account, uint256 _amount) external onlyGov {
    IERC20(_token).safeTransfer(_account, _amount);
  }

  function batchStakeCecForAccount(
    address[] memory _accounts,
    uint256[] memory _amounts
  ) external nonReentrant onlyGov {
    address _cec = cec;
    for (uint256 i = 0; i < _accounts.length; i++) {
      _stakeCec(msg.sender, _accounts[i], _cec, _amounts[i]);
    }
  }

  function stakeCecForAccount(address _account, uint256 _amount) external nonReentrant onlyGov {
    _stakeCec(msg.sender, _account, cec, _amount);
  }

  function stakeCec(uint256 _amount) external nonReentrant {
    _stakeCec(msg.sender, msg.sender, cec, _amount);
  }

  function stakeEsCec(uint256 _amount) external nonReentrant {
    _stakeCec(msg.sender, msg.sender, esCec, _amount);
  }

  function unstakeCec(uint256 _amount) external nonReentrant {
    _unstakeCec(msg.sender, cec, _amount);
  }

  function unstakeEsCec(uint256 _amount) external nonReentrant {
    _unstakeCec(msg.sender, esCec, _amount);
  }

  function claim() external nonReentrant {
    address account = msg.sender;

    IRewardTracker(stakedCecTracker).claimForAccount(account, account);
  }

  function claimEsCec() external nonReentrant {
    address account = msg.sender;
    IRewardTracker(stakedCecTracker).claimForAccount(account, account);
  }

  function compound() external nonReentrant {
    _compound(msg.sender);
  }

  function compoundForAccount(address _account) external nonReentrant onlyGov {
    _compound(_account);
  }

  function handleRewards(
    bool _shouldClaimCec,
    bool _shouldStakeCec,
    bool _shouldClaimEsCec,
    bool _shouldStakeEsCec
  ) external nonReentrant {
    address account = msg.sender;

    uint256 cecAmount = 0;
    if (_shouldClaimCec) {
      cecAmount = IVester(cecVester).claimForAccount(account, account);
    }

    if (_shouldStakeCec && cecAmount > 0) {
      _stakeCec(account, account, cec, cecAmount);
    }

    uint256 esCecAmount = 0;
    if (_shouldClaimEsCec) {
      esCecAmount = IRewardTracker(stakedCecTracker).claimForAccount(account, account);
    }

    if (_shouldStakeEsCec && esCecAmount > 0) {
      _stakeCec(account, account, esCec, esCecAmount);
    }
  }

  function batchCompoundForAccounts(address[] memory _accounts) external nonReentrant onlyGov {
    for (uint256 i = 0; i < _accounts.length; i++) {
      _compound(_accounts[i]);
    }
  }

  function signalTransfer(address _receiver) external nonReentrant {
    require(IERC20(cecVester).balanceOf(msg.sender) == 0, "sender has vested tokens");

    _validateReceiver(_receiver);

    pendingReceivers[msg.sender] = _receiver;
  }

  function acceptTransfer(address _sender) external nonReentrant {
    require(IERC20(cecVester).balanceOf(_sender) == 0, "sender has vested tokens");

    address receiver = msg.sender;
    require(pendingReceivers[_sender] == receiver, "transfer not signalled");
    delete pendingReceivers[_sender];

    _validateReceiver(receiver);
    _compound(_sender);

    uint256 stakedCec = IRewardTracker(stakedCecTracker).depositBalances(_sender, cec);
    if (stakedCec > 0) {
      _unstakeCec(_sender, cec, stakedCec);
      _stakeCec(_sender, receiver, cec, stakedCec);
    }

    uint256 stakedEsCec = IRewardTracker(stakedCecTracker).depositBalances(_sender, esCec);
    if (stakedEsCec > 0) {
      _unstakeCec(_sender, esCec, stakedEsCec);
      _stakeCec(_sender, receiver, esCec, stakedEsCec);
    }

    uint256 esCecBalance = IERC20(esCec).balanceOf(_sender);
    if (esCecBalance > 0) {
      IERC20(esCec).transferFrom(_sender, receiver, esCecBalance);
    }

    IVester(cecVester).transferStakeValues(_sender, receiver);
  }

  function _validateReceiver(address _receiver) private view {
    require(
      IRewardTracker(stakedCecTracker).averageStakedAmounts(_receiver) == 0,
      "stakedCecTracker.averageStakedAmounts > 0"
    );
    require(
      IRewardTracker(stakedCecTracker).cumulativeRewards(_receiver) == 0,
      "stakedCecTracker.cumulativeRewards > 0"
    );

    require(
      IVester(cecVester).transferredAverageStakedAmounts(_receiver) == 0,
      "cecVester.transferredAverageStakedAmounts > 0"
    );
    require(
      IVester(cecVester).transferredCumulativeRewards(_receiver) == 0,
      "cecVester.transferredCumulativeRewards > 0"
    );

    require(IERC20(cecVester).balanceOf(_receiver) == 0, "cecVester.balance > 0");
  }

  function _compound(address _account) private {
    _compoundCec(_account);
  }

  function _compoundCec(address _account) private {
    uint256 esCecAmount = IRewardTracker(stakedCecTracker).claimForAccount(_account, _account);
    if (esCecAmount > 0) {
      _stakeCec(_account, _account, esCec, esCecAmount);
    }
  }

  function _stakeCec(address _fundingAccount, address _account, address _token, uint256 _amount) private {
    require(_amount > 0, "invalid _amount");

    IRewardTracker(stakedCecTracker).stakeForAccount(_fundingAccount, _account, _token, _amount);

    emit StakeCec(_account, _token, _amount);
  }

  function _unstakeCec(address _account, address _token, uint256 _amount) private {
    require(_amount > 0, "invalid _amount");
    // uint256 balance = IRewardTracker(stakedCecTracker).stakedAmounts(_account);
    IRewardTracker(stakedCecTracker).unstakeForAccount(_account, _token, _amount, _account);

    emit UnstakeCec(_account, _token, _amount);
  }
}