no_five/contracts/governance/policies/PolicyProviderConcentration.sol

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

import "../IPolicyModule.sol";
import "@openzeppelin/contracts/access/Ownable.sol";
import "../../interfaces/IFlashLoanRouter.sol";

/**
 * @title PolicyProviderConcentration
 * @notice Policy module that prevents over-concentration in single providers
 * @dev Ensures diversification across flash loan providers
 */
contract PolicyProviderConcentration is IPolicyModule, Ownable {
    string public constant override name = "ProviderConcentration";

    bool private _enabled = true;

    // Maximum percentage of total flash loans from a single provider (basis points)
    uint256 public maxProviderConcentrationBps = 5000; // 50%
    uint256 private constant BPS_SCALE = 10000;

    // Time window for concentration tracking
    uint256 public trackingWindow = 7 days;

    // Provider usage tracking
    struct ProviderUsage {
        uint256 totalVolume;
        uint256 lastResetTime;
        mapping(IFlashLoanRouter.FlashLoanProvider => uint256) providerVolumes;
    }

    mapping(address => ProviderUsage) private vaultProviderUsage; // vault => ProviderUsage
    ProviderUsage private globalProviderUsage;

    event MaxConcentrationUpdated(uint256 oldMax, uint256 newMax);
    event TrackingWindowUpdated(uint256 oldWindow, uint256 newWindow);

    modifier onlyEnabled() {
        require(_enabled, "Policy disabled");
        _;
    }

    constructor(address initialOwner) Ownable(initialOwner) {}

    /**
     * @notice Check if module is enabled
     */
    function isEnabled() external view override returns (bool) {
        return _enabled;
    }

    /**
     * @notice Enable or disable the module
     */
    function setEnabled(bool enabled) external override onlyOwner {
        _enabled = enabled;
    }

    /**
     * @notice Evaluate policy for proposed action
     * @param actionType Action type (FLASH_LOAN, etc.)
     * @param actionData Encoded action data: (vault, asset, amount, provider)
     */
    function evaluate(
        bytes32 actionType,
        bytes memory actionData
    ) external view override onlyEnabled returns (PolicyDecision memory) {
        if (actionType != keccak256("FLASH_LOAN")) {
            return PolicyDecision({
                allowed: true,
                reason: ""
            });
        }

        (
            address vault,
            address asset,
            uint256 amount,
            IFlashLoanRouter.FlashLoanProvider provider
        ) = abi.decode(actionData, (address, address, uint256, IFlashLoanRouter.FlashLoanProvider));

        // Reset usage if window expired
        ProviderUsage storage vaultUsage = vaultProviderUsage[vault];
        if (block.timestamp - vaultUsage.lastResetTime > trackingWindow) {
            // Would reset in actual implementation, but for evaluation assume fresh window
            vaultUsage = globalProviderUsage; // Use global as proxy for "reset" state
        }

        // Calculate new provider volume
        uint256 newProviderVolume = vaultUsage.providerVolumes[provider] + amount;
        uint256 newTotalVolume = vaultUsage.totalVolume + amount;

        if (newTotalVolume > 0) {
            uint256 newConcentration = (newProviderVolume * BPS_SCALE) / newTotalVolume;
            
            if (newConcentration > maxProviderConcentrationBps) {
                return PolicyDecision({
                    allowed: false,
                    reason: "Provider concentration limit exceeded"
                });
            }
        }

        return PolicyDecision({
            allowed: true,
            reason: ""
        });
    }

    /**
     * @notice Record flash loan usage
     */
    function recordUsage(
        address vault,
        address asset,
        uint256 amount,
        IFlashLoanRouter.FlashLoanProvider provider
    ) external {
        // Reset if window expired
        ProviderUsage storage vaultUsage = vaultProviderUsage[vault];
        if (block.timestamp - vaultUsage.lastResetTime > trackingWindow) {
            _resetUsage(vault);
            vaultUsage = vaultProviderUsage[vault];
        }

        // Update usage
        vaultUsage.providerVolumes[provider] += amount;
        vaultUsage.totalVolume += amount;

        // Update global usage
        ProviderUsage storage global = globalProviderUsage;
        if (block.timestamp - global.lastResetTime > trackingWindow) {
            _resetGlobalUsage();
            global = globalProviderUsage;
        }
        global.providerVolumes[provider] += amount;
        global.totalVolume += amount;
    }

    /**
     * @notice Reset usage for a vault
     */
    function _resetUsage(address vault) internal {
        ProviderUsage storage usage = vaultProviderUsage[vault];
        usage.totalVolume = 0;
        usage.lastResetTime = block.timestamp;
        
        // Reset all provider volumes
        for (uint256 i = 0; i <= uint256(IFlashLoanRouter.FlashLoanProvider.DAI_FLASH); i++) {
            usage.providerVolumes[IFlashLoanRouter.FlashLoanProvider(i)] = 0;
        }
    }

    /**
     * @notice Reset global usage
     */
    function _resetGlobalUsage() internal {
        globalProviderUsage.totalVolume = 0;
        globalProviderUsage.lastResetTime = block.timestamp;
        
        for (uint256 i = 0; i <= uint256(IFlashLoanRouter.FlashLoanProvider.DAI_FLASH); i++) {
            globalProviderUsage.providerVolumes[IFlashLoanRouter.FlashLoanProvider(i)] = 0;
        }
    }

    /**
     * @notice Update maximum provider concentration
     */
    function setMaxConcentration(uint256 newMaxBps) external onlyOwner {
        require(newMaxBps <= BPS_SCALE, "Invalid concentration");
        uint256 oldMax = maxProviderConcentrationBps;
        maxProviderConcentrationBps = newMaxBps;
        emit MaxConcentrationUpdated(oldMax, newMaxBps);
    }

    /**
     * @notice Update tracking window
     */
    function setTrackingWindow(uint256 newWindow) external onlyOwner {
        require(newWindow > 0, "Invalid window");
        uint256 oldWindow = trackingWindow;
        trackingWindow = newWindow;
        emit TrackingWindowUpdated(oldWindow, newWindow);
    }

    /**
     * @notice Get provider concentration for a vault
     */
    function getProviderConcentration(
        address vault,
        IFlashLoanRouter.FlashLoanProvider provider
    ) external view returns (uint256 concentrationBps) {
        ProviderUsage storage usage = vaultProviderUsage[vault];
        if (usage.totalVolume == 0) {
            return 0;
        }
        return (usage.providerVolumes[provider] * BPS_SCALE) / usage.totalVolume;
    }
}