assetxContracts/contracts/ytLp/core/YTRewardRouter.sol

// SPDX-License-Identifier: MIT
pragma solidity ^0.8.0;

import "@openzeppelin/contracts-upgradeable/utils/ReentrancyGuardUpgradeable.sol";
import "@openzeppelin/contracts-upgradeable/proxy/utils/Initializable.sol";
import "@openzeppelin/contracts-upgradeable/proxy/utils/UUPSUpgradeable.sol";
import "@openzeppelin/contracts-upgradeable/utils/PausableUpgradeable.sol";
import "@openzeppelin/contracts/token/ERC20/IERC20.sol";
import "@openzeppelin/contracts/token/ERC20/utils/SafeERC20.sol";
import "../../interfaces/IYTPoolManager.sol";
import "../../interfaces/IYTVault.sol";

contract YTRewardRouter is Initializable, UUPSUpgradeable, ReentrancyGuardUpgradeable, PausableUpgradeable {
    using SafeERC20 for IERC20;
    
    /// @custom:oz-upgrades-unsafe-allow constructor
    constructor() {
        _disableInitializers();
    }
    
    error Forbidden();
    error AlreadyInitialized();
    error InvalidAddress();
    error InvalidAmount();
    error InsufficientOutput();
    
    address public gov;
    address public usdy;
    address public ytLP;
    address public ytPoolManager;
    address public ytVault;
    
    event Swap(
        address indexed account,
        address tokenIn,
        address tokenOut,
        uint256 amountIn,
        uint256 amountOut
    );
    
    modifier onlyGov() {
        if (msg.sender != gov) revert Forbidden();
        _;
    }
    
    function initialize(
        address _usdy,
        address _ytLP,
        address _ytPoolManager,
        address _ytVault
    ) external initializer {
        if (_usdy == address(0)) revert InvalidAddress();
        if (_ytLP == address(0)) revert InvalidAddress();
        if (_ytPoolManager == address(0)) revert InvalidAddress();
        if (_ytVault == address(0)) revert InvalidAddress();
        
        __ReentrancyGuard_init();
        __UUPSUpgradeable_init();
        __Pausable_init();
        
        gov = msg.sender;

        usdy = _usdy;
        ytLP = _ytLP;
        ytPoolManager = _ytPoolManager;
        ytVault = _ytVault;
    }
    
    function _authorizeUpgrade(address newImplementation) internal override onlyGov {}
    
    function pause() external onlyGov {
        _pause();
    }
    
    function unpause() external onlyGov {
        _unpause();
    }
    
    function addLiquidity(
        address _token,
        uint256 _amount,
        uint256 _minUsdy,
        uint256 _minYtLP
    ) external nonReentrant whenNotPaused returns (uint256) {
        if (_amount == 0) revert InvalidAmount();
        
        address account = msg.sender;
        
        IERC20(_token).safeTransferFrom(account, address(this), _amount);
        IERC20(_token).approve(ytPoolManager, _amount);
        
        uint256 ytLPAmount = IYTPoolManager(ytPoolManager).addLiquidityForAccount(
            address(this),
            account,
            _token,
            _amount,
            _minUsdy,
            _minYtLP
        );
        
        return ytLPAmount;
    }
    
    function removeLiquidity(
        address _tokenOut,
        uint256 _ytLPAmount,
        uint256 _minOut,
        address _receiver
    ) external nonReentrant whenNotPaused returns (uint256) {
        if (_ytLPAmount == 0) revert InvalidAmount();
        
        address account = msg.sender;
        
        uint256 amountOut = IYTPoolManager(ytPoolManager).removeLiquidityForAccount(
            account,
            _tokenOut,
            _ytLPAmount,
            _minOut,
            _receiver
        );
        
        return amountOut;
    }
    
    function swapYT(
        address _tokenIn,
        address _tokenOut,
        uint256 _amountIn,
        uint256 _minOut,
        address _receiver
    ) external nonReentrant whenNotPaused returns (uint256) {
        if (_amountIn == 0) revert InvalidAmount();
        
        address account = msg.sender;
        
        IERC20(_tokenIn).safeTransferFrom(account, ytVault, _amountIn);
        
        uint256 amountOut = IYTVault(ytVault).swap(_tokenIn, _tokenOut, _receiver);
        
        if (amountOut < _minOut) revert InsufficientOutput();
        
        emit Swap(account, _tokenIn, _tokenOut, _amountIn, amountOut);
        
        return amountOut;
    }
    
    function getYtLPPrice() external view returns (uint256) {
        return IYTPoolManager(ytPoolManager).getPrice(true);
    }
    
    function getAccountValue(address _account) external view returns (uint256) {
        uint256 ytLPBalance = IERC20(ytLP).balanceOf(_account);
        uint256 ytLPPrice = IYTPoolManager(ytPoolManager).getPrice(true);
        return ytLPBalance * ytLPPrice / (10 ** 18);
    }
    
    uint256[50] private __gap;
}