👷🏻 ✅ Support pulling funds from wallet

Kwenta · Jul 10, 2024 · 00369f7 · 00369f7
1 parent 51c9fe0
commit 00369f7
Show file tree

Hide file tree

Showing 4 changed files with 81 additions and 11 deletions.
diff --git a/src/Account.sol b/src/Account.sol
@@ -35,6 +35,34 @@ contract Account is IAccount, IERC721Receiver {
         usdc = IERC20(_usdc);
     }
 
+    function execute(
+        address dest,
+        uint256 value,
+        bytes calldata func
+    ) external {
+        _call(dest, value, func);
+    }
+
+    function _call(address target, uint256 value, bytes memory data) internal {
+        assembly ("memory-safe") {
+            let succ := call(
+                gas(),
+                target,
+                value,
+                add(data, 0x20),
+                mload(data),
+                0x00,
+                0
+            )
+
+            if iszero(succ) {
+                let fmp := mload(0x40)
+                returndatacopy(fmp, 0x00, returndatasize())
+                revert(fmp, returndatasize())
+            }
+        }
+    }
+
     function setupAccount() external {
         accountId = perpsMarketSNXV3.createAccount();
         perpsMarketSNXV3.grantPermission({

diff --git a/src/MarginPaymaster.sol b/src/MarginPaymaster.sol
@@ -10,6 +10,7 @@ import {Account} from "src/Account.sol";
 import {Zap} from "lib/zap/src/Zap.sol";
 import {OracleLibrary} from "src/libraries/OracleLibrary.sol";
 import {IUniswapV3Pool} from "src/interfaces/external/IUniswapV3Pool.sol";
+import {IERC20} from "lib/openzeppelin-contracts/contracts/token/ERC20/IERC20.sol";
 
 import {console} from "forge-std/console.sol";
 
@@ -75,21 +76,35 @@ contract MarginPaymaster is IPaymaster, Zap {
 
         uint256 costOfGasInUSDC = (OracleLibrary.getQuoteAtTick(
             tick,
-            uint128(actualGasCostInWei),
+            uint128(actualGasCostInWei), // TODO: account for gas costs of postOp func
             address(weth),
             address(_USDC)
-        ) * 110) / 100; // TODO: account for gas costs of postOp func
-        uint256 costOfGasInsUSD = costOfGasInUSDC * 1e12;
+        ) * 110) / 100; // 10% slippage
 
         address sender = abi.decode(context, (address));
-        uint128 accountId = Account(sender).accountId();
-        // TODO: also support pulling USDC from the wallet directly
-        perpsMarketSNXV3.modifyCollateral(
-            accountId,
-            sUSDId,
-            -int256(costOfGasInsUSD)
+
+        uint256 walletBalance = _USDC.balanceOf(sender);
+        uint256 allowance = IERC20(address(_USDC)).allowance(
+            sender,
+            address(this)
         );
-        uint256 usdcWithdrawn = _zapOut(costOfGasInsUSD);
+
+        if (walletBalance >= costOfGasInUSDC && allowance >= costOfGasInUSDC) {
+            // pull funds from wallet
+            _USDC.transferFrom(sender, address(this), costOfGasInUSDC);
+        } else {
+            // pull funds from margin
+            uint256 costOfGasInsUSD = costOfGasInUSDC * 1e12;
+            uint128 accountId = Account(sender).accountId();
+            // TODO: also support pulling USDC from the wallet directly
+            perpsMarketSNXV3.modifyCollateral(
+                accountId,
+                sUSDId,
+                -int256(costOfGasInsUSD)
+            );
+            costOfGasInUSDC = _zapOut(costOfGasInsUSD);
+        }
+
         console.log("actualGasCostInWei", actualGasCostInWei); // 43350920000000 = 0.00004335092 ETH = 0.13 USD
 
         IV3SwapRouter.ExactInputSingleParams memory params = IV3SwapRouter
@@ -98,7 +113,7 @@ contract MarginPaymaster is IPaymaster, Zap {
                 tokenOut: address(weth),
                 fee: 500, // 0.05%, top uni pool for USDC/WETH liquidity based on https://www.geckoterminal.com/base/uniswap-v3-base/pools
                 recipient: address(this),
-                amountIn: usdcWithdrawn,
+                amountIn: costOfGasInUSDC,
                 amountOutMinimum: actualGasCostInWei, // TODO: should this be required? -> could cause failures
                 sqrtPriceLimitX96: 0
             });

diff --git a/test/MarginPaymaster.t.sol b/test/MarginPaymaster.t.sol
@@ -123,6 +123,31 @@ contract MarginPaymasterTest is Bootstrap {
         assertLt(colAmount, 5 ether);
     }
 
+    function testTransferToWalletAndApprove() public {
+        bytes memory approvalCalldata = abi.encodeWithSelector(
+            usdc.approve.selector,
+            marginPaymasterAddress,
+            type(uint256).max
+        );
+        userOp.callData = abi.encodeWithSelector(
+            Account.execute.selector,
+            address(usdc),
+            0,
+            approvalCalldata
+        );
+        ops.push(userOp);
+
+        mintUSDC(sender, 1000 * 1e6);
+
+        vm.prank(bundler);
+        entryPoint.handleOps(ops, bundler);
+
+        assertGt(usdc.allowance(sender, marginPaymasterAddress), 0);
+        uint256 usdcLeftInWallet = usdc.balanceOf(sender);
+        assertLt(usdcLeftInWallet, 1000 * 1e6);
+        assertGt(usdcLeftInWallet, 0);
+    }
+
     function testOnlyEntryPointCanCallValidatePaymasterUserOp() public {
         // Create a dummy UserOperation
         UserOperation memory op = getDummyUserOp();

diff --git a/test/utils/interfaces/IUSDC.sol b/test/utils/interfaces/IUSDC.sol
@@ -44,4 +44,6 @@ interface IUSDC {
         address to,
         uint256 amount
     ) external returns (bool);
+
+    function allowance(address owner, address spender) external view returns (uint256);
 }