From fe1454b06c7c361c36c04761225621f7e0ddf195 Mon Sep 17 00:00:00 2001
From: salimtb <salim.toubal@outlook.com>
Date: Sun, 3 Nov 2024 00:12:46 +0100
Subject: [PATCH] fix: multi-chain-detection

---
 .../src/TokenDetectionController.ts           | 196 +++++++++++-------
 1 file changed, 124 insertions(+), 72 deletions(-)

diff --git a/packages/assets-controllers/src/TokenDetectionController.ts b/packages/assets-controllers/src/TokenDetectionController.ts
index 52a84bfc01..e0af458854 100644
--- a/packages/assets-controllers/src/TokenDetectionController.ts
+++ b/packages/assets-controllers/src/TokenDetectionController.ts
@@ -147,7 +147,7 @@ export type TokenDetectionControllerMessenger = RestrictedControllerMessenger<
 
 /** The input to start polling for the {@link TokenDetectionController} */
 type TokenDetectionPollingInput = {
-  networkClientId: NetworkClientId;
+  chainIds: Hex[];
   address: string;
 };
 
@@ -394,28 +394,6 @@ export class TokenDetectionController extends StaticIntervalPollingController<To
         }
       },
     );
-
-    this.messagingSystem.subscribe(
-      'NetworkController:networkDidChange',
-      // TODO: Either fix this lint violation or explain why it's necessary to ignore.
-      // eslint-disable-next-line @typescript-eslint/no-misused-promises
-      async ({ selectedNetworkClientId }) => {
-        const isNetworkClientIdChanged =
-          this.#networkClientId !== selectedNetworkClientId;
-
-        const { chainId: newChainId } =
-          this.#getCorrectChainIdAndNetworkClientId(selectedNetworkClientId);
-        this.#isDetectionEnabledForNetwork =
-          isTokenDetectionSupportedForNetwork(newChainId);
-
-        if (isNetworkClientIdChanged && this.#isDetectionEnabledForNetwork) {
-          this.#networkClientId = selectedNetworkClientId;
-          await this.#restartTokenDetection({
-            networkClientId: this.#networkClientId,
-          });
-        }
-      },
-    );
   }
 
   /**
@@ -501,6 +479,60 @@ export class TokenDetectionController extends StaticIntervalPollingController<To
     return isEqualValues;
   }
 
+  #getCorrectNetworkClientIdByChainId(
+    chainIds: Hex[] | undefined,
+  ): { chainId: Hex; networkClientId: NetworkClientId }[] | null {
+    // Extract necessary state properties
+    const { networkConfigurationsByChainId, selectedNetworkClientId } =
+      this.messagingSystem.call('NetworkController:getState');
+
+    // only return selected network if chainIds is not defined
+    if (!chainIds) {
+      const networkConfiguration = this.messagingSystem.call(
+        'NetworkController:getNetworkConfigurationByNetworkClientId',
+        selectedNetworkClientId,
+      );
+
+      if (networkConfiguration) {
+        return [
+          {
+            chainId: networkConfiguration.chainId,
+            networkClientId: selectedNetworkClientId,
+          },
+        ];
+      }
+      return null;
+    }
+
+    // Map over each chainId in the input array and build the mapping
+    return chainIds.map((chainId) => {
+      // Get the network configuration for the specified chainId
+      const configuration = networkConfigurationsByChainId[chainId];
+
+      // Check if a configuration exists for the specified chainId
+      if (!configuration) {
+        console.error(`No configuration found for chainId: ${chainId}`);
+        return { chainId, networkClientId: '' }; // TODO: better handle this case
+      }
+
+      // Get the rpcEndpoints array from the configuration
+      const { rpcEndpoints } = configuration;
+
+      // Check if selectedNetworkClientId exists in rpcEndpoints
+      const matchingEndpoint = rpcEndpoints.find(
+        (endpoint) => endpoint.networkClientId === selectedNetworkClientId,
+      );
+
+      // Return an object with both chainId and networkClientId
+      return {
+        chainId,
+        networkClientId: matchingEndpoint
+          ? matchingEndpoint.networkClientId
+          : rpcEndpoints[0].networkClientId,
+      };
+    });
+  }
+
   #getCorrectChainIdAndNetworkClientId(networkClientId?: NetworkClientId): {
     chainId: Hex;
     networkClientId: NetworkClientId;
@@ -533,14 +565,14 @@ export class TokenDetectionController extends StaticIntervalPollingController<To
   }
 
   async _executePoll({
-    networkClientId,
+    chainIds,
     address,
   }: TokenDetectionPollingInput): Promise<void> {
     if (!this.isActive) {
       return;
     }
     await this.detectTokens({
-      networkClientId,
+      chainIds,
       selectedAddress: address,
     });
   }
@@ -551,17 +583,17 @@ export class TokenDetectionController extends StaticIntervalPollingController<To
    *
    * @param options - Options for restart token detection.
    * @param options.selectedAddress - the selectedAddress against which to detect for token balances
-   * @param options.networkClientId - The ID of the network client to use.
+   * @param options.chainIds - The chain IDs of the network client to use.
    */
   async #restartTokenDetection({
     selectedAddress,
-    networkClientId,
+    chainIds,
   }: {
     selectedAddress?: string;
-    networkClientId?: NetworkClientId;
+    chainIds?: Hex[];
   } = {}): Promise<void> {
     await this.detectTokens({
-      networkClientId,
+      chainIds,
       selectedAddress,
     });
     this.setIntervalLength(DEFAULT_INTERVAL);
@@ -572,72 +604,92 @@ export class TokenDetectionController extends StaticIntervalPollingController<To
    * On mainnet, if token detection is disabled in preferences, ERC20 token auto detection will be triggered for each contract address in the legacy token list from the @metamask/contract-metadata repo.
    *
    * @param options - Options for token detection.
-   * @param options.networkClientId - The ID of the network client to use.
+   * @param options.chainIds - The chain IDs of the network client to use.
    * @param options.selectedAddress - the selectedAddress against which to detect for token balances.
    */
   async detectTokens({
-    networkClientId,
+    chainIds, // todo: replace with chainIds
     selectedAddress,
   }: {
-    networkClientId?: NetworkClientId;
+    chainIds?: Hex[];
     selectedAddress?: string;
   } = {}): Promise<void> {
     if (!this.isActive) {
       return;
     }
 
+    // Address used by the user
     const addressAgainstWhichToDetect =
       selectedAddress ?? this.#getSelectedAddress();
-    const { chainId, networkClientId: selectedNetworkClientId } =
-      this.#getCorrectChainIdAndNetworkClientId(networkClientId);
-    const chainIdAgainstWhichToDetect = chainId;
-    const networkClientIdAgainstWhichToDetect = selectedNetworkClientId;
 
-    if (!isTokenDetectionSupportedForNetwork(chainIdAgainstWhichToDetect)) {
-      return;
-    }
-    if (
-      !this.#isDetectionEnabledFromPreferences &&
-      chainIdAgainstWhichToDetect !== ChainId.mainnet
-    ) {
+    // Will be an array of objects in the form [{chainId : "0x1", networkClientId: "8e8b86f2-f949-4f29-90bc-f7409f0832d5"}]
+    const clientNetworksIdByChainId =
+      this.#getCorrectNetworkClientIdByChainId(chainIds);
+
+    if (!clientNetworksIdByChainId) {
       return;
     }
-    const isTokenDetectionInactiveInMainnet =
-      !this.#isDetectionEnabledFromPreferences &&
-      chainIdAgainstWhichToDetect === ChainId.mainnet;
-    const { tokensChainsCache } = this.messagingSystem.call(
-      'TokenListController:getState',
-    );
-    this.#tokensChainsCache = isTokenDetectionInactiveInMainnet
-      ? this.#getConvertedStaticMainnetTokenList()
-      : tokensChainsCache ?? {};
 
-    const tokenCandidateSlices = this.#getSlicesOfTokensToDetect({
+    // Execute the rest in a loop for each pair of chainId and networkClientId in the result array
+    for (const {
       chainId: chainIdAgainstWhichToDetect,
-      selectedAddress: addressAgainstWhichToDetect,
-    });
+      networkClientId: networkClientIdAgainstWhichToDetect,
+    } of clientNetworksIdByChainId) {
+      // Check if token detection is supported for this chainId
+      if (!isTokenDetectionSupportedForNetwork(chainIdAgainstWhichToDetect)) {
+        continue;
+      }
 
-    // Attempt Accounts API Detection
-    const accountAPIResult = await this.#addDetectedTokensViaAPI({
-      chainId: chainIdAgainstWhichToDetect,
-      selectedAddress: addressAgainstWhichToDetect,
-      tokenCandidateSlices,
-    });
-    if (accountAPIResult?.result === 'success') {
-      return;
-    }
+      // Check if detection is enabled based on preferences and chain ID
+      if (
+        !this.#isDetectionEnabledFromPreferences &&
+        chainIdAgainstWhichToDetect !== ChainId.mainnet
+      ) {
+        continue;
+      }
 
-    // Attempt RPC Detection
-    const tokenDetectionPromises = tokenCandidateSlices.map((tokensSlice) =>
-      this.#addDetectedTokens({
-        tokensSlice,
+      const isTokenDetectionInactiveInMainnet =
+        !this.#isDetectionEnabledFromPreferences &&
+        chainIdAgainstWhichToDetect === ChainId.mainnet;
+
+      // Get tokens chains cache from TokenListController or initialize mainnet tokens if inactive
+      const { tokensChainsCache } = this.messagingSystem.call(
+        'TokenListController:getState',
+      );
+      this.#tokensChainsCache = isTokenDetectionInactiveInMainnet
+        ? this.#getConvertedStaticMainnetTokenList()
+        : tokensChainsCache ?? {};
+
+      // Get slices of tokens to detect for the given chainId and selected address
+      const tokenCandidateSlices = this.#getSlicesOfTokensToDetect({
+        chainId: chainIdAgainstWhichToDetect,
         selectedAddress: addressAgainstWhichToDetect,
-        networkClientId: networkClientIdAgainstWhichToDetect,
+      });
+
+      // Attempt Accounts API Detection
+      const accountAPIResult = await this.#addDetectedTokensViaAPI({
         chainId: chainIdAgainstWhichToDetect,
-      }),
-    );
+        selectedAddress: addressAgainstWhichToDetect,
+        tokenCandidateSlices,
+      });
+
+      if (accountAPIResult?.result === 'success') {
+        continue; // Skip to the next iteration if detection via API was successful
+      }
+
+      // Attempt RPC Detection if API detection wasn't successful
+      const tokenDetectionPromises = tokenCandidateSlices.map((tokensSlice) =>
+        this.#addDetectedTokens({
+          tokensSlice,
+          selectedAddress: addressAgainstWhichToDetect,
+          networkClientId: networkClientIdAgainstWhichToDetect,
+          chainId: chainIdAgainstWhichToDetect,
+        }),
+      );
 
-    await Promise.all(tokenDetectionPromises);
+      // Await all token detection promises for this chainId
+      await Promise.all(tokenDetectionPromises);
+    }
   }
 
   #getSlicesOfTokensToDetect({