smom-dbis-138/contracts/utils/FeeCollector.sol

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

import "@openzeppelin/contracts/access/AccessControl.sol";
import "@openzeppelin/contracts/token/ERC20/IERC20.sol";
import "@openzeppelin/contracts/token/ERC20/utils/SafeERC20.sol";
import "@openzeppelin/contracts/utils/ReentrancyGuard.sol";

/**
 * @title FeeCollector
 * @notice Collects and distributes protocol fees
 * @dev Supports multiple tokens and multiple fee recipients
 */
contract FeeCollector is AccessControl, ReentrancyGuard {
    using SafeERC20 for IERC20;
    
    bytes32 public constant FEE_MANAGER_ROLE = keccak256("FEE_MANAGER_ROLE");
    
    /**
     * @notice Fee recipient information
     */
    struct FeeRecipient {
        address recipient;
        uint256 shareBps; // Share in basis points (10000 = 100%)
        bool isActive;
    }
    
    mapping(address => FeeRecipient[]) private _feeRecipients; // token => recipients
    mapping(address => uint256) private _totalCollected; // token => total collected
    mapping(address => uint256) private _totalDistributed; // token => total distributed
    
    event FeesCollected(
        address indexed token,
        address indexed from,
        uint256 amount,
        uint256 timestamp
    );
    
    event FeesDistributed(
        address indexed token,
        address indexed recipient,
        uint256 amount,
        uint256 timestamp
    );
    
    event FeeRecipientAdded(
        address indexed token,
        address indexed recipient,
        uint256 shareBps,
        uint256 timestamp
    );
    
    event FeeRecipientRemoved(
        address indexed token,
        address indexed recipient,
        uint256 timestamp
    );
    
    /**
     * @notice Constructor
     * @param admin Address that will receive DEFAULT_ADMIN_ROLE
     */
    constructor(address admin) {
        _grantRole(DEFAULT_ADMIN_ROLE, admin);
        _grantRole(FEE_MANAGER_ROLE, admin);
    }
    
    /**
     * @notice Collect fees in a token
     * @param token Token address (address(0) for native ETH)
     * @param amount Amount to collect
     * @dev Can be called by anyone, typically called by contracts that collect fees
     */
    function collectFees(address token, uint256 amount) external payable nonReentrant {
        if (token == address(0)) {
            // Native ETH
            require(msg.value == amount, "FeeCollector: ETH amount mismatch");
        } else {
            // ERC20 token
            require(msg.value == 0, "FeeCollector: no ETH expected");
            IERC20(token).safeTransferFrom(msg.sender, address(this), amount);
        }
        
        _totalCollected[token] += amount;
        
        emit FeesCollected(token, msg.sender, amount, block.timestamp);
    }
    
    /**
     * @notice Distribute collected fees to recipients
     * @param token Token address (address(0) for native ETH)
     * @dev Requires FEE_MANAGER_ROLE
     */
    function distributeFees(address token) external onlyRole(FEE_MANAGER_ROLE) nonReentrant {
        FeeRecipient[] memory recipients = _feeRecipients[token];
        require(recipients.length > 0, "FeeCollector: no recipients configured");
        
        uint256 balance = token == address(0)
            ? address(this).balance
            : IERC20(token).balanceOf(address(this));
        
        require(balance > 0, "FeeCollector: no fees to distribute");
        
        uint256 totalDistributed = 0;
        
        for (uint256 i = 0; i < recipients.length; i++) {
            if (!recipients[i].isActive) continue;
            
            uint256 amount = (balance * recipients[i].shareBps) / 10000;
            
            if (token == address(0)) {
                // Native ETH
                (bool success, ) = recipients[i].recipient.call{value: amount}("");
                require(success, "FeeCollector: ETH transfer failed");
            } else {
                // ERC20 token
                IERC20(token).safeTransfer(recipients[i].recipient, amount);
            }
            
            totalDistributed += amount;
            _totalDistributed[token] += amount;
            
            emit FeesDistributed(token, recipients[i].recipient, amount, block.timestamp);
        }
        
        // Ensure we distributed exactly the balance (within rounding)
        require(totalDistributed <= balance, "FeeCollector: distribution overflow");
    }
    
    /**
     * @notice Add a fee recipient for a token
     * @param token Token address (address(0) for native ETH)
     * @param recipient Recipient address
     * @param shareBps Share in basis points (10000 = 100%)
     * @dev Requires FEE_MANAGER_ROLE
     */
    function addFeeRecipient(
        address token,
        address recipient,
        uint256 shareBps
    ) external onlyRole(FEE_MANAGER_ROLE) {
        require(recipient != address(0), "FeeCollector: zero recipient");
        require(shareBps > 0 && shareBps <= 10000, "FeeCollector: invalid share");
        
        // Check if recipient already exists
        FeeRecipient[] storage recipients = _feeRecipients[token];
        for (uint256 i = 0; i < recipients.length; i++) {
            require(recipients[i].recipient != recipient, "FeeCollector: recipient already exists");
        }
        
        recipients.push(FeeRecipient({
            recipient: recipient,
            shareBps: shareBps,
            isActive: true
        }));
        
        emit FeeRecipientAdded(token, recipient, shareBps, block.timestamp);
    }
    
    /**
     * @notice Remove a fee recipient
     * @param token Token address
     * @param recipient Recipient address to remove
     * @dev Requires FEE_MANAGER_ROLE
     */
    function removeFeeRecipient(address token, address recipient) external onlyRole(FEE_MANAGER_ROLE) {
        FeeRecipient[] storage recipients = _feeRecipients[token];
        
        for (uint256 i = 0; i < recipients.length; i++) {
            if (recipients[i].recipient == recipient) {
                recipients[i] = recipients[recipients.length - 1];
                recipients.pop();
                emit FeeRecipientRemoved(token, recipient, block.timestamp);
                return;
            }
        }
        
        revert("FeeCollector: recipient not found");
    }
    
    /**
     * @notice Get fee recipients for a token
     * @param token Token address
     * @return Array of fee recipients
     */
    function getFeeRecipients(address token) external view returns (FeeRecipient[] memory) {
        return _feeRecipients[token];
    }
    
    /**
     * @notice Get total collected fees for a token
     * @param token Token address
     * @return Total collected amount
     */
    function getTotalCollected(address token) external view returns (uint256) {
        return _totalCollected[token];
    }
    
    /**
     * @notice Get total distributed fees for a token
     * @param token Token address
     * @return Total distributed amount
     */
    function getTotalDistributed(address token) external view returns (uint256) {
        return _totalDistributed[token];
    }
    
    /**
     * @notice Get current balance for a token
     * @param token Token address (address(0) for native ETH)
     * @return Current balance
     */
    function getBalance(address token) external view returns (uint256) {
        if (token == address(0)) {
            return address(this).balance;
        } else {
            return IERC20(token).balanceOf(address(this));
        }
    }
    
    /**
     * @notice Emergency withdraw (admin only)
     * @param token Token address (address(0) for native ETH)
     * @param to Recipient address
     * @param amount Amount to withdraw
     * @dev Requires DEFAULT_ADMIN_ROLE
     */
    function emergencyWithdraw(
        address token,
        address to,
        uint256 amount
    ) external onlyRole(DEFAULT_ADMIN_ROLE) nonReentrant {
        require(to != address(0), "FeeCollector: zero recipient");
        
        if (token == address(0)) {
            (bool success, ) = to.call{value: amount}("");
            require(success, "FeeCollector: ETH transfer failed");
        } else {
            IERC20(token).safeTransfer(to, amount);
        }
    }
}