Treasury

[
  {
    "inputs": [],
    "name": "LZ_Treasury_LzTokenNotEnabled",
    "type": "error"
  },
  {
    "inputs": [
      {
        "internalType": "address",
        "name": "_to",
        "type": "address"
      },
      {
        "internalType": "uint256",
        "name": "_value",
        "type": "uint256"
      }
    ],
    "name": "Transfer_NativeFailed",
    "type": "error"
  },
  {
    "inputs": [],
    "name": "Transfer_ToAddressIsZero",
    "type": "error"
  },
  {
    "anonymous": false,
    "inputs": [
      {
        "indexed": true,
        "internalType": "address",
        "name": "previousOwner",
        "type": "address"
      },
      {
        "indexed": true,
        "internalType": "address",
        "name": "newOwner",
        "type": "address"
      }
    ],
    "name": "OwnershipTransferred",
    "type": "event"
  },
  {
    "inputs": [
      {
        "internalType": "address",
        "name": "",
        "type": "address"
      },
      {
        "internalType": "uint32",
        "name": "",
        "type": "uint32"
      },
      {
        "internalType": "uint256",
        "name": "_totalFee",
        "type": "uint256"
      },
      {
        "internalType": "bool",
        "name": "_payInLzToken",
        "type": "bool"
      }
    ],
    "name": "getFee",
    "outputs": [
      {
        "internalType": "uint256",
        "name": "",
        "type": "uint256"
      }
    ],
    "stateMutability": "view",
    "type": "function"
  },
  {
    "inputs": [],
    "name": "lzTokenEnabled",
    "outputs": [
      {
        "internalType": "bool",
        "name": "",
        "type": "bool"
      }
    ],
    "stateMutability": "view",
    "type": "function"
  },
  {
    "inputs": [],
    "name": "lzTokenFee",
    "outputs": [
      {
        "internalType": "uint256",
        "name": "",
        "type": "uint256"
      }
    ],
    "stateMutability": "view",
    "type": "function"
  },
  {
    "inputs": [],
    "name": "nativeBP",
    "outputs": [
      {
        "internalType": "uint256",
        "name": "",
        "type": "uint256"
      }
    ],
    "stateMutability": "view",
    "type": "function"
  },
  {
    "inputs": [],
    "name": "owner",
    "outputs": [
      {
        "internalType": "address",
        "name": "",
        "type": "address"
      }
    ],
    "stateMutability": "view",
    "type": "function"
  },
  {
    "inputs": [
      {
        "internalType": "address",
        "name": "",
        "type": "address"
      },
      {
        "internalType": "uint32",
        "name": "",
        "type": "uint32"
      },
      {
        "internalType": "uint256",
        "name": "_totalFee",
        "type": "uint256"
      },
      {
        "internalType": "bool",
        "name": "_payInLzToken",
        "type": "bool"
      }
    ],
    "name": "payFee",
    "outputs": [
      {
        "internalType": "uint256",
        "name": "",
        "type": "uint256"
      }
    ],
    "stateMutability": "payable",
    "type": "function"
  },
  {
    "inputs": [],
    "name": "renounceOwnership",
    "outputs": [],
    "stateMutability": "nonpayable",
    "type": "function"
  },
  {
    "inputs": [
      {
        "internalType": "bool",
        "name": "_lzTokenEnabled",
        "type": "bool"
      }
    ],
    "name": "setLzTokenEnabled",
    "outputs": [],
    "stateMutability": "nonpayable",
    "type": "function"
  },
  {
    "inputs": [
      {
        "internalType": "uint256",
        "name": "_lzTokenFee",
        "type": "uint256"
      }
    ],
    "name": "setLzTokenFee",
    "outputs": [],
    "stateMutability": "nonpayable",
    "type": "function"
  },
  {
    "inputs": [
      {
        "internalType": "uint256",
        "name": "_nativeBP",
        "type": "uint256"
      }
    ],
    "name": "setNativeFeeBP",
    "outputs": [],
    "stateMutability": "nonpayable",
    "type": "function"
  },
  {
    "inputs": [
      {
        "internalType": "address",
        "name": "newOwner",
        "type": "address"
      }
    ],
    "name": "transferOwnership",
    "outputs": [],
    "stateMutability": "nonpayable",
    "type": "function"
  },
  {
    "inputs": [
      {
        "internalType": "address",
        "name": "_messageLib",
        "type": "address"
      },
      {
        "internalType": "address",
        "name": "_lzToken",
        "type": "address"
      },
      {
        "internalType": "address",
        "name": "_to",
        "type": "address"
      },
      {
        "internalType": "uint256",
        "name": "_amount",
        "type": "uint256"
      }
    ],
    "name": "withdrawLzToken",
    "outputs": [],
    "stateMutability": "nonpayable",
    "type": "function"
  },
  {
    "inputs": [
      {
        "internalType": "address",
        "name": "_messageLib",
        "type": "address"
      },
      {
        "internalType": "address payable",
        "name": "_to",
        "type": "address"
      },
      {
        "internalType": "uint256",
        "name": "_amount",
        "type": "uint256"
      }
    ],
    "name": "withdrawNativeFee",
    "outputs": [],
    "stateMutability": "nonpayable",
    "type": "function"
  },
  {
    "inputs": [
      {
        "internalType": "address",
        "name": "_token",
        "type": "address"
      },
      {
        "internalType": "address",
        "name": "_to",
        "type": "address"
      },
      {
        "internalType": "uint256",
        "name": "_amount",
        "type": "uint256"
      }
    ],
    "name": "withdrawToken",
    "outputs": [],
    "stateMutability": "nonpayable",
    "type": "function"
  }
]

// SPDX-License-Identifier: LZBL-1.2

pragma solidity ^0.8.20;

import { Ownable } from "@openzeppelin/contracts/access/Ownable.sol";

import { ISendLib } from "@layerzerolabs/lz-evm-protocol-v2/contracts/interfaces/ISendLib.sol";
import { Transfer } from "@layerzerolabs/lz-evm-protocol-v2/contracts/libs/Transfer.sol";

import { ILayerZeroTreasury } from "./interfaces/ILayerZeroTreasury.sol";

contract Treasury is Ownable, ILayerZeroTreasury {
    uint256 public nativeBP;
    uint256 public lzTokenFee;
    bool public lzTokenEnabled;

    error LZ_Treasury_LzTokenNotEnabled();

    function getFee(
        address /*_sender*/,
        uint32 /*_eid*/,
        uint256 _totalFee,
        bool _payInLzToken
    ) external view override returns (uint256) {
        return _getFee(_totalFee, _payInLzToken);
    }

    function payFee(
        address /*_sender*/,
        uint32 /*_eid*/,
        uint256 _totalFee,
        bool _payInLzToken
    ) external payable override returns (uint256) {
        return _getFee(_totalFee, _payInLzToken);
    }

    function setLzTokenEnabled(bool _lzTokenEnabled) external onlyOwner {
        lzTokenEnabled = _lzTokenEnabled;
    }

    function setNativeFeeBP(uint256 _nativeBP) external onlyOwner {
        nativeBP = _nativeBP;
    }

    function setLzTokenFee(uint256 _lzTokenFee) external onlyOwner {
        lzTokenFee = _lzTokenFee;
    }

    function withdrawLzToken(address _messageLib, address _lzToken, address _to, uint256 _amount) external onlyOwner {
        ISendLib(_messageLib).withdrawLzTokenFee(_lzToken, _to, _amount);
    }

    function withdrawNativeFee(address _messageLib, address payable _to, uint256 _amount) external onlyOwner {
        ISendLib(_messageLib).withdrawFee(_to, _amount);
    }

    // this is for withdrawing lz token sent to this contract by uln301 and fee handler
    // and to withdraw any native sent over via payFee
    function withdrawToken(address _token, address _to, uint256 _amount) external onlyOwner {
        // transfers native if _token is address(0x0)
        Transfer.nativeOrToken(_token, _to, _amount);
    }

    // ======================= Internal =======================

    function _getFee(uint256 _totalFee, bool _payInLzToken) internal view returns (uint256) {
        if (_payInLzToken) {
            if (!lzTokenEnabled) revert LZ_Treasury_LzTokenNotEnabled();
            return lzTokenFee;
        } else {
            return (_totalFee * nativeBP) / 10000;
        }
    }
}