Lido Split Audit Fix

src/lido/LidoSplit.sol

@@ -15,52 +15,51 @@ interface IwSTETH {
 /// stETH token to wstETH token because stETH is a rebasing token
 /// @dev Wraps stETH to wstETH and transfers to defined SplitWallet address
 contract LidoSplit is Clone {
+
+  error Invalid_Address();
+
   /// -----------------------------------------------------------------------
   /// libraries
   /// -----------------------------------------------------------------------
   using SafeTransferLib for ERC20;
+  using SafeTransferLib for address;
+
+  address internal constant ETH_ADDRESS = address(0);
 
   /// -----------------------------------------------------------------------
   /// storage - cwia offsets
   /// -----------------------------------------------------------------------
 
-  // stETH (address, 20 bytes),
-  // 0; first item
-  uint256 internal constant ST_ETH_ADDRESS_OFFSET = 0;
-  // wstETH (address, 20 bytees)
-  // 20 = st_eth_offset(0) + st_eth_address_size(address, 20 bytes)
-  uint256 internal constant WST_ETH_ADDRESS_OFFSET = 20;
   // splitWallet (adress, 20 bytes)
-  // 40 = wst_eth_offset(20) + wst_eth_size(address, 20 bytes)
-  uint256 internal constant SPLIT_WALLET_ADDRESS_OFFSET = 40;
+  // 0; first item
+  uint256 internal constant SPLIT_WALLET_ADDRESS_OFFSET = 0;
 
-  constructor() {}
+
+  /// -----------------------------------------------------------------------
+  /// storage
+  /// -----------------------------------------------------------------------
+
+  /// @notice stETH token
+  ERC20 public immutable stETH;
+
+  /// @notice wstETH token
+  ERC20 public immutable wstETH;
+
+  constructor(ERC20 _stETH, ERC20 _wstETH) {
+    stETH = _stETH;
+    wstETH = _wstETH;
+  }
 
   /// Address of split wallet to send funds to to
   /// @dev equivalent to address public immutable splitWallet
   function splitWallet() public pure returns (address) {
     return _getArgAddress(SPLIT_WALLET_ADDRESS_OFFSET);
   }
 
-  /// Address of stETH token
-  /// @dev equivalent to address public immutable stETHAddress
-  function stETHAddress() public pure returns (address) {
-    return _getArgAddress(ST_ETH_ADDRESS_OFFSET);
-  }
-
-  /// Address of wstETH token
-  /// @dev equivalent to address public immutable wstETHAddress
-  function wstETHAddress() public pure returns (address) {
-    return _getArgAddress(WST_ETH_ADDRESS_OFFSET);
-  }
-
   /// Wraps the current stETH token balance to wstETH
   /// transfers the wstETH balance to splitWallet for distribution
   /// @return amount Amount of wstETH transferred to splitWallet
   function distribute() external returns (uint256 amount) {
-    ERC20 stETH = ERC20(stETHAddress());
-    ERC20 wstETH = ERC20(wstETHAddress());
-
     // get current balance
     uint256 balance = stETH.balanceOf(address(this));
     // approve the wstETH
@@ -70,4 +69,19 @@ contract LidoSplit is Clone {
     // transfer to split wallet
     ERC20(wstETH).safeTransfer(splitWallet(), amount);
   }
+
+  /// @notice Rescue stuck ETH
+  /// Uses token == address(0) to represent ETH
+  /// @return balance Amount of ETH rescued
+  function rescueFunds(address token) external returns (uint256 balance) {
+    if (token == address(stETH)) revert Invalid_Address();  
+
+    if (token == ETH_ADDRESS) {
+      balance = address(this).balance;
+      if (balance > 0) splitWallet().safeTransferETH(balance);
+    } else {
+      balance = ERC20(token).balanceOf(address(this));
+      if (balance > 0) ERC20(token).transfer(splitWallet(), balance);
+    }
+  }
 }

src/lido/LidoSplitFactory.sol

@@ -33,19 +33,11 @@ contract LidoSplitFactory {
   /// storage
   /// -----------------------------------------------------------------------
 
-  /// @notice stETH token address
-  ERC20 public immutable stETH;
-
-  /// @notice wstETH token address
-  ERC20 public immutable wstETH;
-
   /// @dev lido split implementation
   LidoSplit public immutable lidoSplitImpl;
 
   constructor(ERC20 _stETH, ERC20 _wstETH) {
-    stETH = _stETH;
-    wstETH = _wstETH;
-    lidoSplitImpl = new LidoSplit();
+    lidoSplitImpl = new LidoSplit(_stETH, _wstETH);
   }
 
   /// Creates a wrapper for splitWallet that transforms stETH token into
@@ -55,7 +47,7 @@ contract LidoSplitFactory {
   function createSplit(address splitWallet) external returns (address lidoSplit) {
     if (splitWallet == address(0)) revert Invalid_Wallet();
 
-    lidoSplit = address(lidoSplitImpl).clone(abi.encodePacked(stETH, wstETH, splitWallet));
+    lidoSplit = address(lidoSplitImpl).clone(abi.encodePacked(splitWallet));
 
     emit CreateLidoSplit(lidoSplit);
   }

src/test/lido/LIdoSplitFactory.t.sol

@@ -18,9 +18,9 @@ contract LidoSplitFactoryTest is LidoSplitTestHelper, Test {
     vm.createSelectFork(getChain("mainnet").rpcUrl, mainnetBlock);
 
     lidoSplitFactory = new LidoSplitFactory(
-            ERC20(STETH_MAINNET_ADDRESS),
-            ERC20(WSTETH_MAINNET_ADDRESS)
-        );
+      ERC20(STETH_MAINNET_ADDRESS),
+      ERC20(WSTETH_MAINNET_ADDRESS)
+    );
 
     demoSplit = makeAddr("demoSplit");
   }

src/test/lido/LidoSplit.t.sol

@@ -5,13 +5,17 @@ import "forge-std/Test.sol";
 import {LidoSplitFactory, LidoSplit} from "src/lido/LidoSplitFactory.sol";
 import {ERC20} from "solmate/tokens/ERC20.sol";
 import {LidoSplitTestHelper} from "./LidoSplitTestHelper.sol";
+import { MockERC20 } from "src/test/utils/mocks/MockERC20.sol";
+
 
 contract LidoSplitTest is LidoSplitTestHelper, Test {
   LidoSplitFactory internal lidoSplitFactory;
   LidoSplit internal lidoSplit;
 
   address demoSplit;
 
+  MockERC20 mERC20;
+
   function setUp() public {
     uint256 mainnetBlock = 17_421_005;
     vm.createSelectFork(getChain("mainnet").rpcUrl, mainnetBlock);
@@ -24,12 +28,40 @@ contract LidoSplitTest is LidoSplitTestHelper, Test {
     demoSplit = makeAddr("demoSplit");
 
     lidoSplit = LidoSplit(lidoSplitFactory.createSplit(demoSplit));
+
+    mERC20 = new MockERC20("Test Token", "TOK", 18);
+    mERC20.mint(type(uint256).max);
   }
 
   function test_CloneArgsIsCorrect() public {
     assertEq(lidoSplit.splitWallet(), demoSplit, "invalid address");
-    assertEq(lidoSplit.stETHAddress(), STETH_MAINNET_ADDRESS, "invalid stETH address");
-    assertEq(lidoSplit.wstETHAddress(), WSTETH_MAINNET_ADDRESS, "invalid wstETH address");
+    assertEq(address(lidoSplit.stETH()), STETH_MAINNET_ADDRESS, "invalid stETH address");
+    assertEq(address(lidoSplit.wstETH()), WSTETH_MAINNET_ADDRESS, "invalid wstETH address");
+  }
+
+  function test_CanRescueFunds() public {
+    // rescue ETH
+    uint256 amountOfEther = 1 ether;
+    deal(address(lidoSplit), amountOfEther);
+
+    uint256 balance = lidoSplit.rescueFunds(address(0));
+    assertEq(balance, amountOfEther, "balance not rescued");
+    assertEq(address(lidoSplit).balance, 0, "balance is not zero");
+    assertEq(address(lidoSplit.splitWallet()).balance, amountOfEther, "rescue not successful");
+
+    // rescue tokens
+    mERC20.transfer(address(lidoSplit), amountOfEther);
+    uint256 tokenBalance = lidoSplit.rescueFunds(address(mERC20));
+    assertEq(tokenBalance, amountOfEther, "token - balance not rescued");
+    assertEq(mERC20.balanceOf(address(lidoSplit)), 0, "token - balance is not zero");
+    assertEq(mERC20.balanceOf(lidoSplit.splitWallet()), amountOfEther, "token - rescue not successful");
+  }
+
+  function testCannot_RescueLidoTokens() public {
+    vm.expectRevert(
+      LidoSplit.Invalid_Address.selector
+    );
+    lidoSplit.rescueFunds(address(STETH_MAINNET_ADDRESS));
   }
 
   function test_CanDistribute() public {
@@ -46,10 +78,6 @@ contract LidoSplitTest is LidoSplitTestHelper, Test {
 
     uint256 afterBalance = ERC20(WSTETH_MAINNET_ADDRESS).balanceOf(demoSplit);
 
-    console.log("checking");
-    console.log(afterBalance);
-    console.log(prevBalance);
-
     assertGe(afterBalance, prevBalance, "after balance greater");
   }
 }