From ec54702df4d899c05a4a367ce64b7ad823a4c101 Mon Sep 17 00:00:00 2001 From: Igor Yalovoy Date: Mon, 23 Dec 2024 15:23:00 -0600 Subject: [PATCH 1/5] Add handling for max uint256 amount in withdrawERC20 and corresponding test case. --- src/access/workflows/WithdrawWorkflow.sol | 5 +++++ test/unit/access/WithdrawWorkflow.t.sol | 12 ++++++++++++ 2 files changed, 17 insertions(+) diff --git a/src/access/workflows/WithdrawWorkflow.sol b/src/access/workflows/WithdrawWorkflow.sol index 3cbd00a0e..b242835f3 100644 --- a/src/access/workflows/WithdrawWorkflow.sol +++ b/src/access/workflows/WithdrawWorkflow.sol @@ -18,6 +18,11 @@ contract WithdrawWorkflow { error NativeWithdrawalFailed(); function withdrawERC20(IERC20 asset, uint256 amount) external { + // If amount is max uint256, set it to the entire balance + if (amount == type(uint256).max) { + amount = asset.balanceOf(address(this)); + } + address owner = _getOwner(); asset.safeTransfer({to: owner, value: amount}); } diff --git a/test/unit/access/WithdrawWorkflow.t.sol b/test/unit/access/WithdrawWorkflow.t.sol index 4281eaa67..a1171f9c8 100644 --- a/test/unit/access/WithdrawWorkflow.t.sol +++ b/test/unit/access/WithdrawWorkflow.t.sol @@ -69,6 +69,18 @@ contract WithdrawWorkflowTest is BaseTest { assertEq(token.balanceOf(_user), defaultAmount); } + function testWithdrawERC20__WhenAmountMax() public { + token.mint(address(accessPoint), defaultAmount); + + bytes memory data = + abi.encodeWithSelector(WithdrawWorkflow.withdrawERC20.selector, IERC20(token), type(uint256).max); + + vm.prank(_user); + accessPoint.execute(address(withdrawWorkflow), data); + + assertEq(token.balanceOf(_user), defaultAmount); + } + function testWithdrawNative() public { vm.deal(address(accessPoint), defaultAmount); From 606a1dc4cbb853d19fba194c691661f09ca51f88 Mon Sep 17 00:00:00 2001 From: Igor Yalovoy Date: Tue, 24 Dec 2024 08:09:06 -0600 Subject: [PATCH 2/5] Update WithdrawWorkflow address in addresses.json for test artifacts 42161. --- test/artifacts/42161/addresses.json | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/test/artifacts/42161/addresses.json b/test/artifacts/42161/addresses.json index 83511364b..df1a7ed8f 100644 --- a/test/artifacts/42161/addresses.json +++ b/test/artifacts/42161/addresses.json @@ -29,7 +29,7 @@ "Viewer": "0x8888886e1d7c1468d7300cF08db89FFE68F29830", "Viewer-impl": "0x80338A3f75614491c8DC383fFaA663b9a27CD05d", "WethWorkflow": "0x7F7c594eE170a62d7e7615972831038Cf7d4Fc1A", - "WithdrawWorkflow": "0x36e6cA034958B2E0D4dC7Ea9a8151f15Ba0B27D2", + "WithdrawWorkflow": "0xbc22c860C1ED7330271eeF19FB47Eb08548f1723", "AaveLendWorkflow": "0xB47Ed636c8296729E81463109FEbf833CeEa71fb", "AaveRepayWorkflow": "0x24f71379C39b515Ff5182F4b0cc298793EC5998c", "AaveWithdrawWorkflow": "0xef4D6687372172c4af1802C208Ab40673b014309", @@ -37,4 +37,4 @@ "AccessManager": "0xacc0003a4aAE5dA4ba12F771C7350D40147Cd7D4", "BridgerV13-impl": "0xbA6FD752CE93879c381fb7ffdbe7baB233D6e6e4", "BridgerV14-impl": "0x363EFf1981E664107EF4E8568Cc4321B74558DAA" -} \ No newline at end of file +} From 27c1b74ddf1e9cac793698ccc168684f743c13ca Mon Sep 17 00:00:00 2001 From: Igor Yalovoy Date: Tue, 24 Dec 2024 08:39:13 -0600 Subject: [PATCH 3/5] Enhance WithdrawWorkflow with WETH support, update constructor, add tests for native withdrawal and error handling. --- .../48-deploy_access_protocol.s.sol | 4 +- src/access/workflows/WithdrawWorkflow.sol | 58 +++++++++++++++++-- test/unit/access/WithdrawWorkflow.t.sol | 33 ++++++++++- 3 files changed, 87 insertions(+), 8 deletions(-) diff --git a/script/migrations/48-deploy_access_protocol.s.sol b/script/migrations/48-deploy_access_protocol.s.sol index 504fdcc74..d2600fbf1 100644 --- a/script/migrations/48-deploy_access_protocol.s.sol +++ b/script/migrations/48-deploy_access_protocol.s.sol @@ -69,7 +69,9 @@ contract DeployAccessProtocolScript is Script, MigrationHelper { ); console2.log("SafeBeaconProxy at: %s", address(safeBeaconProxy)); - withdrawWorkflow = WithdrawWorkflow(create2(abi.encodePacked(type(WithdrawWorkflow).creationCode))); + withdrawWorkflow = WithdrawWorkflow( + create2(abi.encodePacked(type(WithdrawWorkflow).creationCode, abi.encode(getWethByChainId(block.chainid)))) + ); registry.allowWorkflow(address(withdrawWorkflow)); wethWorkflow = WethWorkflow( diff --git a/src/access/workflows/WithdrawWorkflow.sol b/src/access/workflows/WithdrawWorkflow.sol index b242835f3..4ad044a60 100644 --- a/src/access/workflows/WithdrawWorkflow.sol +++ b/src/access/workflows/WithdrawWorkflow.sol @@ -3,20 +3,46 @@ pragma solidity ^0.8.18; import {IERC20} from "@openzeppelin-5.0.1/contracts/token/ERC20/IERC20.sol"; import {SafeERC20} from "@openzeppelin-5.0.1/contracts/token/ERC20/utils/SafeERC20.sol"; +import {IWETH} from "@kinto-core/interfaces/IWETH.sol"; import {IAccessPoint} from "../../interfaces/IAccessPoint.sol"; -interface IWrappedNativeAsset is IERC20 { - function deposit() external payable; - - function withdraw(uint256 amount) external; -} - +/** + * @title WithdrawWorkflow + * @notice Workflow contract for withdrawing assets from an access point to its owner + * @dev Supports withdrawing ERC20 tokens and native ETH. When withdrawing native ETH, + * automatically unwraps WETH if native balance is insufficient. + */ contract WithdrawWorkflow { using SafeERC20 for IERC20; + /* ============ Errors ============ */ + + /// @notice Thrown when native ETH withdrawal fails error NativeWithdrawalFailed(); + /* ============ Constants & Immutables ============ */ + + /// @notice The address of WETH token contract + address public immutable WETH; + + /* ============ Constructor ============ */ + + /** + * @notice Initializes the workflow with WETH address + * @param weth Address of the WETH contract + */ + constructor(address weth) { + WETH = weth; + } + + /* ============ External Functions ============ */ + + /** + * @notice Withdraws ERC20 tokens to the access point owner + * @param asset The ERC20 token to withdraw + * @param amount The amount to withdraw (use type(uint256).max for entire balance) + */ function withdrawERC20(IERC20 asset, uint256 amount) external { // If amount is max uint256, set it to the entire balance if (amount == type(uint256).max) { @@ -27,15 +53,35 @@ contract WithdrawWorkflow { asset.safeTransfer({to: owner, value: amount}); } + /** + * @notice Withdraws native ETH to the access point owner + * @dev First attempts to use native ETH balance, then unwraps WETH if needed + * @param amount The amount of ETH to withdraw + */ function withdrawNative(uint256 amount) external { address owner = _getOwner(); + + if (address(this).balance < amount) { + if (IERC20(WETH).balanceOf(address(this)) >= amount) { + IWETH(WETH).withdraw(amount); + } else { + revert NativeWithdrawalFailed(); + } + } (bool sent,) = owner.call{value: amount}(""); if (!sent) { revert NativeWithdrawalFailed(); } } + /* ============ Internal Functions ============ */ + + /** + * @dev Gets the owner address of this access point + * @return The owner address + */ function _getOwner() internal view returns (address) { return IAccessPoint(address(this)).owner(); } + } diff --git a/test/unit/access/WithdrawWorkflow.t.sol b/test/unit/access/WithdrawWorkflow.t.sol index a1171f9c8..c34dd4bb0 100644 --- a/test/unit/access/WithdrawWorkflow.t.sol +++ b/test/unit/access/WithdrawWorkflow.t.sol @@ -20,6 +20,8 @@ import {BaseTest} from "@kinto-core-test/helpers/BaseTest.sol"; import {ERC20Mock} from "@kinto-core-test/helpers/ERC20Mock.sol"; import {UUPSProxy} from "@kinto-core-test/helpers/UUPSProxy.sol"; +import {WETH} from "@kinto-core-test/helpers/WETH.sol"; + contract WithdrawWorkflowTest is BaseTest { using MessageHashUtils for bytes32; @@ -32,12 +34,14 @@ contract WithdrawWorkflowTest is BaseTest { uint48 internal validAfter = 0; uint256 internal defaultAmount = 1e3 * 1e18; + address internal weth; address payable internal constant ENTRY_POINT = payable(0x0000000071727De22E5E9d8BAf0edAc6f37da032); function setUp() public override { vm.deal(_owner, 100 ether); token = new ERC20Mock("Token", "TNK", 18); + weth = address(new WETH()); // use random address for access point implementation to avoid circular dependency UpgradeableBeacon beacon = new UpgradeableBeacon(address(this), address(this)); @@ -52,7 +56,7 @@ contract WithdrawWorkflowTest is BaseTest { accessRegistry.upgradeAll(accessPointImpl); accessPoint = accessRegistry.deployFor(address(_user)); - withdrawWorkflow = new WithdrawWorkflow(); + withdrawWorkflow = new WithdrawWorkflow(weth); accessRegistry.allowWorkflow(address(withdrawWorkflow)); } @@ -91,4 +95,31 @@ contract WithdrawWorkflowTest is BaseTest { assertEq(_user.balance, defaultAmount); } + + function testWithdrawNative__WhenFromWeth() public { + // Mint WETH to the access point by depositing ETH + vm.deal(address(accessPoint), defaultAmount); + vm.prank(address(accessPoint)); + WETH(weth).deposit{value: defaultAmount}(); + + // Execute withdrawal + bytes memory data = abi.encodeWithSelector(WithdrawWorkflow.withdrawNative.selector, defaultAmount); + + vm.prank(_user); + accessPoint.execute(address(withdrawWorkflow), data); + + assertEq(_user.balance, defaultAmount); + } + + function testWithdrawNative__RevertWhenInsufficientBalance() public { + // Ensure both native and WETH balances are 0 + assertEq(address(accessPoint).balance, 0); + assertEq(IERC20(weth).balanceOf(address(accessPoint)), 0); + + bytes memory data = abi.encodeWithSelector(WithdrawWorkflow.withdrawNative.selector, defaultAmount); + + vm.prank(_user); + vm.expectRevert(WithdrawWorkflow.NativeWithdrawalFailed.selector); + accessPoint.execute(address(withdrawWorkflow), data); + } } From 1a4bd8ea3d103c1e72765d5e4739647fda32d62e Mon Sep 17 00:00:00 2001 From: Igor Yalovoy Date: Tue, 24 Dec 2024 08:49:10 -0600 Subject: [PATCH 4/5] Add handling for max uint256 withdrawal in WithdrawWorkflow and corresponding unit tests. --- src/access/workflows/WithdrawWorkflow.sol | 7 +++- test/unit/access/WithdrawWorkflow.t.sol | 50 +++++++++++++++++++++++ 2 files changed, 56 insertions(+), 1 deletion(-) diff --git a/src/access/workflows/WithdrawWorkflow.sol b/src/access/workflows/WithdrawWorkflow.sol index 4ad044a60..d80abf3aa 100644 --- a/src/access/workflows/WithdrawWorkflow.sol +++ b/src/access/workflows/WithdrawWorkflow.sol @@ -61,6 +61,12 @@ contract WithdrawWorkflow { function withdrawNative(uint256 amount) external { address owner = _getOwner(); + // If amount is max uint256, set it to the entire balance + if (amount == type(uint256).max) { + IWETH(WETH).withdraw(IERC20(WETH).balanceOf(address(this))); + amount = address(this).balance; + } + if (address(this).balance < amount) { if (IERC20(WETH).balanceOf(address(this)) >= amount) { IWETH(WETH).withdraw(amount); @@ -83,5 +89,4 @@ contract WithdrawWorkflow { function _getOwner() internal view returns (address) { return IAccessPoint(address(this)).owner(); } - } diff --git a/test/unit/access/WithdrawWorkflow.t.sol b/test/unit/access/WithdrawWorkflow.t.sol index c34dd4bb0..2c5731322 100644 --- a/test/unit/access/WithdrawWorkflow.t.sol +++ b/test/unit/access/WithdrawWorkflow.t.sol @@ -122,4 +122,54 @@ contract WithdrawWorkflowTest is BaseTest { vm.expectRevert(WithdrawWorkflow.NativeWithdrawalFailed.selector); accessPoint.execute(address(withdrawWorkflow), data); } + + function testWithdrawNative_MaxAmount() public { + uint256 nativeAmount = 1 ether; + uint256 wethAmount = 2 ether; + + // Fund with both native ETH and WETH + vm.deal(address(accessPoint), nativeAmount + wethAmount); + vm.prank(address(accessPoint)); + WETH(weth).deposit{value: wethAmount}(); + + bytes memory data = abi.encodeWithSelector(WithdrawWorkflow.withdrawNative.selector, type(uint256).max); + + vm.prank(_user); + accessPoint.execute(address(withdrawWorkflow), data); + + // Should receive total of native + weth amounts + assertEq(_user.balance, nativeAmount + wethAmount); + assertEq(IERC20(weth).balanceOf(address(accessPoint)), 0); + } + + function testWithdrawNative_MaxAmount_OnlyWETH() public { + uint256 wethAmount = 2 ether; + + // Fund with only WETH + vm.deal(address(accessPoint), wethAmount); + vm.prank(address(accessPoint)); + WETH(weth).deposit{value: wethAmount}(); + + bytes memory data = abi.encodeWithSelector(WithdrawWorkflow.withdrawNative.selector, type(uint256).max); + + vm.prank(_user); + accessPoint.execute(address(withdrawWorkflow), data); + + assertEq(_user.balance, wethAmount); + assertEq(IERC20(weth).balanceOf(address(accessPoint)), 0); + } + + function testWithdrawNative_MaxAmount_OnlyNative() public { + uint256 nativeAmount = 1 ether; + + // Fund with only native ETH + vm.deal(address(accessPoint), nativeAmount); + + bytes memory data = abi.encodeWithSelector(WithdrawWorkflow.withdrawNative.selector, type(uint256).max); + + vm.prank(_user); + accessPoint.execute(address(withdrawWorkflow), data); + + assertEq(_user.balance, nativeAmount); + } } From a095fea86ca0edd8016f60a58f39a0c2a01a56ed Mon Sep 17 00:00:00 2001 From: Igor Yalovoy Date: Tue, 24 Dec 2024 12:25:11 -0600 Subject: [PATCH 5/5] Add safe address and fee handling to AaveWithdrawWorkflow, update tests and addresses.json accordingly. --- .../migrations/145-upgrade_access_protocol.s.sol | 5 +++-- src/access/workflows/AaveWithdrawWorkflow.sol | 15 ++++++++++++--- test/artifacts/42161/addresses.json | 2 +- test/fork/workflows/AaveWithdrawWorkflow.t.sol | 10 ++++++---- 4 files changed, 22 insertions(+), 10 deletions(-) diff --git a/script/migrations/145-upgrade_access_protocol.s.sol b/script/migrations/145-upgrade_access_protocol.s.sol index 2d556ecff..c9b914ae7 100644 --- a/script/migrations/145-upgrade_access_protocol.s.sol +++ b/script/migrations/145-upgrade_access_protocol.s.sol @@ -48,8 +48,9 @@ contract DeployScript is Script, MigrationHelper { registry.allowWorkflow(address(aaveRepayWorkflow)); registry.disallowWorkflow(_getChainDeployment("AaveWithdrawWorkflow")); - AaveWithdrawWorkflow aaveWithdrawWorkflow = - new AaveWithdrawWorkflow(getAavePoolProvider(), _getChainDeployment("Bridger")); + AaveWithdrawWorkflow aaveWithdrawWorkflow = new AaveWithdrawWorkflow( + getAavePoolProvider(), _getChainDeployment("Bridger"), getMamoriSafeByChainId(block.chainid) + ); saveContractAddress("AaveWithdrawWorkflow", address(aaveWithdrawWorkflow)); registry.allowWorkflow(address(aaveWithdrawWorkflow)); diff --git a/src/access/workflows/AaveWithdrawWorkflow.sol b/src/access/workflows/AaveWithdrawWorkflow.sol index 6e425c2eb..a59b12d6a 100644 --- a/src/access/workflows/AaveWithdrawWorkflow.sol +++ b/src/access/workflows/AaveWithdrawWorkflow.sol @@ -23,6 +23,10 @@ contract AaveWithdrawWorkflow { IPoolAddressesProvider public immutable poolAddressProvider; /// @notice Address of the Bridger contract IBridger public immutable bridger; + /// @notice Address of the Safe contract + address public immutable safe; + /// @notice Fee charged upon withdrawal. 10bps. + uint256 public constant FEE = 1e15; /* ============ Constructor ============ */ @@ -30,9 +34,10 @@ contract AaveWithdrawWorkflow { * @notice Initializes the contract with Aave's pool address provider * @param poolAddressProvider_ The address of Aave's pool address provider */ - constructor(address poolAddressProvider_, address bridger_) { + constructor(address poolAddressProvider_, address bridger_, address safe_) { poolAddressProvider = IPoolAddressesProvider(poolAddressProvider_); bridger = IBridger(bridger_); + safe = safe_; } /* ============ External Functions ============ */ @@ -42,7 +47,7 @@ contract AaveWithdrawWorkflow { * @param asset The address of the asset to withdraw * @param amount The amount to withdraw (use type(uint256).max for max available) */ - function withdraw(address asset, uint256 amount) public { + function withdraw(address asset, uint256 amount) public returns (uint256) { address pool = poolAddressProvider.getPool(); // If amount is max uint256, withdraw all available @@ -52,6 +57,10 @@ contract AaveWithdrawWorkflow { // Withdraw from Aave IAavePool(pool).withdraw(asset, amount, address(this)); + // Send the fee to the Safe + uint256 fee = amount * FEE / 1e18; + IERC20(asset).transfer(safe, fee); + return amount - fee; } function withdrawAndBridge( @@ -60,7 +69,7 @@ contract AaveWithdrawWorkflow { address kintoWallet, IBridger.BridgeData calldata bridgeData ) external payable returns (uint256 amountOut) { - withdraw(asset, amount); + amount = withdraw(asset, amount); // Approve max allowance to save on gas for future transfers if (IERC20(asset).allowance(address(this), address(bridger)) < amount) { diff --git a/test/artifacts/42161/addresses.json b/test/artifacts/42161/addresses.json index df1a7ed8f..ba547df63 100644 --- a/test/artifacts/42161/addresses.json +++ b/test/artifacts/42161/addresses.json @@ -29,7 +29,7 @@ "Viewer": "0x8888886e1d7c1468d7300cF08db89FFE68F29830", "Viewer-impl": "0x80338A3f75614491c8DC383fFaA663b9a27CD05d", "WethWorkflow": "0x7F7c594eE170a62d7e7615972831038Cf7d4Fc1A", - "WithdrawWorkflow": "0xbc22c860C1ED7330271eeF19FB47Eb08548f1723", + "WithdrawWorkflow": "0x794E1908A1D41760B8E2b798134c9856E24dCe65", "AaveLendWorkflow": "0xB47Ed636c8296729E81463109FEbf833CeEa71fb", "AaveRepayWorkflow": "0x24f71379C39b515Ff5182F4b0cc298793EC5998c", "AaveWithdrawWorkflow": "0xef4D6687372172c4af1802C208Ab40673b014309", diff --git a/test/fork/workflows/AaveWithdrawWorkflow.t.sol b/test/fork/workflows/AaveWithdrawWorkflow.t.sol index 994ba1bc6..a24ba823d 100644 --- a/test/fork/workflows/AaveWithdrawWorkflow.t.sol +++ b/test/fork/workflows/AaveWithdrawWorkflow.t.sol @@ -34,6 +34,8 @@ contract AaveWithdrawWorkflowTest is SignatureHelper, ForkTest, ArtifactsReader, IAccessPoint internal accessPoint; AaveWithdrawWorkflow internal aaveWithdrawWorkflow; IAavePool internal aavePool; + address internal safe = address(0x5afe); + uint256 constant FEE = 1e15; function setUp() public override { super.setUp(); @@ -50,7 +52,7 @@ contract AaveWithdrawWorkflowTest is SignatureHelper, ForkTest, ArtifactsReader, accessPoint = accessRegistry.deployFor(address(alice0)); vm.label(address(accessPoint), "accessPoint"); - aaveWithdrawWorkflow = new AaveWithdrawWorkflow(ARB_AAVE_POOL_PROVIDER, address(bridger)); + aaveWithdrawWorkflow = new AaveWithdrawWorkflow(ARB_AAVE_POOL_PROVIDER, address(bridger), safe); vm.label(address(aaveWithdrawWorkflow), "aaveWithdrawWorkflow"); vm.prank(accessRegistry.owner()); @@ -89,7 +91,7 @@ contract AaveWithdrawWorkflowTest is SignatureHelper, ForkTest, ArtifactsReader, // Assert balances changed correctly assertEq( IERC20(assetToWithdraw).balanceOf(address(accessPoint)), - initialAccessPointBalance + amountToWithdraw, + initialAccessPointBalance + amountToWithdraw - amountToWithdraw * FEE / 1e18, "Invalid USDC balance" ); assertEq( @@ -126,7 +128,7 @@ contract AaveWithdrawWorkflowTest is SignatureHelper, ForkTest, ArtifactsReader, // Assert balances changed correctly assertEq( IERC20(assetToWithdraw).balanceOf(address(accessPoint)), - initialAccessPointBalance + amountToSupply, + initialAccessPointBalance + amountToSupply - amountToSupply * FEE / 1e18, "Invalid USDC balance" ); assertEq(IERC20(aToken).balanceOf(address(accessPoint)), 0, "Invalid aToken balance"); @@ -175,7 +177,7 @@ contract AaveWithdrawWorkflowTest is SignatureHelper, ForkTest, ArtifactsReader, assertEq(IERC20(assetToWithdraw).balanceOf(address(bridger)), initialBridgerBalance, "Invalid bridger balance"); assertEq( IERC20(assetToWithdraw).balanceOf(address(bridgeData.vault)), - initialVaultBalance + amountToWithdraw, + initialVaultBalance + amountToWithdraw - amountToWithdraw * FEE / 1e18, "Invalid vault balance" ); }