From f5838ae5e31ee3a6cfabb506c119cbe4f72917e0 Mon Sep 17 00:00:00 2001 From: samparsky <8148384+samparsky@users.noreply.github.com> Date: Tue, 26 Sep 2023 10:43:29 +0100 Subject: [PATCH 1/6] chore: move stETH and wstETH addresses to LidoSplit --- src/lido/LidoSplit.sol | 40 +++++++++++----------------- src/lido/LidoSplitFactory.sol | 12 ++------- src/test/lido/LIdoSplitFactory.t.sol | 6 ++--- src/test/lido/LidoSplit.t.sol | 8 ++---- 4 files changed, 23 insertions(+), 43 deletions(-) diff --git a/src/lido/LidoSplit.sol b/src/lido/LidoSplit.sol index 8156d3e..3c68bc5 100644 --- a/src/lido/LidoSplit.sol +++ b/src/lido/LidoSplit.sol @@ -24,17 +24,24 @@ contract LidoSplit is Clone { /// 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; + + /// ----------------------------------------------------------------------- + /// storage + /// ----------------------------------------------------------------------- + + /// @notice stETH token + ERC20 public immutable stETH; + + /// @notice wstETH token + ERC20 public immutable wstETH; - constructor() {} + 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 @@ -42,25 +49,10 @@ contract LidoSplit is Clone { 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 diff --git a/src/lido/LidoSplitFactory.sol b/src/lido/LidoSplitFactory.sol index a31310e..922827c 100644 --- a/src/lido/LidoSplitFactory.sol +++ b/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); } diff --git a/src/test/lido/LIdoSplitFactory.t.sol b/src/test/lido/LIdoSplitFactory.t.sol index a42f44c..9ffb7cb 100644 --- a/src/test/lido/LIdoSplitFactory.t.sol +++ b/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"); } diff --git a/src/test/lido/LidoSplit.t.sol b/src/test/lido/LidoSplit.t.sol index d39b4e6..89c3d00 100644 --- a/src/test/lido/LidoSplit.t.sol +++ b/src/test/lido/LidoSplit.t.sol @@ -28,8 +28,8 @@ contract LidoSplitTest is LidoSplitTestHelper, Test { 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_CanDistribute() public { @@ -46,10 +46,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"); } } From 4e007c36a71b496f742b013679fc9bf47ce7a1cf Mon Sep 17 00:00:00 2001 From: samparsky <8148384+samparsky@users.noreply.github.com> Date: Tue, 26 Sep 2023 10:48:11 +0100 Subject: [PATCH 2/6] chore: add rescueETH function to LidoSplit --- src/lido/LidoSplit.sol | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/src/lido/LidoSplit.sol b/src/lido/LidoSplit.sol index 3c68bc5..4b86033 100644 --- a/src/lido/LidoSplit.sol +++ b/src/lido/LidoSplit.sol @@ -19,6 +19,7 @@ contract LidoSplit is Clone { /// libraries /// ----------------------------------------------------------------------- using SafeTransferLib for ERC20; + using SafeTransferLib for address; /// ----------------------------------------------------------------------- /// storage - cwia offsets @@ -62,4 +63,11 @@ contract LidoSplit is Clone { // transfer to split wallet ERC20(wstETH).safeTransfer(splitWallet(), amount); } + + /// @notice Rescue stuck ETH + /// @return balance Amount of ETH rescued + function rescueETH() external returns (uint256 balance) { + balance = address(this).balance; + splitWallet().safeTransferETH(balance); + } } From 8b0138afdbb0d77d070de0834bd640f46b04b1a6 Mon Sep 17 00:00:00 2001 From: samparsky <8148384+samparsky@users.noreply.github.com> Date: Tue, 26 Sep 2023 11:00:39 +0100 Subject: [PATCH 3/6] chore: change rescueETH to rescueFunds to enable rescue of tokens --- src/lido/LidoSplit.sol | 22 +++++++++++++++++++--- src/test/lido/LidoSplit.t.sol | 6 ++++++ 2 files changed, 25 insertions(+), 3 deletions(-) diff --git a/src/lido/LidoSplit.sol b/src/lido/LidoSplit.sol index 4b86033..3b40106 100644 --- a/src/lido/LidoSplit.sol +++ b/src/lido/LidoSplit.sol @@ -15,12 +15,17 @@ 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 /// ----------------------------------------------------------------------- @@ -29,6 +34,7 @@ contract LidoSplit is Clone { // 0; first item uint256 internal constant SPLIT_WALLET_ADDRESS_OFFSET = 0; + /// ----------------------------------------------------------------------- /// storage /// ----------------------------------------------------------------------- @@ -65,9 +71,19 @@ contract LidoSplit is Clone { } /// @notice Rescue stuck ETH + /// Uses token == address(0) to represent ETH /// @return balance Amount of ETH rescued - function rescueETH() external returns (uint256 balance) { - balance = address(this).balance; - splitWallet().safeTransferETH(balance); + function rescueFunds(address token) external returns (uint256 balance) { + if (token == address(stETH) || token == address(wstETH)) { + 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); + } } } diff --git a/src/test/lido/LidoSplit.t.sol b/src/test/lido/LidoSplit.t.sol index 89c3d00..ca7d83f 100644 --- a/src/test/lido/LidoSplit.t.sol +++ b/src/test/lido/LidoSplit.t.sol @@ -32,6 +32,12 @@ contract LidoSplitTest is LidoSplitTestHelper, Test { assertEq(address(lidoSplit.wstETH()), WSTETH_MAINNET_ADDRESS, "invalid wstETH address"); } + function test_CanRescueETH() public { + deal(lidoSplit.splitWallet(), 1 ether); + + + } + function test_CanDistribute() public { // we use a random account on Etherscan to credit the lidoSplit address // with 10 ether worth of stETH on mainnet From d781b30971c572c78ec73d9d40b5c371d7b16039 Mon Sep 17 00:00:00 2001 From: samparsky <8148384+samparsky@users.noreply.github.com> Date: Tue, 26 Sep 2023 11:11:15 +0100 Subject: [PATCH 4/6] test: fix LidoSplit test suite; add rescueFunds test cases --- src/test/lido/LidoSplit.t.sol | 37 ++++++++++++++++++++++++++++++++--- 1 file changed, 34 insertions(+), 3 deletions(-) diff --git a/src/test/lido/LidoSplit.t.sol b/src/test/lido/LidoSplit.t.sol index ca7d83f..8bd504f 100644 --- a/src/test/lido/LidoSplit.t.sol +++ b/src/test/lido/LidoSplit.t.sol @@ -5,6 +5,8 @@ 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; @@ -12,6 +14,8 @@ contract LidoSplitTest is LidoSplitTestHelper, Test { address demoSplit; + MockERC20 mERC20; + function setUp() public { uint256 mainnetBlock = 17_421_005; vm.createSelectFork(getChain("mainnet").rpcUrl, mainnetBlock); @@ -24,6 +28,9 @@ 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 { @@ -32,10 +39,34 @@ contract LidoSplitTest is LidoSplitTestHelper, Test { assertEq(address(lidoSplit.wstETH()), WSTETH_MAINNET_ADDRESS, "invalid wstETH address"); } - function test_CanRescueETH() public { - deal(lidoSplit.splitWallet(), 1 ether); + 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)); + + vm.expectRevert( + LidoSplit.Invalid_Address.selector + ); + lidoSplit.rescueFunds(address(WSTETH_MAINNET_ADDRESS)); } function test_CanDistribute() public { From bf088de1feeac93afbcbda41920bed36beeae5e0 Mon Sep 17 00:00:00 2001 From: samparsky <8148384+samparsky@users.noreply.github.com> Date: Wed, 27 Sep 2023 07:10:07 +0100 Subject: [PATCH 5/6] chore: enable rescueFunds for wstETH --- lib/solady | 2 +- src/lido/LidoSplit.sol | 4 +--- 2 files changed, 2 insertions(+), 4 deletions(-) diff --git a/lib/solady b/lib/solady index 2b33744..77809c1 160000 --- a/lib/solady +++ b/lib/solady @@ -1 +1 @@ -Subproject commit 2b33744067c2afa74a6ee364004e99035cae692e +Subproject commit 77809c18e010b914dde9518956a4ae7cb507d383 diff --git a/src/lido/LidoSplit.sol b/src/lido/LidoSplit.sol index 3b40106..ef69379 100644 --- a/src/lido/LidoSplit.sol +++ b/src/lido/LidoSplit.sol @@ -74,9 +74,7 @@ contract LidoSplit is Clone { /// 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) || token == address(wstETH)) { - revert Invalid_Address(); - } + if (token == address(stETH)) revert Invalid_Address(); if (token == ETH_ADDRESS) { balance = address(this).balance; From 91388dabaf376e5c791dd77830a7c8a2cd311b0d Mon Sep 17 00:00:00 2001 From: samparsky <8148384+samparsky@users.noreply.github.com> Date: Wed, 27 Sep 2023 07:10:47 +0100 Subject: [PATCH 6/6] test: fix testCannot_RescueLidoTokens test case --- src/test/lido/LidoSplit.t.sol | 5 ----- 1 file changed, 5 deletions(-) diff --git a/src/test/lido/LidoSplit.t.sol b/src/test/lido/LidoSplit.t.sol index 8bd504f..37a5d33 100644 --- a/src/test/lido/LidoSplit.t.sol +++ b/src/test/lido/LidoSplit.t.sol @@ -62,11 +62,6 @@ contract LidoSplitTest is LidoSplitTestHelper, Test { LidoSplit.Invalid_Address.selector ); lidoSplit.rescueFunds(address(STETH_MAINNET_ADDRESS)); - - vm.expectRevert( - LidoSplit.Invalid_Address.selector - ); - lidoSplit.rescueFunds(address(WSTETH_MAINNET_ADDRESS)); } function test_CanDistribute() public {