feat: add AllocationExtension.sol

.github/workflows/forge-test.yml

@@ -27,6 +27,11 @@ jobs:
         uses: foundry-rs/[email protected]
         with:
           version: nightly
+
+      - name: Run Smock
+        run: |
+          bun smock
+        id: smock
 
       - name: Run Forge Format
         run: |

contracts/extensions/contracts/AllocationExtension.sol

@@ -0,0 +1,141 @@
+// SPDX-License-Identifier: AGPL-3.0-only
+pragma solidity 0.8.19;
+
+import {IAllocationExtension} from "contracts/extensions/interfaces/IAllocationExtension.sol";
+import {CoreBaseStrategy} from "contracts/strategies/CoreBaseStrategy.sol";
+
+abstract contract AllocationExtension is CoreBaseStrategy, IAllocationExtension {
+    /// ================================
+    /// ========== Storage =============
+    /// ================================
+
+    /// @notice The start and end times for allocations
+    uint64 public allocationStartTime;
+    uint64 public allocationEndTime;
+
+    /// @notice Defines if the strategy is sending Metadata struct in the data parameter
+    bool public isUsingAllocationMetadata;
+
+    /// @notice token -> isAllowed
+    mapping(address => bool) public allowedTokens;
+
+    /// ===============================
+    /// ========= Initialize ==========
+    /// ===============================
+
+    /// @notice This initializes the Alocation Extension
+    /// @dev This function MUST be called by the 'initialize' function in the strategy.
+    /// @param _allowedTokens The allowed tokens
+    /// @param _allocationStartTime The start time for the allocation period
+    /// @param _allocationEndTime The end time for the allocation period
+    /// @param _isUsingAllocationMetadata Defines if the strategy is sending Metadata struct in the data parameter
+    function __AllocationExtension_init(
+        address[] memory _allowedTokens,
+        uint64 _allocationStartTime,
+        uint64 _allocationEndTime,
+        bool _isUsingAllocationMetadata
+    ) internal virtual {
+        if (_allowedTokens.length == 0) {
+            // all tokens
+            allowedTokens[address(0)] = true;
+        } else {
+            for (uint256 i; i < _allowedTokens.length; i++) {
+                allowedTokens[_allowedTokens[i]] = true;
+            }
+        }
+
+        isUsingAllocationMetadata = _isUsingAllocationMetadata;
+
+        _updateAllocationTimestamps(_allocationStartTime, _allocationEndTime);
+    }
+
+    /// ====================================
+    /// =========== Modifiers ==============
+    /// ====================================
+
+    /// @notice Modifier to check if allocation has ended
+    /// @dev Reverts if allocation has not ended
+    modifier onlyAfterAllocation() {
+        _checkOnlyAfterAllocation();
+        _;
+    }
+
+    /// @notice Modifier to check if allocation is active
+    /// @dev Reverts if allocation is not active
+    modifier onlyActiveAllocation() {
+        _checkOnlyActiveAllocation();
+        _;
+    }
+
+    /// @notice Modifier to check if allocation has started
+    /// @dev Reverts if allocation has started
+    modifier onlyBeforeAllocation() {
+        _checkBeforeAllocation();
+        _;
+    }
+
+    /// ====================================
+    /// ============ Internal ==============
+    /// ====================================
+
+    /// @notice Checks if the allocator is valid
+    /// @param _allocator The allocator address
+    /// @return 'true' if the allocator is valid, otherwise 'false'
+    function _isValidAllocator(address _allocator) internal view virtual returns (bool);
+
+    /// @notice Returns TRUE if the token is allowed
+    /// @param _token The token to check
+    function _isAllowedToken(address _token) internal view virtual returns (bool) {
+        // all tokens allowed
+        if (allowedTokens[address(0)]) return true;
+
+        if (allowedTokens[_token]) return true;
+
+        return false;
+    }
+
+    /// @notice Sets the start and end dates for allocation.
+    /// @dev The 'msg.sender' must be a pool manager.
+    /// @param _allocationStartTime The start time for the allocation
+    /// @param _allocationEndTime The end time for the allocation
+    function _updateAllocationTimestamps(uint64 _allocationStartTime, uint64 _allocationEndTime) internal virtual {
+        if (_allocationStartTime > _allocationEndTime) revert INVALID_ALLOCATION_TIMESTAMPS();
+
+        allocationStartTime = _allocationStartTime;
+        allocationEndTime = _allocationEndTime;
+
+        emit AllocationTimestampsUpdated(allocationStartTime, allocationEndTime, msg.sender);
+    }
+
+    /// @dev Ensure the function is called before allocation start time
+    function _checkBeforeAllocation() internal virtual {
+        if (block.timestamp >= allocationStartTime) revert ALLOCATION_HAS_STARTED();
+    }
+
+    /// @dev Ensure the function is called during allocation times
+    function _checkOnlyActiveAllocation() internal virtual {
+        if (block.timestamp < allocationStartTime) revert ALLOCATION_NOT_ACTIVE();
+        if (block.timestamp > allocationEndTime) revert ALLOCATION_NOT_ACTIVE();
+    }
+
+    /// @dev Ensure the function is called after allocation start time
+    function _checkOnlyAfterAllocation() internal virtual {
+        if (block.timestamp <= allocationEndTime) revert ALLOCATION_NOT_ENDED();
+    }
+
+    //  ====================================
+    //  ==== External/Public Functions =====
+    //  ====================================
+
+    /// @notice Sets the start and end dates for allocation.
+    /// @dev The 'msg.sender' must be a pool manager.
+    /// @param _allocationStartTime The start time for the allocation
+    /// @param _allocationEndTime The end time for the allocation
+    function updateAllocationTimestamps(uint64 _allocationStartTime, uint64 _allocationEndTime)
+        external
+        virtual
+        onlyPoolManager(msg.sender)
+    {
+        _updateAllocationTimestamps(_allocationStartTime, _allocationEndTime);
+    }
+}

contracts/extensions/interfaces/IAllocationExtension.sol

@@ -0,0 +1,39 @@
+// SPDX-License-Identifier: AGPL-3.0-only
+pragma solidity 0.8.19;
+
+interface IAllocationExtension {
+    /// @dev Error thrown when the allocation timestamps are invalid
+    error INVALID_ALLOCATION_TIMESTAMPS();
+
+    /// @dev Error thrown when trying to call the function when the allocation has started
+    error ALLOCATION_HAS_STARTED();
+
+    /// @dev Error thrown when trying to call the function when the allocation is not active
+    error ALLOCATION_NOT_ACTIVE();
+
+    /// @dev Error thrown when trying to call the function when the allocation has ended
+    error ALLOCATION_NOT_ENDED();
+
+    /// @notice Emitted when the allocation timestamps are updated
+    /// @param allocationStartTime The start time for the allocation period
+    /// @param allocationEndTime The end time for the allocation period
+    /// @param sender The sender of the transaction
+    event AllocationTimestampsUpdated(uint64 allocationStartTime, uint64 allocationEndTime, address sender);
+
+    /// @notice The start time for the allocation period
+    function allocationStartTime() external view returns (uint64);
+
+    /// @notice The end time for the allocation period
+    function allocationEndTime() external view returns (uint64);
+
+    /// @notice Defines if the strategy is sending Metadata struct in the data parameter
+    function isUsingAllocationMetadata() external view returns (bool);
+
+    /// @notice Returns TRUE if the token is allowed, FALSE otherwise
+    function allowedTokens(address _token) external view returns (bool);
+
+    /// @notice Sets the start and end dates for allocation.
+    /// @param _allocationStartTime The start time for the allocation
+    /// @param _allocationEndTime The end time for the allocation
+    function updateAllocationTimestamps(uint64 _allocationStartTime, uint64 _allocationEndTime) external;
+}

contracts/strategies/CoreBaseStrategy.sol

@@ -178,7 +178,7 @@ abstract contract CoreBaseStrategy is IBaseStrategy, Transfer {
     /// @notice Checks if the '_sender' is a pool manager.
     /// @dev Reverts if the '_sender' is not a pool manager.
     /// @param _sender The address to check if they are a pool manager
-    function _checkOnlyPoolManager(address _sender) internal view {
+    function _checkOnlyPoolManager(address _sender) internal view virtual {
         if (!allo.isPoolManager(poolId, _sender)) revert BaseStrategy_UNAUTHORIZED();
     }
 

package.json

@@ -30,7 +30,7 @@
     "/////// deploy-test ///////": "echo 'deploy test scripts'",
     "create-profile": "npx hardhat run scripts/test/createProfile.ts --network",
     "create-pool": "npx hardhat run scripts/test/createPool.ts --network",
-    "smock": "smock-foundry --contracts contracts/core"
+    "smock": "smock-foundry --contracts contracts/core --contracts test/utils/mocks/"
   },
   "devDependencies": {
     "@matterlabs/hardhat-zksync-deploy": "^1.2.1",