fix: simulation should handle NFT mints (#4217)

MetaMask · Apr 29, 2024 · 342023d · 342023d
1 parent 9349754
commit 342023d
Show file tree

Hide file tree

Showing 2 changed files with 111 additions and 23 deletions.
diff --git a/packages/transaction-controller/src/utils/simulation.test.ts b/packages/transaction-controller/src/utils/simulation.test.ts
@@ -559,6 +559,61 @@ describe('Simulation Utils', () => {
         });
       });
 
+      it('on NFT mint', async () => {
+        mockParseLog({
+          erc721: {
+            ...PARSED_ERC721_TRANSFER_EVENT_MOCK,
+            args: [
+              '0x0000000000000000000000000000000000000000',
+              USER_ADDRESS_MOCK,
+              TOKEN_ID_MOCK,
+            ],
+          },
+        });
+
+        simulateTransactionsMock
+          .mockResolvedValueOnce(
+            createEventResponseMock([createLogMock(CONTRACT_ADDRESS_MOCK)]),
+          )
+          .mockResolvedValueOnce(
+            createBalanceOfResponse([], [USER_ADDRESS_MOCK]),
+          );
+
+        const simulationData = await getSimulationData(REQUEST_MOCK);
+
+        expect(simulateTransactionsMock).toHaveBeenCalledTimes(2);
+        // The second call should only simulate the minting of the NFT and
+        // check the balance after, and not before.
+        expect(simulateTransactionsMock).toHaveBeenNthCalledWith(
+          2,
+          REQUEST_MOCK.chainId,
+          {
+            transactions: [
+              REQUEST_MOCK,
+              {
+                from: REQUEST_MOCK.from,
+                to: CONTRACT_ADDRESS_MOCK,
+                data: expect.any(String),
+              },
+            ],
+          },
+        );
+        expect(simulationData).toStrictEqual({
+          nativeBalanceChange: undefined,
+          tokenBalanceChanges: [
+            {
+              standard: SimulationTokenStandard.erc721,
+              address: CONTRACT_ADDRESS_MOCK,
+              id: TOKEN_ID_MOCK,
+              previousBalance: '0x0',
+              newBalance: '0x1',
+              difference: '0x1',
+              isDecrease: false,
+            },
+          ],
+        });
+      });
+
       it('as empty if events cannot be parsed', async () => {
         mockParseLog({});
 

diff --git a/packages/transaction-controller/src/utils/simulation.ts b/packages/transaction-controller/src/utils/simulation.ts
@@ -91,6 +91,8 @@ const SUPPORTED_TOKEN_ABIS = {
 
 const REVERTED_ERRORS = ['execution reverted', 'insufficient funds for gas'];
 
+type BalanceTransactionMap = Map<SimulationToken, SimulationRequestTransaction>;
+
 /**
  * Generate simulation data for a transaction.
  * @param request - The transaction to simulate.
@@ -199,7 +201,7 @@ function getNativeBalanceChange(
  * @param response - The simulation response.
  * @returns The parsed events.
  */
-function getEvents(response: SimulationResponse): ParsedEvent[] {
+export function getEvents(response: SimulationResponse): ParsedEvent[] {
   /* istanbul ignore next */
   const logs = extractLogs(
     response.transactions[0]?.callTrace ?? ({} as SimulationResponseCallTrace),
@@ -284,41 +286,45 @@ async function getTokenBalanceChanges(
   request: GetSimulationDataRequest,
   events: ParsedEvent[],
 ): Promise<SimulationTokenBalanceChange[]> {
-  const balanceTransactionsByToken = getTokenBalanceTransactions(
-    request,
-    events,
-  );
+  const balanceTxs = getTokenBalanceTransactions(request, events);
 
-  const balanceTransactions = [...balanceTransactionsByToken.values()];
+  log('Generated balance transactions', [...balanceTxs.after.values()]);
 
-  log('Generated balance transactions', balanceTransactions);
+  const transactions = [
+    ...balanceTxs.before.values(),
+    request,
+    ...balanceTxs.after.values(),
+  ];
 
-  if (!balanceTransactions.length) {
+  if (transactions.length === 1) {
     return [];
   }
 
   const response = await simulateTransactions(request.chainId as Hex, {
-    transactions: [...balanceTransactions, request, ...balanceTransactions],
+    transactions,
   });
 
   log('Balance simulation response', response);
 
-  if (response.transactions.length !== balanceTransactions.length * 2 + 1) {
+  if (response.transactions.length !== transactions.length) {
     throw new SimulationInvalidResponseError();
   }
 
-  return [...balanceTransactionsByToken.keys()]
+  return [...balanceTxs.after.keys()]
     .map((token, index) => {
-      const previousBalance = getValueFromBalanceTransaction(
-        request.from,
-        token,
-        response.transactions[index],
-      );
+      const previousBalanceCheckSkipped = !balanceTxs.before.get(token);
+      const previousBalance = previousBalanceCheckSkipped
+        ? '0x0'
+        : getValueFromBalanceTransaction(
+            request.from,
+            token,
+            response.transactions[index],
+          );
 
       const newBalance = getValueFromBalanceTransaction(
         request.from,
         token,
-        response.transactions[index + balanceTransactions.length + 1],
+        response.transactions[index + balanceTxs.before.size + 1],
       );
 
       const balanceChange = getSimulationBalanceChange(
@@ -347,8 +353,13 @@ async function getTokenBalanceChanges(
 function getTokenBalanceTransactions(
   request: GetSimulationDataRequest,
   events: ParsedEvent[],
-): Map<SimulationToken, SimulationRequestTransaction> {
+): {
+  before: BalanceTransactionMap;
+  after: BalanceTransactionMap;
+} {
   const tokenKeys = new Set();
+  const before = new Map();
+  const after = new Map();
 
   const userEvents = events.filter(
     (event) =>
@@ -358,7 +369,7 @@ function getTokenBalanceTransactions(
 
   log('Filtered user events', userEvents);
 
-  return userEvents.reduce((result, event) => {
+  for (const event of userEvents) {
     const tokenIds = getEventTokenIds(event);
 
     log('Extracted token ids', tokenIds);
@@ -388,15 +399,37 @@ function getTokenBalanceTransactions(
         tokenId,
       );
 
-      result.set(simulationToken, {
+      const transaction: SimulationRequestTransaction = {
         from: request.from,
         to: event.contractAddress,
         data,
-      });
+      };
+
+      if (skipPriorBalanceCheck(event)) {
+        after.set(simulationToken, transaction);
+      } else {
+        before.set(simulationToken, transaction);
+        after.set(simulationToken, transaction);
+      }
     }
+  }
 
-    return result;
-  }, new Map<SimulationToken, SimulationRequestTransaction>());
+  return { before, after };
+}
+
+/**
+ * Check if an event needs to check the previous balance.
+ * @param event - The parsed event.
+ * @returns True if the prior balance check should be skipped.
+ */
+function skipPriorBalanceCheck(event: ParsedEvent): boolean {
+  // In the case of an NFT mint, we cannot check the NFT owner before the mint
+  // as the balance check transaction would revert.
+  return (
+    event.name === 'Transfer' &&
+    event.tokenStandard === SimulationTokenStandard.erc721 &&
+    parseInt(event.args.from as string, 16) === 0
+  );
 }
 
 /**