contracts-imtbl/contracts/staking/RewardTracker.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 {Governable} from "../core/Governable.sol";

contract RewardTracker is ReentrancyGuard, IRewardTracker, Governable {
  using SafeERC20 for IERC20;

  mapping(address token => bool status) public isDepositToken;
  mapping(address token => uint256 amount) public totalDepositSupply;

  mapping(address account => uint256 amount) public override stakedAmounts;
  mapping(address account => uint256 amount) public claimableReward;
  mapping(address account => uint256 amount) public override cumulativeRewards;

  mapping(address account => uint256 time) public lastUpdateTimes;

  uint256 public tokensPerInterval;
  address public rewardToken;
  uint8 public decimals = 18;

  bool public inPrivateStakingMode;
  bool public inPrivateClaimingMode;
  mapping(address handler => bool status) public isHandler;

  event Claim(address receiver, uint256 amount);
  event TokensPerIntervalChange(uint256 amount);

  constructor(address _rewardToken, address _depositToken, uint256 _tokensPerInterval, uint8 _decimals) {
    rewardToken = _rewardToken;
    isDepositToken[_depositToken] = true;
    tokensPerInterval = _tokensPerInterval;
    decimals = _decimals;
  }

  function setDepositToken(address _depositToken, bool _isDepositToken) external onlyGov {
    isDepositToken[_depositToken] = _isDepositToken;
  }
  //TODO:: 是否更新用户未领取的奖励
  function setTokensPerInterval(uint256 _amount) external onlyGov {
    tokensPerInterval = _amount;
    emit TokensPerIntervalChange(_amount);
  }

  function setInPrivateStakingMode(bool _inPrivateStakingMode) external onlyGov {
    inPrivateStakingMode = _inPrivateStakingMode;
  }

  function setInPrivateClaimingMode(bool _inPrivateClaimingMode) external onlyGov {
    inPrivateClaimingMode = _inPrivateClaimingMode;
  }

  function setHandler(address _handler, bool _isActive) external onlyGov {
    isHandler[_handler] = _isActive;
  }

  // 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 stake(address _depositToken, uint256 _amount) external override nonReentrant {
    if (inPrivateStakingMode) {
      revert("RewardTracker: action not enabled");
    }
    _stake(msg.sender, msg.sender, _depositToken, _amount);
  }

  function stakeForAccount(
    address _fundingAccount,
    address _account,
    address _depositToken,
    uint256 _amount
  ) external override nonReentrant {
    _validateHandler();
    _stake(_fundingAccount, _account, _depositToken, _amount);
  }

  function unstake(address _depositToken, uint256 _amount) external override nonReentrant {
    if (inPrivateStakingMode) {
      revert("RewardTracker: action not enabled");
    }
    _unstake(msg.sender, _depositToken, _amount, msg.sender);
  }

  function unstakeForAccount(
    address _account,
    address _depositToken,
    uint256 _amount,
    address _receiver
  ) external override nonReentrant {
    _validateHandler();
    _unstake(_account, _depositToken, _amount, _receiver);
  }

  function claim(address _receiver) external override nonReentrant returns (uint256) {
    if (inPrivateClaimingMode) {
      revert("RewardTracker: action not enabled");
    }
    return _claim(msg.sender, _receiver);
  }

  function claimForAccount(address _account, address _receiver) external override nonReentrant returns (uint256) {
    _validateHandler();
    return _claim(_account, _receiver);
  }

  function claimable(address _account) public view override returns (uint256) {
    uint256 stakedAmount = stakedAmounts[_account];
    if (stakedAmount == 0) {
      return claimableReward[_account];
    }
    uint256 pendingRewards = _getNextClaimableAmount(_account);
    return claimableReward[_account] + pendingRewards;
  }

  function _claim(address _account, address _receiver) private returns (uint256) {
    _updateRewards(_account);
    uint256 tokenAmount = claimableReward[_account];
    claimableReward[_account] = 0;
    if (tokenAmount > 0) {
      IERC20(rewardToken).safeTransfer(_receiver, tokenAmount);
      emit Claim(_account, tokenAmount);
    }
    return tokenAmount;
  }

  function _validateHandler() private view {
    require(isHandler[msg.sender], "RewardTracker: forbidden");
  }

  function _stake(address _fundingAccount, address _account, address _depositToken, uint256 _amount) private {
    require(_amount > 0, "RewardTracker: invalid _amount");
    require(isDepositToken[_depositToken], "RewardTracker: invalid _depositToken");
    IERC20(_depositToken).safeTransferFrom(_fundingAccount, address(this), _amount);
    _updateRewards(_account);
    stakedAmounts[_account] = stakedAmounts[_account] + _amount;
    totalDepositSupply[_depositToken] = totalDepositSupply[_depositToken] + _amount;
  }

  function _unstake(address _account, address _depositToken, uint256 _amount, address _receiver) private {
    require(_amount > 0, "RewardTracker: invalid _amount");
    require(isDepositToken[_depositToken], "RewardTracker: invalid _depositToken");
    _updateRewards(_account);
    uint256 stakedAmount = stakedAmounts[_account];
    require(stakedAmount >= _amount, "RewardTracker: _amount exceeds stakedAmount");
    stakedAmounts[_account] = stakedAmount - _amount;
    totalDepositSupply[_depositToken] = totalDepositSupply[_depositToken] - _amount;
    IERC20(_depositToken).safeTransfer(_receiver, _amount);
  }

  function updateRewards(address _account) public {
    _updateRewards(_account);
  }

  function _updateRewards(address _account) private {
    uint256 accountReward = _getNextClaimableAmount(_account);
    lastUpdateTimes[_account] = block.timestamp;
    uint256 _claimableReward = claimableReward[_account] + accountReward;
    claimableReward[_account] = _claimableReward;
    if (_claimableReward > 0 && stakedAmounts[_account] > 0) {
      uint256 nextCumulativeReward = cumulativeRewards[_account] + accountReward;
      cumulativeRewards[_account] = nextCumulativeReward;
    }
  }

  function _getNextClaimableAmount(address _account) private view returns (uint256) {
    uint256 timeDiff = block.timestamp - lastUpdateTimes[_account];
    uint256 stakedAmount = stakedAmounts[_account];
    if (stakedAmount == 0) {
      return 0;
    }
    uint256 accountReward = (stakedAmount / (10 ** decimals)) * tokensPerInterval * timeDiff;
    return accountReward;
  }
}