From ec54702df4d899c05a4a367ce64b7ad823a4c101 Mon Sep 17 00:00:00 2001
From: Igor Yalovoy <to@ylv.io>
Date: Mon, 23 Dec 2024 15:23:00 -0600
Subject: [PATCH] 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);