fix spearbit-audits/review-kiln#26

kilnfi · May 16, 2023 · dc6f050 · dc6f050
1 parent 8def7d6
commit dc6f050
Show file tree

Hide file tree

Showing 3 changed files with 71 additions and 11 deletions.
diff --git a/src/contracts/StakingContract.sol b/src/contracts/StakingContract.sol
@@ -45,6 +45,7 @@ contract StakingContract {
     error FundedValidatorDeletionAttempt();
     error OperatorLimitTooHigh(uint256 limit, uint256 keyCount);
     error MaximumOperatorCountAlreadyReached();
+    error LastEditAfterSnapshot();
 
     struct ValidatorAllocationCache {
         bool used;
@@ -443,7 +444,12 @@ contract StakingContract {
     /// @dev Allows all keys to be verified by the system admin before limit is increased
     /// @param _operatorIndex Operator Index
     /// @param _limit New staking limit
-    function setOperatorLimit(uint256 _operatorIndex, uint256 _limit) external onlyAdmin {
+    /// @param _snapshot Block number at which verification was done
+    function setOperatorLimit(
+        uint256 _operatorIndex,
+        uint256 _limit,
+        uint256 _snapshot
+    ) external onlyAdmin {
         StakingContractStorageLib.OperatorsSlot storage operators = StakingContractStorageLib.getOperators();
         if (operators.value[_operatorIndex].deactivated) {
             revert Deactivated();
@@ -452,6 +458,12 @@ contract StakingContract {
         if (publicKeyCount < _limit) {
             revert OperatorLimitTooHigh(_limit, publicKeyCount);
         }
+        if (
+            operators.value[_operatorIndex].limit < _limit &&
+            StakingContractStorageLib.getLastValidatorEdit() > _snapshot
+        ) {
+            revert LastEditAfterSnapshot();
+        }
         operators.value[_operatorIndex].limit = _limit;
         _updateAvailableValidatorCount(_operatorIndex);
         emit ChangedOperatorLimit(_operatorIndex, _limit);
@@ -555,6 +567,7 @@ contract StakingContract {
 
         emit ValidatorKeysAdded(_operatorIndex, _publicKeys, _signatures);
 
+        StakingContractStorageLib.setLastValidatorEdit(block.number);
         _updateAvailableValidatorCount(_operatorIndex);
     }
 
@@ -615,6 +628,7 @@ contract StakingContract {
             emit ChangedOperatorLimit(_operatorIndex, _indexes[_indexes.length - 1]);
         }
 
+        StakingContractStorageLib.setLastValidatorEdit(block.number);
         _updateAvailableValidatorCount(_operatorIndex);
     }
 

diff --git a/src/contracts/libs/StakingContractStorageLib.sol b/src/contracts/libs/StakingContractStorageLib.sol
@@ -393,4 +393,18 @@ library StakingContractStorageLib {
     function setDepositStopped(bool val) internal {
         setBool(DEPOSIT_STOPPED_SLOT, val);
     }
+
+    /* ========================================
+    ===========================================
+    =========================================*/
+
+    bytes32 internal constant LAST_VALIDATOR_EDIT_SLOT = keccak256("StakingContract.lastValidatorsEdit");
+
+    function getLastValidatorEdit() internal view returns (uint256) {
+        return getUint256(LAST_VALIDATOR_EDIT_SLOT);
+    }
+
+    function setLastValidatorEdit(uint256 value) internal {
+        setUint256(LAST_VALIDATOR_EDIT_SLOT, value);
+    }
 }
diff --git a/src/test/StakingContract.t.sol b/src/test/StakingContract.t.sol
@@ -157,7 +157,7 @@ contract StakingContractTest is DSTestPlus {
 
         {
             vm.startPrank(admin);
-            stakingContract.setOperatorLimit(0, 10);
+            stakingContract.setOperatorLimit(0, 10, block.number);
             vm.stopPrank();
         }
     }
@@ -319,22 +319,54 @@ contract StakingContractTest is DSTestPlus {
         vm.startPrank(admin);
         vm.expectEmit(true, true, true, true);
         emit ChangedOperatorLimit(operatorIndex, _limit);
-        stakingContract.setOperatorLimit(operatorIndex, _limit);
+        stakingContract.setOperatorLimit(operatorIndex, _limit, block.number);
         vm.stopPrank();
 
         (, , limit, , , , ) = stakingContract.getOperator(operatorIndex);
         assertEq(limit, _limit);
     }
 
+    function testSetOperatorLimit_snapshotRevert(uint256 _operatorSalt, uint8 _limit) public {
+        vm.assume(_limit > 0);
+        address newOperator = uf._new(_operatorSalt);
+        address newOperatorFeeRecipient = uf._new(_operatorSalt);
+
+        uint256 operatorIndex;
+
+        vm.startPrank(admin);
+        operatorIndex = stakingContract.addOperator(newOperator, newOperatorFeeRecipient);
+        vm.stopPrank();
+
+        (, , uint256 limit, , , , ) = stakingContract.getOperator(operatorIndex);
+        assertEq(limit, 0);
+
+        vm.roll(1000);
+        if (_limit > 0) {
+            vm.startPrank(newOperator);
+            stakingContract.addValidators(
+                operatorIndex,
+                _limit,
+                genBytes(48 * uint256(_limit)),
+                genBytes(96 * uint256(_limit))
+            );
+            vm.stopPrank();
+        }
+
+        vm.startPrank(admin);
+        vm.expectRevert(abi.encodeWithSignature("LastEditAfterSnapshot()"));
+        stakingContract.setOperatorLimit(operatorIndex, _limit, block.number - 10);
+        vm.stopPrank();
+    }
+
     function testSetOperatorLimitUnauthorized() public {
         vm.expectRevert(abi.encodeWithSignature("Unauthorized()"));
-        stakingContract.setOperatorLimit(0, 10);
+        stakingContract.setOperatorLimit(0, 10, block.number);
     }
 
     function testSetOperatorLimitTooHighUnauthorized() public {
         vm.startPrank(admin);
         vm.expectRevert(abi.encodeWithSignature("OperatorLimitTooHigh(uint256,uint256)", 11, 10));
-        stakingContract.setOperatorLimit(0, 11);
+        stakingContract.setOperatorLimit(0, 11, block.number);
         vm.stopPrank();
     }
 
@@ -358,7 +390,7 @@ contract StakingContractTest is DSTestPlus {
         vm.startPrank(admin);
         stakingContract.deactivateOperator(operatorIndex, operatorOne);
         vm.expectRevert(abi.encodeWithSignature("Deactivated()"));
-        stakingContract.setOperatorLimit(operatorIndex, _limit);
+        stakingContract.setOperatorLimit(operatorIndex, _limit, block.number);
         vm.stopPrank();
     }
 
@@ -408,7 +440,7 @@ contract StakingContractTest is DSTestPlus {
         assert(deactivated == false);
 
         vm.startPrank(admin);
-        stakingContract.setOperatorLimit(0, 20);
+        stakingContract.setOperatorLimit(0, 20, block.number);
         vm.stopPrank();
 
         (operatorAddress, feeRecipientAddress, limit, keys, funded, available, deactivated) = stakingContract
@@ -1096,7 +1128,7 @@ contract StakingContractOperatorTest is DSTestPlus {
         stakingContract.addValidators(0, 10, publicKeys, signatures);
         vm.stopPrank();
         vm.startPrank(admin);
-        stakingContract.setOperatorLimit(0, 10);
+        stakingContract.setOperatorLimit(0, 10, block.number);
         vm.stopPrank();
 
         vm.deal(address(this), 32 ether);
@@ -1185,7 +1217,7 @@ contract StakingContractDistributionTest is DSTestPlus {
             stakingContract.addValidators(i, keyPerOperator, publicKeys, signatures);
             vm.stopPrank();
             vm.startPrank(admin);
-            stakingContract.setOperatorLimit(i, keyPerOperator);
+            stakingContract.setOperatorLimit(i, keyPerOperator, block.number);
             vm.stopPrank();
         }
 
@@ -1293,7 +1325,7 @@ contract StakingContractOneValidatorTest is Test {
 
         {
             vm.startPrank(admin);
-            stakingContract.setOperatorLimit(0, 10);
+            stakingContract.setOperatorLimit(0, 10, block.number);
             vm.stopPrank();
         }
     }
@@ -2077,7 +2109,7 @@ contract StakingContractBehindProxyTest is Test {
 
         {
             vm.startPrank(admin);
-            stakingContract.setOperatorLimit(0, 10);
+            stakingContract.setOperatorLimit(0, 10, block.number);
             vm.stopPrank();
         }
     }