diff --git a/src/WrapperScripts.sol b/src/WrapperScripts.sol index bf5f888..66ddf83 100644 --- a/src/WrapperScripts.sol +++ b/src/WrapperScripts.sol @@ -10,10 +10,24 @@ contract WrapperActions { IWETH(weth).deposit{value: amount}(); } + function wrapETHUpTo(address weth, uint256 targetAmount) external payable { + uint256 currentBalance = IERC20(weth).balanceOf(address(this)); + if (currentBalance < targetAmount) { + IWETH(weth).deposit{value: targetAmount - currentBalance}(); + } + } + function unwrapWETH(address weth, uint256 amount) external { IWETH(weth).withdraw(amount); } + function unwrapWETHUpTo(address weth, uint256 targetAmount) external { + uint256 currentBalance = address(this).balance; + if (currentBalance < targetAmount) { + IWETH(weth).withdraw(targetAmount - currentBalance); + } + } + function wrapLidoStETH(address wstETH, address stETH, uint256 amount) external { IERC20(stETH).approve(wstETH, amount); IWstETH(wstETH).wrap(amount); diff --git a/test/WrapperScripts.t.sol b/test/WrapperScripts.t.sol index 320f577..2f87c83 100644 --- a/test/WrapperScripts.t.sol +++ b/test/WrapperScripts.t.sol @@ -65,6 +65,56 @@ contract WrapperScriptsTest is Test { assertEq(address(wallet).balance, 0 ether); } + function testWrapETHUpTo() public { + vm.pauseGasMetering(); + QuarkWallet wallet = QuarkWallet(factory.create(alice, address(0))); + + deal(address(wallet), 10 ether); + deal(WETH, address(wallet), 7 ether); + + QuarkWallet.QuarkOperation memory op = new QuarkOperationHelper().newBasicOpWithCalldata( + wallet, + wrapperScript, + abi.encodeWithSelector(WrapperActions.wrapETHUpTo.selector, WETH, 10 ether), + ScriptType.ScriptSource + ); + bytes memory signature = new SignatureHelper().signOp(alicePrivateKey, wallet, op); + + assertEq(IERC20(WETH).balanceOf(address(wallet)), 7 ether); + assertEq(address(wallet).balance, 10 ether); + + vm.resumeGasMetering(); + wallet.executeQuarkOperation(op, signature); + + assertEq(IERC20(WETH).balanceOf(address(wallet)), 10 ether); + assertEq(address(wallet).balance, 7 ether); + } + + function testWrapETHUpToDoesNotWrapIfNotNeeded() public { + vm.pauseGasMetering(); + QuarkWallet wallet = QuarkWallet(factory.create(alice, address(0))); + + deal(address(wallet), 10 ether); + deal(WETH, address(wallet), 10 ether); + + QuarkWallet.QuarkOperation memory op = new QuarkOperationHelper().newBasicOpWithCalldata( + wallet, + wrapperScript, + abi.encodeWithSelector(WrapperActions.wrapETHUpTo.selector, WETH, 10 ether), + ScriptType.ScriptSource + ); + bytes memory signature = new SignatureHelper().signOp(alicePrivateKey, wallet, op); + + assertEq(IERC20(WETH).balanceOf(address(wallet)), 10 ether); + assertEq(address(wallet).balance, 10 ether); + + vm.resumeGasMetering(); + wallet.executeQuarkOperation(op, signature); + + assertEq(IERC20(WETH).balanceOf(address(wallet)), 10 ether); + assertEq(address(wallet).balance, 10 ether); + } + function testUnwrapWETH() public { vm.pauseGasMetering(); QuarkWallet wallet = QuarkWallet(factory.create(alice, address(0))); @@ -87,6 +137,56 @@ contract WrapperScriptsTest is Test { assertEq(address(wallet).balance, 10 ether); } + function testUnwrapWETHUpTo() public { + vm.pauseGasMetering(); + QuarkWallet wallet = QuarkWallet(factory.create(alice, address(0))); + + deal(WETH, address(wallet), 10 ether); + deal(address(wallet), 7 ether); + + QuarkWallet.QuarkOperation memory op = new QuarkOperationHelper().newBasicOpWithCalldata( + wallet, + wrapperScript, + abi.encodeWithSelector(WrapperActions.unwrapWETHUpTo.selector, WETH, 10 ether), + ScriptType.ScriptSource + ); + bytes memory signature = new SignatureHelper().signOp(alicePrivateKey, wallet, op); + + assertEq(IERC20(WETH).balanceOf(address(wallet)), 10 ether); + assertEq(address(wallet).balance, 7 ether); + + vm.resumeGasMetering(); + wallet.executeQuarkOperation(op, signature); + + assertEq(IERC20(WETH).balanceOf(address(wallet)), 7 ether); + assertEq(address(wallet).balance, 10 ether); + } + + function testUnwrapWETHUpToDoesNotUnwrapIfNotNeeded() public { + vm.pauseGasMetering(); + QuarkWallet wallet = QuarkWallet(factory.create(alice, address(0))); + + deal(WETH, address(wallet), 10 ether); + deal(address(wallet), 10 ether); + + QuarkWallet.QuarkOperation memory op = new QuarkOperationHelper().newBasicOpWithCalldata( + wallet, + wrapperScript, + abi.encodeWithSelector(WrapperActions.unwrapWETHUpTo.selector, WETH, 10 ether), + ScriptType.ScriptSource + ); + bytes memory signature = new SignatureHelper().signOp(alicePrivateKey, wallet, op); + + assertEq(IERC20(WETH).balanceOf(address(wallet)), 10 ether); + assertEq(address(wallet).balance, 10 ether); + + vm.resumeGasMetering(); + wallet.executeQuarkOperation(op, signature); + + assertEq(IERC20(WETH).balanceOf(address(wallet)), 10 ether); + assertEq(address(wallet).balance, 10 ether); + } + function testWrapStETH() public { vm.pauseGasMetering(); QuarkWallet wallet = QuarkWallet(factory.create(alice, address(0)));