From 10002da084c8e3c27f09506d598bae1c8b3ead6c Mon Sep 17 00:00:00 2001
From: Kevin Ingersoll <kingersoll@gmail.com>
Date: Wed, 22 May 2024 18:41:07 +0100
Subject: [PATCH] refactor with proper fee calculation

---
 .../src/steps/deposit/WithdrawButton.tsx      | 103 +++---------------
 .../src/steps/deposit/usePrepareWithdraw.ts   |  88 +++++++++++++++
 2 files changed, 103 insertions(+), 88 deletions(-)
 create mode 100644 packages/account-kit/src/steps/deposit/usePrepareWithdraw.ts

diff --git a/packages/account-kit/src/steps/deposit/WithdrawButton.tsx b/packages/account-kit/src/steps/deposit/WithdrawButton.tsx
index cef7a087a0..88400cdfbd 100644
--- a/packages/account-kit/src/steps/deposit/WithdrawButton.tsx
+++ b/packages/account-kit/src/steps/deposit/WithdrawButton.tsx
@@ -1,96 +1,32 @@
-import { useMutation, useQuery } from "@tanstack/react-query";
+import { useMutation } from "@tanstack/react-query";
 import { useAppAccountClient } from "../../useAppAccountClient";
-import { prepareTransactionRequest, readContract, sendTransaction, waitForTransactionReceipt } from "viem/actions";
-import { getAction, getChainContractAddress, serializeTransaction } from "viem/utils";
+import { sendTransaction, waitForTransactionReceipt } from "viem/actions";
+import { getAction } from "viem/utils";
 import { Button } from "../../ui/Button";
 import { AppAccountClient } from "../../common";
-import { useAccount, useBalance, useEstimateFeesPerGas, useEstimateGas } from "wagmi";
-import { EstimateFeesPerGasReturnType, Hex, maxUint256 } from "viem";
-import { gasPriceOracleAbi } from "viem/op-stack/abis";
+import { SendTransactionParameters } from "viem";
 import { useInvalidateBalance } from "./useInvalidateBalance";
-import { estimateL1Fee } from "viem/op-stack";
-import { useConfig } from "../../AccountKitConfigProvider";
-
-// TODO: switch network
+import { usePrepareWithdraw } from "./usePrepareWithdraw";
+import { useAccount } from "wagmi";
 
 export function WithdrawButton() {
   const invalidateBalance = useInvalidateBalance();
-  const { chainId } = useConfig();
-  const { address: userAddress } = useAccount();
   const { data: appAccountClient } = useAppAccountClient();
-  const { data: balance } = useBalance({
-    chainId: appAccountClient?.chain.id,
-    address: appAccountClient?.account.address,
-  });
-
-  const estimateL2Fee = useEstimateFeesPerGas({ chainId });
-  // sending value is a constant 21k gas
-  const withdrawL2Gas = 21000n;
-
-  const queryKey = [
-    "withdrawL1Fee",
-    appAccountClient?.chain.id,
-    appAccountClient?.account.address,
-    userAddress,
-    balance!.value.toString(),
-  ];
-  const withdrawL1Fee = useQuery(
-    appAccountClient && userAddress && balance
-      ? {
-          queryKey,
-          queryFn: async () => {
-            return estimateL1Fee(appAccountClient, {
-              account: appAccountClient.account,
-              to: userAddress,
-              value: balance.value,
-              // Skip gas estimation of `prepareTransactionRequest` inside `estimateL1Fee`,
-              // otherwise this will fail due to `balance` not being enough to cover `value` + `fee`.
-              gas: 0n,
-            });
-          },
-          // TODO: figure out a good refresh rate (L1 block time? L2 commit time?)
-          refetchInterval: 4_000,
-        }
-      : { queryKey, enabled: false },
-  );
-
-  console.log("balance", balance?.value);
-  console.log("withdrawL1Fee", withdrawL1Fee.status, withdrawL1Fee.data, withdrawL1Fee.error);
-
-  const withdrawL2Fee = estimateL2Fee.isSuccess ? withdrawL2Gas * estimateL2Fee.data.maxFeePerGas : undefined;
-  const withdrawFee = withdrawL1Fee.isSuccess && withdrawL2Fee != null ? withdrawL1Fee.data + withdrawL2Fee : undefined;
-  const withdrawAmount = balance && withdrawFee != null ? balance.value - withdrawFee : undefined;
+  const { address: userAddress } = useAccount();
 
-  // TODO: lift this up so we can conditionally hide the button
+  const prepared = usePrepareWithdraw({ appAccountClient, userAddress });
 
   const withdraw = useMutation({
-    mutationKey: ["withdraw", appAccountClient?.account.address],
+    mutationKey: ["withdraw"],
     mutationFn: async ({
       appAccountClient,
-      userAddress,
-      amount,
-      gas,
-      fees,
+      params,
     }: {
       appAccountClient: AppAccountClient;
-      userAddress: Hex;
-      amount: bigint;
-      gas: bigint;
-      fees: EstimateFeesPerGasReturnType;
+      params: SendTransactionParameters;
     }) => {
       console.log("withdrawing funds");
-      const hash = await getAction(
-        appAccountClient,
-        sendTransaction,
-        "sendTransaction",
-      )({
-        chain: appAccountClient.chain,
-        account: appAccountClient.account,
-        to: userAddress,
-        value: amount,
-        gas,
-        ...fees,
-      });
+      const hash = await getAction(appAccountClient, sendTransaction, "sendTransaction")(params);
 
       const receipt = await getAction(
         appAccountClient,
@@ -108,21 +44,12 @@ export function WithdrawButton() {
     <Button
       variant="secondary"
       className="p-2 text-sm"
-      pending={
-        userAddress == null ||
-        appAccountClient == null ||
-        withdrawAmount == null ||
-        withdraw.isPending ||
-        estimateL2Fee.isPending
-      }
-      disabled={withdrawAmount ? withdrawAmount <= 0n : false}
+      pending={appAccountClient == null || prepared.isPending || withdraw.isPending}
+      disabled={prepared.isError}
       onClick={() =>
         withdraw.mutateAsync({
-          userAddress: userAddress!,
           appAccountClient: appAccountClient!,
-          amount: withdrawAmount!,
-          gas: withdrawL2Gas,
-          fees: estimateL2Fee.data!,
+          params: prepared.data!,
         })
       }
     >
diff --git a/packages/account-kit/src/steps/deposit/usePrepareWithdraw.ts b/packages/account-kit/src/steps/deposit/usePrepareWithdraw.ts
new file mode 100644
index 0000000000..853c98bbfb
--- /dev/null
+++ b/packages/account-kit/src/steps/deposit/usePrepareWithdraw.ts
@@ -0,0 +1,88 @@
+import { useQuery } from "@tanstack/react-query";
+import {
+  SendTransactionParameters,
+  estimateFeesPerGas,
+  getBalance,
+  prepareTransactionRequest,
+  readContract,
+} from "viem/actions";
+import { getChainContractAddress, serializeTransaction } from "viem/utils";
+import { AppAccountClient } from "../../common";
+import { Hex } from "viem";
+
+const withdrawGas = 21000n;
+
+export type UsePrepareWithdrawOptions = {
+  appAccountClient: AppAccountClient | undefined;
+  userAddress: Hex | undefined;
+};
+
+export function usePrepareWithdraw({ appAccountClient, userAddress }: UsePrepareWithdrawOptions) {
+  const queryKey = ["prepareWithdraw", appAccountClient?.chain.id, appAccountClient?.account.address, userAddress];
+  return useQuery(
+    appAccountClient && userAddress
+      ? {
+          queryKey,
+          queryFn: async (): Promise<SendTransactionParameters> => {
+            const [balance, fees] = await Promise.all([
+              getBalance(appAccountClient, { address: appAccountClient.account.address }),
+              estimateFeesPerGas(appAccountClient),
+            ]);
+            let fee = withdrawGas * fees.maxFeePerGas;
+
+            // If this is an L2, add L1 fee
+            const gasPriceOracleAddress = getChainContractAddress({
+              chain: appAccountClient.chain,
+              contract: "gasPriceOracle",
+            });
+            if (gasPriceOracleAddress) {
+              const request = await prepareTransactionRequest(appAccountClient, {
+                chain: appAccountClient.chain,
+                account: appAccountClient.account,
+                to: userAddress,
+                value: balance,
+                gas: withdrawGas,
+                ...fees,
+              });
+
+              const transaction = serializeTransaction({
+                ...request,
+                type: "eip1559",
+              });
+
+              fee += await readContract(appAccountClient, {
+                address: gasPriceOracleAddress,
+                abi: [
+                  {
+                    inputs: [{ internalType: "bytes", name: "_data", type: "bytes" }],
+                    name: "getL1Fee",
+                    outputs: [{ internalType: "uint256", name: "", type: "uint256" }],
+                    stateMutability: "view",
+                    type: "function",
+                  },
+                ],
+                functionName: "getL1Fee",
+                args: [transaction],
+              });
+            }
+
+            const withdrawAmount = balance - fee;
+            if (withdrawAmount <= 0n) {
+              throw new Error(`Account balance (${balance}) not enough to cover estimated transfer fee (${fee}).`);
+            }
+
+            return {
+              chain: appAccountClient.chain,
+              account: appAccountClient.account,
+              to: userAddress,
+              value: withdrawAmount,
+              gas: withdrawGas,
+              ...fees,
+            };
+          },
+          // TODO: figure out a good refresh rate (L1 block time? L2 commit time?)
+          refetchInterval: 4_000,
+        }
+      : { queryKey, enabled: false },
+  );
+}