From ffe2bd315e79edfc7c47aecd7f2bf00b21bc9d0d Mon Sep 17 00:00:00 2001
From: Zach Pomerantz <zzmp@uniswap.org>
Date: Mon, 24 Jan 2022 17:52:45 -0800
Subject: [PATCH] fix: track swap approvals (#3183)

* fix: track swap approvals

* fix: type ambiguous return value
---
 src/hooks/useApproveCallback.ts | 35 ++++++++++++++++-----------------
 src/lib/hooks/useApproval.ts    | 12 +++++++++--
 2 files changed, 27 insertions(+), 20 deletions(-)

diff --git a/src/hooks/useApproveCallback.ts b/src/hooks/useApproveCallback.ts
index 8aa148f44b..4c8b5a9654 100644
--- a/src/hooks/useApproveCallback.ts
+++ b/src/hooks/useApproveCallback.ts
@@ -1,4 +1,3 @@
-import { TransactionResponse } from '@ethersproject/providers'
 import { Trade } from '@uniswap/router-sdk'
 import { Currency, CurrencyAmount, Percent, TradeType } from '@uniswap/sdk-core'
 import { Trade as V2Trade } from '@uniswap/v2-sdk'
@@ -6,31 +5,30 @@ import { Trade as V3Trade } from '@uniswap/v3-sdk'
 import useSwapApproval, { useSwapApprovalOptimizedTrade } from 'lib/hooks/swap/useSwapApproval'
 import { ApprovalState, useApproval } from 'lib/hooks/useApproval'
 import { useCallback } from 'react'
-import invariant from 'tiny-invariant'
 
 import { TransactionType } from '../state/transactions/actions'
 import { useHasPendingApproval, useTransactionAdder } from '../state/transactions/hooks'
 export { ApprovalState } from 'lib/hooks/useApproval'
 
+function useGetAndTrackApproval(getApproval: ReturnType<typeof useApproval>[1]) {
+  const addTransaction = useTransactionAdder()
+  return useCallback(() => {
+    return getApproval().then((pending) => {
+      if (pending) {
+        const { response, tokenAddress, spenderAddress: spender } = pending
+        addTransaction(response, { type: TransactionType.APPROVAL, tokenAddress, spender })
+      }
+    })
+  }, [addTransaction, getApproval])
+}
+
 // returns a variable indicating the state of the approval and a function which approves if necessary or early returns
 export function useApproveCallback(
   amountToApprove?: CurrencyAmount<Currency>,
   spender?: string
 ): [ApprovalState, () => Promise<void>] {
-  const token = amountToApprove?.currency?.isToken ? amountToApprove.currency : undefined
-  const addTransaction = useTransactionAdder()
-  const [approval, approvalCallback] = useApproval(amountToApprove, spender, useHasPendingApproval)
-
-  const approveCallback = useCallback(() => {
-    return approvalCallback().then((response?: TransactionResponse) => {
-      if (response) {
-        invariant(token && spender)
-        addTransaction(response, { type: TransactionType.APPROVAL, tokenAddress: token.address, spender })
-      }
-    })
-  }, [approvalCallback, token, spender, addTransaction])
-
-  return [approval, approveCallback]
+  const [approval, getApproval] = useApproval(amountToApprove, spender, useHasPendingApproval)
+  return [approval, useGetAndTrackApproval(getApproval)]
 }
 
 export function useApprovalOptimizedTrade(
@@ -47,6 +45,7 @@ export function useApproveCallbackFromTrade(
     | Trade<Currency, Currency, TradeType>
     | undefined,
   allowedSlippage: Percent
-) {
-  return useSwapApproval(trade, allowedSlippage, useHasPendingApproval)
+): [ApprovalState, () => Promise<void>] {
+  const [approval, getApproval] = useSwapApproval(trade, allowedSlippage, useHasPendingApproval)
+  return [approval, useGetAndTrackApproval(getApproval)]
 }
diff --git a/src/lib/hooks/useApproval.ts b/src/lib/hooks/useApproval.ts
index 97b369a7e0..30e3394a62 100644
--- a/src/lib/hooks/useApproval.ts
+++ b/src/lib/hooks/useApproval.ts
@@ -44,7 +44,10 @@ export function useApproval(
   amountToApprove: CurrencyAmount<Currency> | undefined,
   spender: string | undefined,
   useIsPendingApproval: (token?: Token, spender?: string) => boolean
-): [ApprovalState, () => Promise<TransactionResponse | undefined>] {
+): [
+  ApprovalState,
+  () => Promise<{ response: TransactionResponse; tokenAddress: string; spenderAddress: string } | undefined>
+] {
   const { chainId } = useActiveWeb3React()
   const token = amountToApprove?.currency?.isToken ? amountToApprove.currency : undefined
 
@@ -53,7 +56,7 @@ export function useApproval(
 
   const tokenContract = useTokenContract(token?.address)
 
-  const approve = useCallback(async (): Promise<TransactionResponse | undefined> => {
+  const approve = useCallback(async () => {
     function logFailure(error: Error | string): undefined {
       console.warn(`${token?.symbol || 'Token'} approval failed:`, error)
       return
@@ -85,6 +88,11 @@ export function useApproval(
       .approve(spender, useExact ? amountToApprove.quotient.toString() : MaxUint256, {
         gasLimit: calculateGasMargin(estimatedGas),
       })
+      .then((response) => ({
+        response,
+        tokenAddress: token.address,
+        spenderAddress: spender,
+      }))
       .catch((error: Error) => {
         logFailure(error)
         throw error