Merge pull request #1488 from lidofinance/feat/fix-staking-limit

Author: tamtamchik

contracts/0.4.24/Lido.sol

@@ -1141,10 +1141,9 @@ contract Lido is Versioned, StETHPermit, AragonApp {
         StakeLimitState.Data memory stakeLimitData = STAKING_STATE_POSITION.getStorageStakeLimitStruct();
         if (stakeLimitData.isStakingLimitSet()) {
             uint256 newStakeLimit = stakeLimitData.calculateCurrentStakeLimit() + _amount;
-            uint256 maxStakeLimit = stakeLimitData.maxStakeLimit;
 
             STAKING_STATE_POSITION.setStorageStakeLimitStruct(
-                stakeLimitData.updatePrevStakeLimit(newStakeLimit > maxStakeLimit ? maxStakeLimit : newStakeLimit)
+                stakeLimitData.updatePrevStakeLimit(newStakeLimit)
             );
         }
     }

contracts/0.4.24/lib/StakeLimitUtils.sol

@@ -93,7 +93,8 @@ library StakeLimitUnstructuredStorage {
 library StakeLimitUtils {
     /**
     * @notice Calculate stake limit for the current block.
-    * @dev using `_constGasMin` to make gas consumption independent of the current block number
+    * @dev using `_constGasMin`, `_constGasMax`, `_saturatingSub`, `_constGasLt` to make gas consumption independent
+    *      of the current block number
     */
     function calculateCurrentStakeLimit(StakeLimitState.Data memory _data) internal view returns(uint256 limit) {
         uint256 stakeLimitIncPerBlock;
@@ -102,12 +103,11 @@ library StakeLimitUtils {
         }
 
         uint256 blocksPassed = block.number - _data.prevStakeBlockNumber;
-        uint256 projectedLimit = _data.prevStakeLimit + blocksPassed * stakeLimitIncPerBlock;
+        uint256 change = blocksPassed * stakeLimitIncPerBlock;
 
-        limit = _constGasMin(
-            projectedLimit,
-            _data.maxStakeLimit
-        );
+        limit = _data.prevStakeLimit < _data.maxStakeLimit ?
+            _constGasMin(_data.prevStakeLimit + change, _data.maxStakeLimit) :
+            _constGasMax(_saturatingSub(_data.prevStakeLimit, change), _data.maxStakeLimit);
     }
 
     /**
@@ -215,17 +215,47 @@ library StakeLimitUtils {
         return _data;
     }
 
+    /**
+     * @notice branchless less-than comparison
+     * @param a first value
+     * @param b second value
+     * @return result 1 if a < b, 0 otherwise
+     */
+    function _constGasLt(uint256 a, uint256 b) internal pure returns (uint256 result) {
+        assembly {
+            result := lt(a, b)
+        }
+    }
+
     /**
      * @notice find a minimum of two numbers with a constant gas consumption
      * @dev doesn't use branching logic inside
      * @param _lhs left hand side value
      * @param _rhs right hand side value
      */
     function _constGasMin(uint256 _lhs, uint256 _rhs) internal pure returns (uint256 min) {
-        uint256 lhsIsLess;
-        assembly {
-            lhsIsLess := lt(_lhs, _rhs) // lhsIsLess = (_lhs < _rhs) ? 1 : 0
-        }
+        uint256 lhsIsLess = _constGasLt(_lhs, _rhs);
         min = (_lhs * lhsIsLess) + (_rhs * (1 - lhsIsLess));
     }
+
+    /**
+     * @notice find a maximum of two numbers with a constant gas consumption
+     * @dev doesn't use branching logic inside
+     * @param _lhs left hand side value
+     * @param _rhs right hand side value
+     */
+    function _constGasMax(uint256 _lhs, uint256 _rhs) internal pure returns (uint256 max) {
+        uint256 lhsIsLess = _constGasLt(_lhs, _rhs);
+        max = (_lhs * (1 - lhsIsLess)) + (_rhs * lhsIsLess);
+    }
+
+    /**
+     * @notice unsigned saturating subtraction, bounds to zero instead of overflowing
+     * @param a first value
+     * @param b second value
+     */
+    function _saturatingSub(uint256 a, uint256 b) internal pure returns (uint256 result) {
+        uint256 isUnderflow = _constGasLt(a, b);
+        result = (a - b) * (1 - isUnderflow);
+    }
 }

lib/time.ts

@@ -29,6 +29,10 @@ export async function getNextBlockTimestamp() {
   return nextBlockTimestamp;
 }
 
+export async function getCurrentBlockNumber() {
+  return await ethers.provider.getBlockNumber();
+}
+
 export async function getNextBlockNumber() {
   const latestBlock = BigInt(await time.latestBlock());
   return latestBlock + 1n;

test/0.4.24/lib/stakeLimitUtils.test.ts

@@ -211,17 +211,43 @@ describe("StakeLimitUtils.sol", () => {
         expect(await stakeLimitUtils.calculateCurrentStakeLimit()).to.equal(staticStakeLimit);
       });
 
-      it("the full limit gets restored after growth blocks", async () => {
+      it("the full limit gets restored after growth blocks (increasing to limit)", async () => {
         prevStakeBlockNumber = BigInt(await latestBlock());
         const baseStakeLimit = 0n;
         await stakeLimitUtils.harness_setState(prevStakeBlockNumber, 0n, maxStakeLimitGrowthBlocks, maxStakeLimit);
+
+        const growthPerBlock = maxStakeLimit / maxStakeLimitGrowthBlocks;
+
+        // 1 block passed due to the setter call above
+        expect(await stakeLimitUtils.calculateCurrentStakeLimit()).to.equal(growthPerBlock);
+
+        // growth blocks passed (might be not equal to maxStakeLimit yet due to rounding)
+        await mineUpTo(BigInt(prevStakeBlockNumber) + maxStakeLimitGrowthBlocks);
+        expect(await stakeLimitUtils.calculateCurrentStakeLimit()).to.equal(
+          baseStakeLimit + maxStakeLimitGrowthBlocks * growthPerBlock,
+        );
+
+        // move forward one more block to account for rounding and reach max
+        await mineUpTo(BigInt(prevStakeBlockNumber) + maxStakeLimitGrowthBlocks + 1n);
+        // growth blocks mined, the limit should be full
+        expect(await stakeLimitUtils.calculateCurrentStakeLimit()).to.equal(maxStakeLimit);
+      });
+
+      it("the full limit gets restored after growth blocks (decreasing to limit)", async () => {
+        prevStakeBlockNumber = BigInt(await latestBlock());
+        const initial = maxStakeLimit * 2n;
+
+        await stakeLimitUtils.harness_setState(prevStakeBlockNumber, initial, maxStakeLimitGrowthBlocks, maxStakeLimit);
+
+        const growthPerBlock = maxStakeLimit / maxStakeLimitGrowthBlocks;
+
         // 1 block passed due to the setter call above
-        expect(await stakeLimitUtils.calculateCurrentStakeLimit()).to.equal(maxStakeLimit / maxStakeLimitGrowthBlocks);
+        expect(await stakeLimitUtils.calculateCurrentStakeLimit()).to.equal(initial - growthPerBlock);
 
         // growth blocks passed (might be not equal to maxStakeLimit yet due to rounding)
         await mineUpTo(BigInt(prevStakeBlockNumber) + maxStakeLimitGrowthBlocks);
         expect(await stakeLimitUtils.calculateCurrentStakeLimit()).to.equal(
-          baseStakeLimit + maxStakeLimitGrowthBlocks * (maxStakeLimit / maxStakeLimitGrowthBlocks),
+          initial - maxStakeLimitGrowthBlocks * growthPerBlock,
         );
 
         // move forward one more block to account for rounding and reach max

test/0.4.24/lido/lido.externalShares.test.ts

@@ -6,7 +6,7 @@ import { HardhatEthersSigner } from "@nomicfoundation/hardhat-ethers/signers";
 
 import { ACL, Lido, LidoLocator } from "typechain-types";
 
-import { ether, impersonate, MAX_UINT256 } from "lib";
+import { advanceChainTime, ether, impersonate, MAX_UINT256 } from "lib";
 import { TOTAL_BASIS_POINTS } from "lib/constants";
 
 import { deployLidoDao } from "test/deploy";
@@ -386,17 +386,20 @@ describe("Lido.sol:externalShares", () => {
 
     it("Increases staking limit when burning", async () => {
       await lido.setMaxExternalRatioBP(maxExternalRatioBP);
-      await lido.setStakingLimit(10n, 1n);
+      await lido.setStakingLimit(10n, 10n);
 
       await lido.connect(vaultHubSigner).mintExternalShares(vaultHubSigner, 1n);
 
-      expect(await lido.getCurrentStakeLimit()).to.equal(9n);
+      let limit = 9n;
+      expect(await lido.getCurrentStakeLimit()).to.equal(limit);
 
       await lido.connect(vaultHubSigner).burnExternalShares(1n);
-      expect(await lido.getCurrentStakeLimit()).to.equal(10n);
+      limit += 1n; // for mining block with burning
+
+      expect(await lido.getCurrentStakeLimit()).to.equal(limit + 1n);
     });
 
-    it("Restores staking limit to max when burning more", async () => {
+    it("Bypasses staking limit when burning more than staking limit", async () => {
       await lido.setMaxExternalRatioBP(maxExternalRatioBP);
       await lido.connect(vaultHubSigner).mintExternalShares(vaultHubSigner, 5n);
 
@@ -407,10 +410,15 @@ describe("Lido.sol:externalShares", () => {
       const amountToMint = await lido.getPooledEthByShares(sharesToMint);
       await lido.connect(vaultHubSigner).mintExternalShares(vaultHubSigner, sharesToMint);
 
-      expect(await lido.getCurrentStakeLimit()).to.equal(10n - amountToMint);
+      let limit = 10n - amountToMint;
+      expect(await lido.getCurrentStakeLimit()).to.equal(limit);
 
-      await lido.connect(vaultHubSigner).burnExternalShares(10n);
-      expect(await lido.getCurrentStakeLimit()).to.equal(10n);
+      const sharesToBurn = 10n;
+      const amountToBurn = await lido.getPooledEthByShares(sharesToBurn);
+      await lido.connect(vaultHubSigner).burnExternalShares(sharesToBurn);
+      limit += 1n; // for mining block with burning
+
+      expect(await lido.getCurrentStakeLimit()).to.equal(limit + amountToBurn);
     });
   });
 
@@ -487,18 +495,16 @@ describe("Lido.sol:externalShares", () => {
 
     it("Can mint and burn external shares without limit change after multiple loops", async () => {
       await lido.setMaxExternalRatioBP(maxExternalRatioBP);
-      await lido.setStakingLimit(1000n, 1n);
+      await lido.setStakingLimit(1000n, 100n);
 
       for (let i = 1n; i <= 500n; i++) {
-        const stakingLimitBefore = await lido.getCurrentStakeLimit();
-        expect(stakingLimitBefore).to.equal(1000n);
-
         await lido.connect(vaultHubSigner).mintExternalShares(vaultHubSigner, i);
         await lido.connect(vaultHubSigner).burnExternalShares(i);
-
-        const stakingLimitAfter = await lido.getCurrentStakeLimit();
-        expect(stakingLimitAfter).to.equal(stakingLimitBefore);
       }
+
+      // need to mine a block to update the stake limit otherwise it will be 1000n + 100n (after burning)
+      await advanceChainTime(1n);
+      expect(await lido.getCurrentStakeLimit()).to.equal(1000n);
     });
   });