[Certora][M-07] Vaibhav/certora/fix steth swap

src/LiquidityPool.sol

@@ -195,10 +195,13 @@ contract LiquidityPool is Initializable, OwnableUpgradeable, UUPSUpgradeable, IL
     function withdraw(address _recipient, uint256 _amount) external whenNotPaused returns (uint256) {
         uint256 share = sharesForWithdrawalAmount(_amount);
         require(msg.sender == address(withdrawRequestNFT) || msg.sender == address(membershipManager) || msg.sender == address(liquifier), "Incorrect Caller");
-        if (totalValueInLp < _amount || (msg.sender == address(withdrawRequestNFT) && ethAmountLockedForWithdrawal < _amount) || eETH.balanceOf(msg.sender) < _amount) revert InsufficientLiquidity();
+        if (totalValueInLp < _amount || (msg.sender == address(withdrawRequestNFT) && ethAmountLockedForWithdrawal < _amount) || (msg.sender == address(liquifier) && eETH.balanceOf(msg.sender) < _amount)) revert InsufficientLiquidity();
         if (_amount > type(uint128).max || _amount == 0 || share == 0) revert InvalidAmount();
-
+        if (msg.sender == address(liquifier)) {
+            totalValueOutOfLp -= uint128(_amount);
+        } else {
         totalValueInLp -= uint128(_amount);
+        }
         if (msg.sender == address(withdrawRequestNFT)) {
             ethAmountLockedForWithdrawal -= uint128(_amount);
         }

src/Liquifier.sol

@@ -94,6 +94,7 @@ contract Liquifier is Initializable, UUPSUpgradeable, OwnableUpgradeable, Pausab
     BucketRateLimiter public rateLimiter;
 
     bytes32 public constant LIQUIFIER_ADMIN_ROLE = keccak256("LIQUIFIER_ADMIN_ROLE");
+    bytes32 public constant EETH_STETH_SWAPPER = keccak256("EETH_STETH_SWAPPER");
 
     event Liquified(address _user, uint256 _toEEthAmount, address _fromToken, bool _isRestaked);
     event QueuedStEthWithdrawals(uint256[] _reqIds);
@@ -180,9 +181,13 @@ contract Liquifier is Initializable, UUPSUpgradeable, OwnableUpgradeable, Pausab
     function swapEEthForStEth(uint256 _amount) external whenNotPaused nonReentrant {
         if (_amount > lido.balanceOf(address(this))) revert NotEnoughBalance();
         if (_amount > liquidityPool.eETH().balanceOf(msg.sender)) revert NotEnoughBalance();
-
+        if (!roleRegistry.hasRole(EETH_STETH_SWAPPER, msg.sender)) revert IncorrectRole();
+        uint256 preBalance = liquidityPool.eETH().balanceOf(address(this));
         IERC20(address(liquidityPool.eETH())).safeTransferFrom(msg.sender, address(this), _amount);
+        uint256 postBalance = liquidityPool.eETH().balanceOf(address(this));
+        uint256 transferredAmount = postBalance - preBalance;
         IERC20(address(lido)).safeTransfer(msg.sender, _amount);
+        liquidityPool.withdraw(address(liquidityPool), transferredAmount);
     }
 
     function depositWithERC20WithPermit(address _token, uint256 _amount, address _referral, PermitInput calldata _permit) external whenNotPaused returns (uint256) {
@@ -383,7 +388,6 @@ contract Liquifier is Initializable, UUPSUpgradeable, OwnableUpgradeable, Pausab
         for (uint256 i = 0; i < dummies.length; i++) {
             total += getTotalPooledEther(address(dummies[i]));
         }
-        total += liquidityPool.eETH().balanceOf(address(this));
     }
 
     /// deposited ETH can have 2 states:

test/Liquifier.t.sol

@@ -438,43 +438,42 @@ contract LiquifierTest is TestSetup {
         liquifierInstance.getTotalPooledEther();
     }
 
-    function test_swapEEthForStEth_mainnet() public {
-        initializeRealisticFork(MAINNET_FORK);
-        setUpLiquifier(MAINNET_FORK);
 
-        vm.deal(alice, 100 ether);
-        vm.startPrank(alice);
+    function _swapEEthForStEth_mainnet() public {
+        vm.deal(bob, 100 ether);
+        vm.startPrank(bob);
         liquidityPoolInstance.deposit{value: 100 ether}();
 
         eETHInstance.approve(address(liquifierInstance), 50 ether);
 
         uint256 beforeTVL = liquidityPoolInstance.getTotalPooledEther();
         uint256 beforeLiquifierTotalPooledEther = liquifierInstance.getTotalPooledEther();
-
+        uint256 totalValueOutOfLpBefore = liquidityPoolInstance.totalValueOutOfLp(); 
+        uint256 totalValueInLpBefore = liquidityPoolInstance.totalValueInLp(); 
+        uint256 preBalance = IERC20(liquifierInstance.lido()).balanceOf(address(liquifierInstance));
         liquifierInstance.swapEEthForStEth(50 ether);
-        vm.stopPrank();
-
+        uint256 postBalance = IERC20(liquifierInstance.lido()).balanceOf(address(liquifierInstance));
+        uint256 changeInBalance = preBalance - postBalance;
         uint256 afterTVL = liquidityPoolInstance.getTotalPooledEther();
         uint256 afterLiquifierTotalPooledEther = liquifierInstance.getTotalPooledEther();
-
-        assertApproxEqAbs(afterTVL, beforeTVL, 1);
-        assertApproxEqAbs(beforeLiquifierTotalPooledEther, afterLiquifierTotalPooledEther, 1);
+        uint256 totalValueOutOfLpAfter = liquidityPoolInstance.totalValueOutOfLp();
+        uint256 totalValueInLpAfter = liquidityPoolInstance.totalValueInLp();  
+        vm.stopPrank();
+        // Check LP eth balance remains same, eeth supply decreases, liquifier total pooled ether should decrease
+        assertApproxEqAbs(IERC20(address(liquifierInstance.lido())).balanceOf(bob), changeInBalance, 1);
+        assertApproxEqAbs(afterTVL + changeInBalance, beforeTVL, 1);
+        assertApproxEqAbs(afterLiquifierTotalPooledEther + changeInBalance, beforeLiquifierTotalPooledEther, 1);
+        assertApproxEqAbs(totalValueOutOfLpAfter, totalValueOutOfLpBefore, changeInBalance);
+        assertApproxEqAbs(totalValueInLpAfter, totalValueInLpBefore, 1);
     }
 
-    function test_withdrawEEth() public {
-        test_swapEEthForStEth_mainnet();
-
-        _upgrade_liquidity_pool_contract();
-
-        uint256 beforeTVL = liquidityPoolInstance.getTotalPooledEther();
-        uint256 beforeLiquifierTotalPooledEther = liquifierInstance.getTotalPooledEther();
-
-        vm.prank(alice);
-        liquifierInstance.withdrawEEth(10 ether);
-
-        uint256 afterTVL = liquidityPoolInstance.getTotalPooledEther();
-        uint256 afterLiquifierTotalPooledEther = liquifierInstance.getTotalPooledEther();
-
-        assertApproxEqAbs(int256(afterTVL) - int256(beforeTVL), int256(afterLiquifierTotalPooledEther) - int256(beforeLiquifierTotalPooledEther), 1);
+    //same as no fees for whitelisted user
+    function test_whitelisted() public {
+        initializeRealisticFork(MAINNET_FORK);
+        setUpLiquifier(MAINNET_FORK);
+        vm.startPrank(superAdmin);
+        roleRegistry.grantRole(liquifierInstance.EETH_STETH_SWAPPER(), bob);
+        vm.stopPrank();
+        _swapEEthForStEth_mainnet();
     }
 }

test/TestSetup.sol

@@ -1446,7 +1446,6 @@ contract TestSetup is Test {
         uint256[] memory hints = liquifierInstance.lidoWithdrawalQueue().findCheckpointHints(reqIds, 1, lastCheckPointIndex);
         liquifierInstance.stEthClaimWithdrawals(reqIds, hints);
 
-        // The ether.fi admin withdraws the ETH from the liquifier contract to the liquidity pool contract
         liquifierInstance.withdrawEther();
         vm.stopPrank();
     }