Add handling for max uint256 amount in withdrawERC20

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));
 

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(

src/access/workflows/AaveWithdrawWorkflow.sol

@@ -23,16 +23,21 @@ 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 ============ */
 
     /**
      * @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) {

src/access/workflows/WithdrawWorkflow.sol

@@ -3,33 +3,89 @@ 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) {
+            amount = asset.balanceOf(address(this));
+        }
+
         address owner = _getOwner();
         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 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);
+            } 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();
     }

test/artifacts/42161/addresses.json

@@ -29,12 +29,12 @@
   "Viewer": "0x8888886e1d7c1468d7300cF08db89FFE68F29830",
   "Viewer-impl": "0x80338A3f75614491c8DC383fFaA663b9a27CD05d",
   "WethWorkflow": "0x7F7c594eE170a62d7e7615972831038Cf7d4Fc1A",
-  "WithdrawWorkflow": "0x36e6cA034958B2E0D4dC7Ea9a8151f15Ba0B27D2",
+  "WithdrawWorkflow": "0x794E1908A1D41760B8E2b798134c9856E24dCe65",
   "AaveLendWorkflow": "0xB47Ed636c8296729E81463109FEbf833CeEa71fb",
   "AaveRepayWorkflow": "0x24f71379C39b515Ff5182F4b0cc298793EC5998c",
   "AaveWithdrawWorkflow": "0xef4D6687372172c4af1802C208Ab40673b014309",
   "AaveBorrowWorkflow": "0xD0187Ca378f7B26D53e0A91fBe8D7ba80498EF10",
   "AccessManager": "0xacc0003a4aAE5dA4ba12F771C7350D40147Cd7D4",
   "BridgerV13-impl": "0xbA6FD752CE93879c381fb7ffdbe7baB233D6e6e4",
   "BridgerV14-impl": "0x363EFf1981E664107EF4E8568Cc4321B74558DAA"
-}
+}

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"
         );
     }

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));
     }
@@ -69,6 +73,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);
 
@@ -79,4 +95,81 @@ 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);
+    }
+
+    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);
+    }
 }