Fix accounting bug, code quality

script/DeployPuffETH.s.sol

@@ -11,6 +11,7 @@ import { PufferVault } from "src/PufferVault.sol";
 import { NoImplementation } from "src/NoImplementation.sol";
 import { PufferDeployment } from "src/structs/PufferDeployment.sol";
 import { AccessManager } from "openzeppelin/access/manager/AccessManager.sol";
+import { TimelockController } from "openzeppelin/governance/TimelockController.sol";
 
 /**
  * @title DeployPuffer
@@ -25,7 +26,7 @@ import { AccessManager } from "openzeppelin/access/manager/AccessManager.sol";
  *         Other scripts will fail because addresses will be updated in deployments file, but the deployment never happened.
  *
  *         BaseScript.sol holds the private key logic, if you don't have `PK` ENV variable, it will use the default one PK from `makeAddr("pufferDeployer")`
- * 
+ *
  *         PK=${deployer_pk} forge script script/DeployPuffETH.s.sol:DeployPuffETH -vvvv --rpc-url=... --broadcast
  */
 contract DeployPuffETH is BaseScript {

src/PufferDepositor.sol

@@ -10,11 +10,13 @@ import { ERC20Upgradeable } from "@openzeppelin-contracts-upgradeable/token/ERC2
 import { ERC20PermitUpgradeable } from
     "@openzeppelin-contracts-upgradeable/token/ERC20/extensions/ERC20PermitUpgradeable.sol";
 import { UUPSUpgradeable } from "@openzeppelin-contracts-upgradeable/proxy/utils/UUPSUpgradeable.sol";
-import { IStETH } from "src/interface/IStETH.sol";
+import { IStETH } from "src/interface/Lido/IStETH.sol";
+import { IWstETH } from "src/interface/Lido/IWstETH.sol";
 import { PufferVault } from "src/PufferVault.sol";
 import { SafeERC20 } from "@openzeppelin/contracts/token/ERC20/utils/SafeERC20.sol";
-import { IEigenLayer, IStrategy } from "src/interface/IEigenLayer.sol";
-import { ISushiRouter } from "src/interface/ISushiRouter.sol";
+import { IEigenLayer, IStrategy } from "src/interface/EigenLayer/IEigenLayer.sol";
+import { ISushiRouter } from "src/interface/Other/ISushiRouter.sol";
+import { IPufferDepositor } from "src/interface/IPufferDepositor.sol";
 import { FixedPointMathLib } from "solady/utils/FixedPointMathLib.sol";
 import { console } from "forge-std/console.sol";
 
@@ -23,44 +25,16 @@ import { console } from "forge-std/console.sol";
  * @author Puffer Finance
  * @custom:security-contact [email protected]
  */
-contract PufferDepositor is AccessManagedUpgradeable, UUPSUpgradeable {
+contract PufferDepositor is IPufferDepositor, AccessManagedUpgradeable, UUPSUpgradeable {
     using SafeERC20 for address;
 
-    /**
-     * @dev Error indicating that the token is not allowed.
-     */
-    error TokenNotAllowed(address token);
-
-    /**
-     * @dev Event indicating that the token is allowed.
-     */
-    event TokenAllowed(IERC20 token);
-    /**
-     * @dev Event indicating that the token is disallowed.
-     */
-    event TokenDisallowed(IERC20 token);
-
-    /**
-     * @dev Struct representing a permit for a specific action.
-     */
-    struct Permit {
-        address owner;
-        uint256 deadline;
-        uint256 amount;
-        uint8 v;
-        bytes32 r;
-        bytes32 s;
-    }
-
     IERC20 internal constant _USDT = IERC20(0xdAC17F958D2ee523a2206206994597C13D831ec7);
     IERC20 internal constant _USDC = IERC20(0xA0b86991c6218b36c1d19D4a2e9Eb0cE3606eB48);
 
     IStETH internal constant _ST_ETH = IStETH(0xae7ab96520DE3A18E5e111B5EaAb095312D7fE84);
-    ISushiRouter internal constant _SUSHI_ROUTER = ISushiRouter(0x5550D13389bB70F45fCeF58f19f6b6e87F6e747d);
+    IWstETH internal constant _WST_ETH = IWstETH(0x7f39C581F595B53c5cb19bD0b3f8dA6c935E2Ca0);
 
-    // Sushi router uses this address to represent native ETH
-    // slither-disable-next-line unused-state
-    address constant _ETH_NATIVE_ADDRESS = 0xEeeeeEeeeEeEeeEeEeEeeEEEeeeeEeeeeeeeEEeE;
+    ISushiRouter internal constant _SUSHI_ROUTER = ISushiRouter(0x5550D13389bB70F45fCeF58f19f6b6e87F6e747d);
 
     /**
      * @dev The Puffer Vault contract address
@@ -146,7 +120,7 @@ contract PufferDepositor is AccessManagedUpgradeable, UUPSUpgradeable {
     function swapAndDepositWithPermit(
         address tokenIn,
         uint256 amountOutMin,
-        Permit calldata permitData,
+        IPufferDepositor.Permit calldata permitData,
         bytes calldata routeCode
     ) public virtual returns (uint256 pufETHAmount) {
         DepositorStorage storage $ = _getDepositorStorage();
@@ -182,6 +156,28 @@ contract PufferDepositor is AccessManagedUpgradeable, UUPSUpgradeable {
         return _PUFFER_VAULT.deposit(stETHAmountOut, msg.sender);
     }
 
+    /**
+     * @notice Deposits wrapped stETH (wstETH) into the Puffer Vault
+     * @param permitData The permit data containing the approval information
+     * @return pufETHAmount The amount of pufETH received from the deposit
+     */
+    function depositWstETH(IPufferDepositor.Permit calldata permitData) external returns (uint256 pufETHAmount) {
+        try ERC20Permit(address(_WST_ETH)).permit({
+            owner: permitData.owner,
+            spender: address(this),
+            value: permitData.amount,
+            deadline: permitData.deadline,
+            v: permitData.v,
+            s: permitData.s,
+            r: permitData.r
+        }) { } catch { }
+
+        SafeERC20.safeTransferFrom(IERC20(address(_WST_ETH)), msg.sender, address(this), permitData.amount);
+        uint256 stETHAmount = _WST_ETH.unwrap(permitData.amount);
+
+        return _PUFFER_VAULT.deposit(stETHAmount, msg.sender);
+    }
+
     /**
      * @notice Allows the specified token for deposit
      * Restricted access

src/PufferVault.sol

@@ -1,12 +1,15 @@
 // SPDX-License-Identifier: GPL-3.0
 pragma solidity >=0.8.0 <0.9.0;
 
-import { IStETH } from "src/interface/IStETH.sol";
-import { ILidoWithdrawalQueue } from "src/interface/ILidoWithdrawalQueue.sol";
+import { IPufferVault } from "src/interface/IPufferVault.sol";
+import { IStETH } from "src/interface/Lido/IStETH.sol";
+import { ILidoWithdrawalQueue } from "src/interface/Lido/ILidoWithdrawalQueue.sol";
+import { PufferVaultStorage } from "src/PufferVaultStorage.sol";
 import { SafeERC20 } from "@openzeppelin/contracts/token/ERC20/utils/SafeERC20.sol";
 import { IERC721Receiver } from "@openzeppelin/contracts/token/ERC721/IERC721Receiver.sol";
-import { IEigenLayer, IStrategy } from "src/interface/IEigenLayer.sol";
+import { IEigenLayer, IStrategy } from "src/interface/EigenLayer/IEigenLayer.sol";
 import { UUPSUpgradeable } from "@openzeppelin-contracts-upgradeable/proxy/utils/UUPSUpgradeable.sol";
+import { EnumerableSet } from "openzeppelin/utils/structs/EnumerableSet.sol";
 import { AccessManagedUpgradeable } from
     "@openzeppelin-contracts-upgradeable/access/manager/AccessManagedUpgradeable.sol";
 import {
@@ -20,26 +23,26 @@ import { ERC4626Upgradeable } from "@openzeppelin-contracts-upgradeable/token/ER
  * @author Puffer Finance
  * @custom:security-contact [email protected]
  */
-contract PufferVault is IERC721Receiver, AccessManagedUpgradeable, ERC4626Upgradeable, UUPSUpgradeable {
-    //@todo check what to override on ERC4626Upgradeable
-    //@todo should probably add `restricted` to some of the external/public functions so that we can pause/unpause in state of an emergency
+contract PufferVault is
+    IPufferVault,
+    IERC721Receiver,
+    PufferVaultStorage,
+    AccessManagedUpgradeable,
+    ERC4626Upgradeable,
+    UUPSUpgradeable
+{
+    using EnumerableSet for EnumerableSet.UintSet;
     using SafeERC20 for address;
 
     /**
-     * @notice Emitted when we request withdrawals from Lido
+     * @dev Ethereum Mainnet addresses
      */
-    event RequestedWithdrawals(uint256[]);
-
     IStrategy internal constant _EIGEN_STETH_STRATEGY = IStrategy(0x93c4b944D05dfe6df7645A86cd2206016c51564D);
     IEigenLayer internal constant _EIGEN_STRATEGY_MANAGER = IEigenLayer(0x858646372CC42E1A627fcE94aa7A7033e7CF075A);
     IStETH internal constant _ST_ETH = IStETH(0xae7ab96520DE3A18E5e111B5EaAb095312D7fE84);
     ILidoWithdrawalQueue internal constant _LIDO_WITHDRAWAL_QUEUE =
         ILidoWithdrawalQueue(0x889edC2eDab5f40e902b864aD4d7AdE8E412F9B1);
 
-    //@todo Should support depositing of wstETH
-    // slither-disable-next-line unused-state
-    address internal constant _W_STETH = 0x7f39C581F595B53c5cb19bD0b3f8dA6c935E2Ca0;
-
     constructor() {
         // Approve stETH to Lido && EL
         _disableInitializers();
@@ -53,30 +56,48 @@ contract PufferVault is IERC721Receiver, AccessManagedUpgradeable, ERC4626Upgrad
         SafeERC20.safeIncreaseAllowance(_ST_ETH, address(_EIGEN_STRATEGY_MANAGER), type(uint256).max);
     }
 
-    receive() external payable virtual { }
+    receive() external payable virtual {
+        // If we don't use this pattern, somebody can create a Lido withdrawal and claim it to this contract
+        // Making `$.lidoLockedETH -= msg.value` revert
+        VaultStorage storage $ = _getPufferVaultStorage();
+        if ($.isLidoWithdrawal) {
+            $.lidoLockedETH -= msg.value;
+        }
+    }
 
     /**
-     * @notice Not allowed
+     * @notice Claims ETH withdrawals from Lido
+     * @param requestIds An array of request IDs for the withdrawals
      */
-    function redeem(uint256, address, address) public virtual override returns (uint256) {
-        //@todo figure out
-        revert("Not allowed");
+    function claimWithdrawalsFromLido(uint256[] calldata requestIds) external virtual {
+        VaultStorage storage $ = _getPufferVaultStorage();
+
+        // Tell our receive() that we are doing a Lido claim
+        $.isLidoWithdrawal = true;
+
+        //@todo can probably be optimized
+        for (uint256 i = 0; i < requestIds.length; ++i) {
+            // slither-disable-next-line calls-loop
+            _LIDO_WITHDRAWAL_QUEUE.claimWithdrawal(requestIds[i]);
+        }
+
+        // Reset back the value
+        $.isLidoWithdrawal = false;
+        emit ClaimedWithdrawals(requestIds);
     }
 
     /**
      * @notice Not allowed
      */
-    function withdraw(uint256, address, address) public virtual override returns (uint256) {
-        //@todo figure out
-        revert("Not allowed");
+    function redeem(uint256, address, address) public virtual override returns (uint256) {
+        revert WithdrawalsAreDisabled();
     }
 
     /**
-     * notice Deposits stETH into `stETH` EigenLayer strategy
-     * @param amount the amount of stETH to deposit
+     * @notice Not allowed
      */
-    function depositToEigenLayer(uint256 amount) external virtual restricted {
-        _EIGEN_STRATEGY_MANAGER.depositIntoStrategy({ strategy: _EIGEN_STETH_STRATEGY, token: _ST_ETH, amount: amount });
+    function withdraw(uint256, address, address) public virtual override returns (uint256) {
+        revert WithdrawalsAreDisabled();
     }
 
     /**
@@ -86,33 +107,33 @@ contract PufferVault is IERC721Receiver, AccessManagedUpgradeable, ERC4626Upgrad
      * Sum of EL assets + Vault Assets
      */
     function totalAssets() public view virtual override returns (uint256) {
-        return _ST_ETH.balanceOf(address(this)) + getELBackingEthAmount() + address(this).balance;
+        return _ST_ETH.balanceOf(address(this)) + getELBackingEthAmount() + getPendingLidoETHAmount()
+            + address(this).balance;
     }
 
     /**
      * @notice Returns the ETH amount that is backing this vault locked in EigenLayer stETH strategy
      */
     function getELBackingEthAmount() public view virtual returns (uint256 ethAmount) {
-        //@todo optimize, no loop, look on the strategy directly or something
-        uint256 elShares;
         // EigenLayer returns the number of shares owned in that strategy
-        (IStrategy[] memory strategies, uint256[] memory amounts) = _EIGEN_STRATEGY_MANAGER.getDeposits(address(this));
-        for (uint256 i = 0; i < strategies.length; ++i) {
-            if (address(strategies[i]) == address(_EIGEN_STETH_STRATEGY)) {
-                elShares = amounts[i];
-                break;
-            }
-        }
+        return _EIGEN_STETH_STRATEGY.userUnderlying(address(this));
+    }
 
-        // No deposits to EL
-        if (elShares == 0) {
-            return 0;
-        }
+    /**
+     * @notice Returns the amount of ETH that is pending withdrawal from Lido
+     * @return The amount of ETH pending withdrawal
+     */
+    function getPendingLidoETHAmount() public view virtual returns (uint256) {
+        VaultStorage storage $ = _getPufferVaultStorage();
+        return $.lidoLockedETH;
+    }
 
-        // ETH is 1:1 with stETH
-        // EL Keeps track of deposits in their own shares
-        // This is how we get the stETHAmount owned in EL
-        ethAmount = (elShares * _ST_ETH.balanceOf(address(_EIGEN_STETH_STRATEGY))) / _EIGEN_STETH_STRATEGY.totalShares();
+    /**
+     * notice Deposits stETH into `stETH` EigenLayer strategy
+     * @param amount the amount of stETH to deposit
+     */
+    function depositToEigenLayer(uint256 amount) external virtual restricted {
+        _EIGEN_STRATEGY_MANAGER.depositIntoStrategy({ strategy: _EIGEN_STETH_STRATEGY, token: _ST_ETH, amount: amount });
     }
 
     /**
@@ -125,23 +146,19 @@ contract PufferVault is IERC721Receiver, AccessManagedUpgradeable, ERC4626Upgrad
         restricted
         returns (uint256[] memory requestIds)
     {
+        VaultStorage storage $ = _getPufferVaultStorage();
+
+        uint256 lockedAmount;
+        for (uint256 i = 0; i < amounts.length; ++i) {
+            lockedAmount += amounts[i];
+        }
+        $.lidoLockedETH += lockedAmount;
+
         requestIds = _LIDO_WITHDRAWAL_QUEUE.requestWithdrawals(amounts, address(this));
         emit RequestedWithdrawals(requestIds);
         return requestIds;
     }
 
-    /**
-     * @notice Claims ETH withdrawals from Lido
-     * @param requestIds An array of request IDs for the withdrawals
-     */
-    function claimWithdrawalsFromLido(uint256[] calldata requestIds) external virtual {
-        //@todo can probably be optimized
-        for (uint256 i = 0; i < requestIds.length; ++i) {
-            // slither-disable-next-line calls-loop
-            _LIDO_WITHDRAWAL_QUEUE.claimWithdrawal(requestIds[i]);
-        }
-    }
-
     /**
      * @notice Required by the ERC721 Standard
      */

src/PufferVaultMainnet.sol

@@ -2,9 +2,7 @@
 pragma solidity >=0.8.0 <0.9.0;
 
 import { PufferVault } from "src/PufferVault.sol";
-import { IWETH } from "src/interface/IWETH.sol";
-import { ILidoWithdrawalQueue } from "src/interface/ILidoWithdrawalQueue.sol";
-import { SafeERC20 } from "@openzeppelin/contracts/token/ERC20/utils/SafeERC20.sol";
+import { IWETH } from "src/interface/Other/IWETH.sol";
 
 /**
  * @title PufferVault

src/PufferVaultStorage.sol

@@ -0,0 +1,19 @@
+// SPDX-License-Identifier: GPL-3.0
+pragma solidity >=0.8.0 <0.9.0;
+
+abstract contract PufferVaultStorage {
+    /// @custom:storage-location erc7201:puffervault.storage.ERC4626
+    struct VaultStorage {
+        uint256 lidoLockedETH;
+        bool isLidoWithdrawal;
+    }
+
+    // keccak256(abi.encode(uint256(keccak256("puffervault.storage.ERC4626")) - 1)) & ~bytes32(uint256(0xff))
+    bytes32 private constant ERC4626StorageLocation = 0x6d4971415142040fa945ebf44b5dec920e7693eb61c9c44e4167ab643762ec00;
+
+    function _getPufferVaultStorage() internal pure returns (VaultStorage storage $) {
+        assembly {
+            $.slot := ERC4626StorageLocation
+        }
+    }
+}

src/interface/IEigenLayer.sol → src/interface/EigenLayer/IEigenLayer.sol

@@ -8,6 +8,8 @@ interface IStrategy {
      * @notice Returns the total number of shares in the EL Strategy
      */
     function totalShares() external view returns (uint256);
+
+    function userUnderlying(address user) external view returns (uint256);
 }
 
 // packed struct for queued withdrawals; helps deal with stack-too-deep errors

src/interface/IPufferDepositor.sol

@@ -0,0 +1,37 @@
+// SPDX-License-Identifier: GPL-3.0
+pragma solidity >=0.8.0 <0.9.0;
+
+import { IERC20 } from "openzeppelin/token/ERC20/IERC20.sol";
+
+/**
+ * @title PufferDepositor
+ * @author Puffer Finance
+ * @custom:security-contact [email protected]
+ */
+interface IPufferDepositor {
+    /**
+     * @dev Error indicating that the token is not allowed.
+     */
+    error TokenNotAllowed(address token);
+
+    /**
+     * @dev Event indicating that the token is allowed.
+     */
+    event TokenAllowed(IERC20 token);
+    /**
+     * @dev Event indicating that the token is disallowed.
+     */
+    event TokenDisallowed(IERC20 token);
+
+    /**
+     * @dev Struct representing a permit for a specific action.
+     */
+    struct Permit {
+        address owner;
+        uint256 deadline;
+        uint256 amount;
+        uint8 v;
+        bytes32 r;
+        bytes32 s;
+    }
+}