Fix RewardStreams integration

foundry.toml

@@ -4,7 +4,7 @@ out = "out"
 libs = ["lib"]
 test = 'test'
 optimizer = true
-optimizer_runs = 1000
+optimizer_runs = 500
 # solc = "0.8.0"
 gas_reports = ["*"]
 fs_permissions = [{ access = "read", path = "./"}]

src/AggregationLayerVault.sol

@@ -98,8 +98,7 @@ contract AggregationLayerVault is
         });
         $.totalAllocationPoints = _initParams.initialCashAllocationPoints;
         $.evc = _initParams.evc;
-
-        _setBalanceTracker(_initParams.balanceTracker);
+        $.balanceTracker = _initParams.balanceTracker;
 
         // Setup DEFAULT_ADMIN
         _grantRole(DEFAULT_ADMIN_ROLE, _initParams.aggregationVaultOwner);
@@ -114,92 +113,83 @@ contract AggregationLayerVault is
         _setRoleAdmin(REBALANCER, REBALANCER_ADMIN);
     }
 
-    /// @notice Set performance fee recipient address
-    /// @notice @param _newFeeRecipient Recipient address
+    /// @dev See {FeeModule-setFeeRecipient}.
     function setFeeRecipient(address _newFeeRecipient) external onlyRole(AGGREGATION_VAULT_MANAGER) use(MODULE_FEE) {}
 
-    /// @notice Set performance fee (1e18 == 100%)
-    /// @notice @param _newFee Fee rate
+    /// @dev See {FeeModule-setPerformanceFee}.
     function setPerformanceFee(uint256 _newFee) external onlyRole(AGGREGATION_VAULT_MANAGER) use(MODULE_FEE) {}
 
-    /// @notice Opt in to strategy rewards
-    /// @param _strategy Strategy address
+    /// @dev See {RewardsModule-optInStrategyRewards}.
     function optInStrategyRewards(address _strategy)
         external
         override
         onlyRole(AGGREGATION_VAULT_MANAGER)
         use(MODULE_REWARDS)
     {}
 
-    /// @notice Opt out of strategy rewards
-    /// @param _strategy Strategy address
+    /// @dev See {RewardsModule-optOutStrategyRewards}.
     function optOutStrategyRewards(address _strategy)
         external
         override
         onlyRole(AGGREGATION_VAULT_MANAGER)
         use(MODULE_REWARDS)
     {}
 
-    /// @notice Claim a specific strategy rewards
-    /// @param _strategy Strategy address.
-    /// @param _reward The address of the reward token.
-    /// @param _recipient The address to receive the claimed reward tokens.
-    /// @param _forfeitRecentReward Whether to forfeit the recent rewards and not update the accumulator.
+    /// @dev See {RewardsModule-optOutStrategyRewards}.
+    function enableRewardForStrategy(address _strategy, address _reward)
+        external
+        override
+        onlyRole(AGGREGATION_VAULT_MANAGER)
+        use(MODULE_REWARDS)
+    {}
+
+    /// @dev See {RewardsModule-disableRewardForStrategy}.
+    function disableRewardForStrategy(address _strategy, address _reward, bool _forfeitRecentReward)
+        external
+        override
+        onlyRole(AGGREGATION_VAULT_MANAGER)
+        use(MODULE_REWARDS)
+    {}
+
+    /// @dev See {RewardsModule-claimStrategyReward}.
     function claimStrategyReward(address _strategy, address _reward, address _recipient, bool _forfeitRecentReward)
         external
         override
         onlyRole(AGGREGATION_VAULT_MANAGER)
         use(MODULE_REWARDS)
     {}
 
-    /// @notice Enables balance forwarding for sender
-    /// @dev Should call the IBalanceTracker hook with the current user's balance
+    /// @dev See {RewardsModule-enableBalanceForwarder}.
     function enableBalanceForwarder() external override use(MODULE_REWARDS) {}
 
-    /// @notice Disables balance forwarding for the sender
-    /// @dev Should call the IBalanceTracker hook with the account's balance of 0
+    /// @dev See {RewardsModule-disableBalanceForwarder}.
     function disableBalanceForwarder() external override use(MODULE_REWARDS) {}
 
-    /// @notice Adjust a certain strategy's allocation points.
-    /// @dev Can only be called by an address that have the ALLOCATIONS_MANAGER
-    /// @param _strategy address of strategy
-    /// @param _newPoints new strategy's points
+    /// @dev See {AllocationPointsModule-adjustAllocationPoints}.
     function adjustAllocationPoints(address _strategy, uint256 _newPoints)
         external
         use(MODULE_ALLOCATION_POINTS)
         onlyRole(ALLOCATIONS_MANAGER)
     {}
 
-    /// @notice Set cap on strategy allocated amount.
-    /// @dev By default, cap is set to 0, not activated.
-    /// @param _strategy Strategy address.
-    /// @param _cap Cap amount
+    /// @dev See {AllocationPointsModule-setStrategyCap}.
     function setStrategyCap(address _strategy, uint256 _cap)
         external
         use(MODULE_ALLOCATION_POINTS)
         onlyRole(ALLOCATIONS_MANAGER)
     {}
 
-    /// @notice Add new strategy with it's allocation points.
-    /// @dev Can only be called by an address that have STRATEGY_ADDER.
-    /// @param _strategy Address of the strategy
-    /// @param _allocationPoints Strategy's allocation points
+    /// @dev See {AllocationPointsModule-addStrategy}.
     function addStrategy(address _strategy, uint256 _allocationPoints)
         external
         use(MODULE_ALLOCATION_POINTS)
         onlyRole(STRATEGY_ADDER)
     {}
 
-    /// @notice Remove strategy and set its allocation points to zero.
-    /// @dev This function does not pull funds, `harvest()` needs to be called to withdraw
-    /// @dev Can only be called by an address that have the STRATEGY_REMOVER
-    /// @param _strategy Address of the strategy
+    /// @dev See {AllocationPointsModule-removeStrategy}.
     function removeStrategy(address _strategy) external use(MODULE_ALLOCATION_POINTS) onlyRole(STRATEGY_REMOVER) {}
 
-    /// @notice Set hooks contract and hooked functions.
-    /// @dev This funtion should be overriden to implement access control.
-    /// @param _hooksTarget Hooks contract.
-    /// @param _hookedFns Hooked functions.
+    /// @dev See {HooksModule-setHooksConfig}.
     function setHooksConfig(address _hooksTarget, uint32 _hookedFns)
         external
         override
@@ -423,8 +413,7 @@ contract AggregationLayerVault is
 
     /// @dev Withdraw asset back to the user.
     /// @dev See {IERC4626-_withdraw}.
-    /// @dev if the cash reserve can not cover the amount to withdraw, this function will loop through the strategies
-    ///      to cover the remaining amount. This function will revert if the amount to withdraw is not available
+    /// @dev This function call WithdrawalQueue.callWithdrawalQueue() that should handle the rest of the withdraw execution flow.
     function _withdraw(address caller, address receiver, address owner, uint256 assets, uint256 shares)
         internal
         override

src/AggregationLayerVaultFactory.sol

@@ -1,49 +1,51 @@
 // SPDX-License-Identifier: GPL-2.0-or-later
 pragma solidity ^0.8.0;
 
-import {Clones} from "@openzeppelin/contracts/proxy/Clones.sol";
-// core modules
+// contracts
 import {Rewards} from "./module/Rewards.sol";
 import {Hooks} from "./module/Hooks.sol";
 import {Fee} from "./module/Fee.sol";
-// core plugins
 import {WithdrawalQueue} from "./plugin/WithdrawalQueue.sol";
 import {AggregationLayerVault} from "./AggregationLayerVault.sol";
+// libs
+import {Clones} from "@openzeppelin/contracts/proxy/Clones.sol";
 
 contract AggregationLayerVaultFactory {
     /// core dependencies
     address public immutable evc;
     address public immutable balanceTracker;
-    /// core modules
+    /// core modules implementations addresses
     address public immutable rewardsModuleImpl;
     address public immutable hooksModuleImpl;
     address public immutable feeModuleImpl;
-    address public immutable allocationpointsModuleImpl;
+    address public immutable allocationPointsModuleImpl;
     /// plugins
-    /// @dev Rebalancer periphery, one instance can serve different aggregation vaults
-    address public immutable rebalancerAddr;
-    /// @dev Withdrawal queue perihperhy, need to be deployed per aggregation vault
-    address public immutable withdrawalQueueAddr;
+    /// @dev Rebalancer plugin contract address, one instance can serve different aggregation vaults
+    address public immutable rebalancer;
+    /// @dev WithdrawalQueue plugin implementation address, need to be deployed per aggregation vault
+    address public immutable withdrawalQueueImpl;
+
+    struct FactoryParams {
+        address evc;
+        address balanceTracker;
+        address rewardsModuleImpl;
+        address hooksModuleImpl;
+        address feeModuleImpl;
+        address allocationPointsModuleImpl;
+        address rebalancer;
+        address withdrawalQueueImpl;
+    }
 
-    constructor(
-        address _evc,
-        address _balanceTracker,
-        address _rewardsModuleImpl,
-        address _hooksModuleImpl,
-        address _feeModuleImpl,
-        address _allocationpointsModuleImpl,
-        address _rebalancerAddr,
-        address _withdrawalQueueAddr
-    ) {
-        evc = _evc;
-        balanceTracker = _balanceTracker;
-        rewardsModuleImpl = _rewardsModuleImpl;
-        hooksModuleImpl = _hooksModuleImpl;
-        feeModuleImpl = _feeModuleImpl;
-        allocationpointsModuleImpl = _allocationpointsModuleImpl;
+    constructor(FactoryParams memory _factoryParams) {
+        evc = _factoryParams.evc;
+        balanceTracker = _factoryParams.balanceTracker;
+        rewardsModuleImpl = _factoryParams.rewardsModuleImpl;
+        hooksModuleImpl = _factoryParams.hooksModuleImpl;
+        feeModuleImpl = _factoryParams.feeModuleImpl;
+        allocationPointsModuleImpl = _factoryParams.allocationPointsModuleImpl;
 
-        rebalancerAddr = _rebalancerAddr;
-        withdrawalQueueAddr = _withdrawalQueueAddr;
+        rebalancer = _factoryParams.rebalancer;
+        withdrawalQueueImpl = _factoryParams.withdrawalQueueImpl;
     }
 
     function deployEulerAggregationLayer(
@@ -56,10 +58,10 @@ contract AggregationLayerVaultFactory {
         address rewardsModuleAddr = Clones.clone(rewardsModuleImpl);
         address hooksModuleAddr = Clones.clone(hooksModuleImpl);
         address feeModuleAddr = Clones.clone(feeModuleImpl);
-        address allocationpointsModuleAddr = Clones.clone(allocationpointsModuleImpl);
+        address allocationpointsModuleAddr = Clones.clone(allocationPointsModuleImpl);
 
         // cloning plugins
-        WithdrawalQueue withdrawalQueue = WithdrawalQueue(Clones.clone(withdrawalQueueAddr));
+        WithdrawalQueue withdrawalQueue = WithdrawalQueue(Clones.clone(withdrawalQueueImpl));
 
         // deploy new aggregation vault
         AggregationLayerVault aggregationLayerVault =
@@ -69,7 +71,7 @@ contract AggregationLayerVaultFactory {
             evc: evc,
             balanceTracker: balanceTracker,
             withdrawalQueuePeriphery: address(withdrawalQueue),
-            rebalancerPerihpery: rebalancerAddr,
+            rebalancerPerihpery: rebalancer,
             aggregationVaultOwner: msg.sender,
             asset: _asset,
             name: _name,

src/interface/IWithdrawalQueue.sol

@@ -16,5 +16,4 @@ interface IWithdrawalQueue {
     function reorderWithdrawalQueue(uint8 _index1, uint8 _index2) external;
 
     function withdrawalQueueLength() external view returns (uint256);
-    function withdrawalQueue(uint256 _index) external view returns (address);
 }

src/lib/EventsLib.sol

@@ -6,12 +6,12 @@ library EventsLib {
     event Gulp(uint256 interestLeft, uint256 interestSmearEnd);
     event Harvest(address indexed strategy, uint256 strategyBalanceAmount, uint256 strategyAllocatedAmount);
     event AccruePerformanceFee(address indexed feeRecipient, uint256 yield, uint256 feeAssets);
-    event Rebalance(address indexed strategy, uint256 _amountToRebalance, bool _isDeposit);
+    event Rebalance(address indexed strategy, uint256 amountToRebalance, bool isDeposit);
 
     /// @dev Allocationpoints events
     event AdjustAllocationPoints(address indexed strategy, uint256 oldPoints, uint256 newPoints);
     event AddStrategy(address indexed strategy, uint256 allocationPoints);
-    event RemoveStrategy(address indexed _strategy);
+    event RemoveStrategy(address indexed strategy);
     event SetStrategyCap(address indexed strategy, uint256 cap);
 
     /// @dev Fee events
@@ -24,6 +24,8 @@ library EventsLib {
     /// @dev Rewards events
     event OptInStrategyRewards(address indexed strategy);
     event OptOutStrategyRewards(address indexed strategy);
-    event EnableBalanceForwarder(address indexed _user);
-    event DisableBalanceForwarder(address indexed _user);
+    event EnableBalanceForwarder(address indexed user);
+    event DisableBalanceForwarder(address indexed user);
+    event EnableRewardForStrategy(address indexed strategy, address indexed reward);
+    event DisableRewardForStrategy(address indexed strategy, address indexed reward, bool forfeitRecentReward);
 }

src/lib/StorageLib.sol

@@ -44,7 +44,6 @@ struct AggregationVaultStorage {
 /// allocationPoints: number of points allocated to this strategy
 /// active: a boolean to indice if this strategy is active or not
 /// cap: an optional cap in terms of deposited underlying asset. By default, it is set to 0(not activated)
-
 struct Strategy {
     uint120 allocated;
     uint120 allocationPoints;

src/module/Rewards.sol

@@ -13,10 +13,11 @@ import {StorageLib, AggregationVaultStorage} from "../lib/StorageLib.sol";
 import {ErrorsLib} from "../lib/ErrorsLib.sol";
 import {EventsLib} from "../lib/EventsLib.sol";
 
-/// @title BalanceForwarder contract
+/// @title Rewards module
 /// @custom:security-contact [email protected]
 /// @author Euler Labs (https://www.eulerlabs.com/)
-/// @notice A generic contract to integrate with https://github.com/euler-xyz/reward-streams
+/// @notice A module to provide balancer tracking for reward streats and to integrate with strategies rewards.
+/// @dev See https://github.com/euler-xyz/reward-streams.
 abstract contract RewardsModule is IBalanceForwarder, Shared {
     /// @notice Opt in to strategy rewards
     /// @param _strategy Strategy address
@@ -38,6 +39,39 @@ abstract contract RewardsModule is IBalanceForwarder, Shared {
         emit EventsLib.OptOutStrategyRewards(_strategy);
     }
 
+    /// @notice Enable aggregation layer vault rewards for specific strategy's reward token.
+    /// @param _strategy Strategy address.
+    /// @param _reward Reward token address.
+    function enableRewardForStrategy(address _strategy, address _reward) external virtual nonReentrant {
+        AggregationVaultStorage storage $ = StorageLib._getAggregationVaultStorage();
+
+        if (!$.strategies[_strategy].active) revert ErrorsLib.InactiveStrategy();
+
+        IRewardStreams(IBalanceForwarder(_strategy).balanceTrackerAddress()).enableReward(_strategy, _reward);
+
+        emit EventsLib.EnableRewardForStrategy(_strategy, _reward);
+    }
+
+    /// @notice Disable aggregation layer vault rewards for specific strategy's reward token.
+    /// @param _strategy Strategy address.
+    /// @param _reward Reward token address.
+    /// @param _forfeitRecentReward Whether to forfeit the recent rewards or not.
+    function disableRewardForStrategy(address _strategy, address _reward, bool _forfeitRecentReward)
+        external
+        virtual
+        nonReentrant
+    {
+        AggregationVaultStorage storage $ = StorageLib._getAggregationVaultStorage();
+
+        if (!$.strategies[_strategy].active) revert ErrorsLib.InactiveStrategy();
+
+        IRewardStreams(IBalanceForwarder(_strategy).balanceTrackerAddress()).disableReward(
+            _strategy, _reward, _forfeitRecentReward
+        );
+
+        emit EventsLib.DisableRewardForStrategy(_strategy, _reward, _forfeitRecentReward);
+    }
+
     /// @notice Claim a specific strategy rewards
     /// @param _strategy Strategy address.
     /// @param _reward The address of the reward token.
@@ -114,12 +148,6 @@ abstract contract RewardsModule is IBalanceForwarder, Shared {
         emit EventsLib.DisableBalanceForwarder(_sender);
     }
 
-    function _setBalanceTracker(address _balancerTracker) internal {
-        AggregationVaultStorage storage $ = StorageLib._getAggregationVaultStorage();
-
-        $.balanceTracker = _balancerTracker;
-    }
-
     /// @notice Retrieves boolean indicating if the account opted in to forward balance changes to the rewards contract
     /// @param _account Address to query
     /// @return True if balance forwarder is enabled

src/plugin/WithdrawalQueue.sol

@@ -3,13 +3,13 @@ pragma solidity ^0.8.0;
 
 // interfaces
 import {IERC4626} from "@openzeppelin/contracts/token/ERC20/extensions/ERC4626.sol";
-// internal dep
 import {IFourSixTwoSixAgg} from "../interface/IFourSixTwoSixAgg.sol";
+import {IWithdrawalQueue} from "../interface/IWithdrawalQueue.sol";
 // contracts
 import {AccessControlEnumerableUpgradeable} from
     "@openzeppelin-upgradeable/access/extensions/AccessControlEnumerableUpgradeable.sol";
 
-contract WithdrawalQueue is AccessControlEnumerableUpgradeable {
+contract WithdrawalQueue is AccessControlEnumerableUpgradeable, IWithdrawalQueue {
     error OutOfBounds();
     error SameIndexes();
     error NotEnoughAssets();