fix: correct distribution end epoch calculation and improve button pressability check

0xClandestine · 0xClandestine · commit f61520110174 · 2026-01-19T21:29:07.000-05:00
Fix an off-by-one error where distributions were ending one epoch too late.

A distribution with totalEpochs=N should run for exactly N epochs, not N+1.
diff --git a/src/contracts/core/EmissionsController.sol b/src/contracts/core/EmissionsController.sol
@@ -82,12 +82,12 @@ contract EmissionsController is
         // Prevents minting EIGEN before the first epoch has started.
         require(currentEpoch != type(uint256).max, EmissionsNotStarted());
 
+        // Check if there are any active distributions to process.
+        require(isButtonPressable(), AllDistributionsProcessed());
+
         uint256 totalDistributions = getTotalDistributions();
         uint256 nextDistributionId = _epochs[currentEpoch].totalProcessed;
 
-        // Check if all distributions have already been processed.
-        require(nextDistributionId < totalDistributions, AllDistributionsProcessed());
-
         // Mint the total amount of bEIGEN/EIGEN needed for all distributions.
         if (!_epochs[currentEpoch].minted) {
             // NOTE: Approvals may not be entirely spent.
@@ -118,14 +118,8 @@ contract EmissionsController is
         for (uint256 distributionId = nextDistributionId; distributionId < lastIndex; ++distributionId) {
             Distribution memory distribution = _distributions[distributionId];
 
-            // Skip disabled distributions...
-            if (distribution.distributionType == DistributionType.Disabled) continue;
-            // Skip distributions that haven't started yet...
-            if (distribution.startEpoch > currentEpoch) continue;
-            // Skip distributions that have ended (0 means infinite)...
-            if (distribution.totalEpochs != 0) {
-                if (currentEpoch > distribution.startEpoch + distribution.totalEpochs) continue;
-            }
+            // Skip inactive distributions
+            if (!_isDistributionActive(distribution, currentEpoch)) continue;
 
             _processDistribution({
                 distribution: distribution,
@@ -370,6 +364,26 @@ contract EmissionsController is
         );
     }
 
+    /// @dev Internal helper to check if a distribution is active for the given epoch.
+    /// A distribution is active if it is not disabled, has started, and has not ended.
+    /// @param distribution The distribution to check.
+    /// @param currentEpoch The epoch to check against.
+    /// @return True if the distribution is active for the given epoch, false otherwise.
+    function _isDistributionActive(
+        Distribution memory distribution,
+        uint256 currentEpoch
+    ) internal pure returns (bool) {
+        // Skip disabled distributions
+        if (distribution.distributionType == DistributionType.Disabled) return false;
+        // Skip distributions that haven't started yet
+        if (distribution.startEpoch > currentEpoch) return false;
+        // Skip distributions that have ended (0 means infinite)
+        if (distribution.totalEpochs != 0) {
+            if (currentEpoch >= distribution.startEpoch + distribution.totalEpochs) return false;
+        }
+        return true;
+    }
+
     /// -----------------------------------------------------------------------
     /// View
     /// -----------------------------------------------------------------------
@@ -384,7 +398,18 @@ contract EmissionsController is
 
     /// @inheritdoc IEmissionsController
     function isButtonPressable() public view returns (bool) {
-        return _epochs[getCurrentEpoch()].totalProcessed < getTotalDistributions();
+        uint256 currentEpoch = getCurrentEpoch();
+        uint256 totalDistributions = getTotalDistributions();
+        uint256 nextDistributionId = _epochs[currentEpoch].totalProcessed;
+
+        // Check if any remaining distribution is active for the current epoch
+        for (uint256 i = nextDistributionId; i < totalDistributions; ++i) {
+            if (_isDistributionActive(_distributions[i], currentEpoch)) {
+                return true;
+            }
+        }
+
+        return false;
     }
 
     /// @inheritdoc IEmissionsController
diff --git a/src/test/unit/EmissionsControllerUnit.t.sol b/src/test/unit/EmissionsControllerUnit.t.sol
@@ -137,6 +137,170 @@ contract EmissionsControllerUnitTests_pressButton is EmissionsControllerUnitTest
         cheats.expectRevert(IEmissionsControllerErrors.EmissionsNotStarted.selector);
         emissionsController.pressButton(1);
     }
+
+    function test_revert_pressButton_AfterOneEpoch_AllDistributionsProcessed() public {
+        // Add a distribution that should run for exactly 1 epoch (epoch 0 only)
+        cheats.prank(incentiveCouncil);
+        emissionsController.addDistribution(
+            Distribution({
+                weight: 10_000,
+                startEpoch: 0,
+                totalEpochs: 1,
+                distributionType: DistributionType.RewardsForAllEarners,
+                operatorSet: emptyOperatorSet(),
+                strategiesAndMultipliers: defaultStrategiesAndMultipliers()
+            })
+        );
+
+        // Epoch 0: should process
+        cheats.warp(EMISSIONS_START_TIME);
+        assertEq(emissionsController.getCurrentEpoch(), 0);
+        assertTrue(emissionsController.isButtonPressable());
+        emissionsController.pressButton(1);
+
+        // Epoch 1: should NOT process (distribution ended at epoch 0+1=1)
+        cheats.warp(EMISSIONS_START_TIME + EMISSIONS_EPOCH_LENGTH);
+        assertEq(emissionsController.getCurrentEpoch(), 1);
+        assertFalse(emissionsController.isButtonPressable());
+        cheats.expectRevert(IEmissionsControllerErrors.AllDistributionsProcessed.selector);
+        emissionsController.pressButton(1);
+    }
+
+    function test_revert_pressButton_AfterTwoEpochs_AllDistributionsProcessed() public {
+        // Add a distribution that should run for exactly 2 epochs (epochs 0 and 1)
+        cheats.prank(incentiveCouncil);
+        emissionsController.addDistribution(
+            Distribution({
+                weight: 10_000,
+                startEpoch: 0,
+                totalEpochs: 2,
+                distributionType: DistributionType.RewardsForAllEarners,
+                operatorSet: emptyOperatorSet(),
+                strategiesAndMultipliers: defaultStrategiesAndMultipliers()
+            })
+        );
+
+        // Epoch 0: should process
+        cheats.warp(EMISSIONS_START_TIME);
+        assertEq(emissionsController.getCurrentEpoch(), 0);
+        emissionsController.pressButton(1);
+        // Verify button is not pressable again in same epoch
+        assertFalse(emissionsController.isButtonPressable());
+
+        // Epoch 1: should process
+        cheats.warp(EMISSIONS_START_TIME + EMISSIONS_EPOCH_LENGTH);
+        assertEq(emissionsController.getCurrentEpoch(), 1);
+        assertTrue(emissionsController.isButtonPressable());
+        emissionsController.pressButton(1);
+        assertFalse(emissionsController.isButtonPressable());
+
+        // Epoch 2: should NOT process (distribution ended at epoch 0+2=2)
+        cheats.warp(EMISSIONS_START_TIME + 2 * EMISSIONS_EPOCH_LENGTH);
+        assertEq(emissionsController.getCurrentEpoch(), 2);
+        // Button should not be pressable because all distributions have ended
+        assertFalse(emissionsController.isButtonPressable());
+        cheats.expectRevert(IEmissionsControllerErrors.AllDistributionsProcessed.selector);
+        emissionsController.pressButton(1);
+    }
+
+    function test_revert_pressButton_AfterThreeFutureEpochs_AllDistributionsProcessed() public {
+        // Add a distribution that starts at epoch 2 and runs for 3 epochs (epochs 2, 3, 4)
+        cheats.prank(incentiveCouncil);
+        emissionsController.addDistribution(
+            Distribution({
+                weight: 10_000,
+                startEpoch: 2,
+                totalEpochs: 3,
+                distributionType: DistributionType.RewardsForAllEarners,
+                operatorSet: emptyOperatorSet(),
+                strategiesAndMultipliers: defaultStrategiesAndMultipliers()
+            })
+        );
+
+        // Epoch 0: should not process (hasn't started)
+        cheats.warp(EMISSIONS_START_TIME);
+        assertEq(emissionsController.getCurrentEpoch(), 0);
+        assertFalse(emissionsController.isButtonPressable());
+
+        // Epoch 1: should not process (hasn't started)
+        cheats.warp(EMISSIONS_START_TIME + EMISSIONS_EPOCH_LENGTH);
+        assertEq(emissionsController.getCurrentEpoch(), 1);
+        assertFalse(emissionsController.isButtonPressable());
+
+        // Epoch 2: should process (first epoch)
+        cheats.warp(EMISSIONS_START_TIME + 2 * EMISSIONS_EPOCH_LENGTH);
+        assertEq(emissionsController.getCurrentEpoch(), 2);
+        assertTrue(emissionsController.isButtonPressable());
+        emissionsController.pressButton(1);
+
+        // Epoch 3: should process
+        cheats.warp(EMISSIONS_START_TIME + 3 * EMISSIONS_EPOCH_LENGTH);
+        assertEq(emissionsController.getCurrentEpoch(), 3);
+        assertTrue(emissionsController.isButtonPressable());
+        emissionsController.pressButton(1);
+
+        // Epoch 4: should process (last epoch)
+        cheats.warp(EMISSIONS_START_TIME + 4 * EMISSIONS_EPOCH_LENGTH);
+        assertEq(emissionsController.getCurrentEpoch(), 4);
+        assertTrue(emissionsController.isButtonPressable());
+        emissionsController.pressButton(1);
+
+        // Epoch 5: should NOT process (ended at epoch 2+3=5)
+        cheats.warp(EMISSIONS_START_TIME + 5 * EMISSIONS_EPOCH_LENGTH);
+        assertEq(emissionsController.getCurrentEpoch(), 5);
+        assertFalse(emissionsController.isButtonPressable());
+        cheats.expectRevert(IEmissionsControllerErrors.AllDistributionsProcessed.selector);
+        emissionsController.pressButton(1);
+    }
+
+    function testFuzz_revert_pressButton_AfterDistributionEnds_AllDistributionsProcessed(uint8 _startEpoch, uint8 _totalEpochs) public {
+        // Bound inputs to reasonable values
+        uint64 startEpoch = uint64(bound(_startEpoch, 0, 10));
+        uint64 totalEpochs = uint64(bound(_totalEpochs, 1, 10));
+
+        // Add a distribution with fuzzed parameters
+        cheats.prank(incentiveCouncil);
+        emissionsController.addDistribution(
+            Distribution({
+                weight: 10_000,
+                startEpoch: startEpoch,
+                totalEpochs: totalEpochs,
+                distributionType: DistributionType.RewardsForAllEarners,
+                operatorSet: emptyOperatorSet(),
+                strategiesAndMultipliers: defaultStrategiesAndMultipliers()
+            })
+        );
+
+        uint endEpoch = startEpoch + totalEpochs;
+        uint maxTestEpoch = endEpoch + 2; // Test a couple epochs after the end
+
+        // Test each epoch from 0 to maxTestEpoch
+        for (uint epoch = 0; epoch <= maxTestEpoch; epoch++) {
+            // Warp to the beginning of this epoch
+            cheats.warp(EMISSIONS_START_TIME + epoch * EMISSIONS_EPOCH_LENGTH);
+            assertEq(emissionsController.getCurrentEpoch(), epoch);
+
+            // Determine if the button should be pressable in this epoch
+            bool shouldBeActive = (epoch >= startEpoch && epoch < endEpoch);
+
+            if (shouldBeActive) {
+                // Button should be pressable if we haven't pressed it in this epoch yet
+                assertTrue(emissionsController.isButtonPressable(), "Button should be pressable during active epochs");
+                emissionsController.pressButton(1);
+                // After pressing, button should not be pressable again in the same epoch
+                assertFalse(emissionsController.isButtonPressable(), "Button should not be pressable after pressing in same epoch");
+            } else {
+                // Button should not be pressable before start or after end
+                assertFalse(emissionsController.isButtonPressable(), "Button should not be pressable outside active epochs");
+
+                // If we're at or past the end epoch, verify the correct revert
+                if (epoch >= endEpoch) {
+                    cheats.expectRevert(IEmissionsControllerErrors.AllDistributionsProcessed.selector);
+                    emissionsController.pressButton(1);
+                }
+            }
+        }
+    }
 }
 
 contract EmissionsControllerUnitTests_sweep is EmissionsControllerUnitTests {
