Add handling for max uint256 amount in withdrawERC20 and correspondin…

…g test case.
KintoXYZ · Dec 23, 2024 · ec54702 · ec54702
1 parent 98e9a96
commit ec54702
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 0 deletions.
diff --git 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
@@ -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);