From c84b85f9532e04e69abb8286fffe55f56adec036 Mon Sep 17 00:00:00 2001
From: Elliot Winkler <elliot.winkler@gmail.com>
Date: Thu, 8 Jun 2023 18:00:33 -0600
Subject: [PATCH] Add getEthQuery action to NetworkController (#19420)

The core version of NetworkControler features a controller action,
`NetworkController:getEthQuery`, which, as its name implies, can be used
to get the EthQuery instance that has been created specially to talk to
the currently selected network. As a result there is a corresponding
unit test for this action.

To make the test suites between this version of NetworkController and
core easier to compare, this commit adds the action along with its test.
---
 .../network/network-controller.test.ts        | 47 ++++++++++++++++-
 .../controllers/network/network-controller.ts | 51 +++++++++++++++++--
 2 files changed, 94 insertions(+), 4 deletions(-)

diff --git a/app/scripts/controllers/network/network-controller.test.ts b/app/scripts/controllers/network/network-controller.test.ts
index 724f36771..41d851745 100644
--- a/app/scripts/controllers/network/network-controller.test.ts
+++ b/app/scripts/controllers/network/network-controller.test.ts
@@ -2587,6 +2587,48 @@ describe('NetworkController', () => {
     });
   });
 
+  describe('NetworkController:getEthQuery action', () => {
+    it('returns a EthQuery object that can be used to make requests to the currently selected network', async () => {
+      await withController(async ({ controller, messenger }) => {
+        await setFakeProvider(controller, {
+          stubs: [
+            {
+              request: {
+                method: 'test_method',
+                params: [],
+              },
+              response: {
+                result: 'test response',
+              },
+            },
+          ],
+        });
+
+        const ethQuery = messenger.call('NetworkController:getEthQuery');
+        assert(ethQuery, 'ethQuery is not set');
+
+        const promisifiedSendAsync = promisify(ethQuery.sendAsync).bind(
+          ethQuery,
+        );
+        const result = await promisifiedSendAsync({
+          id: 1,
+          jsonrpc: '2.0',
+          method: 'test_method',
+          params: [],
+        });
+        expect(result).toBe('test response');
+      });
+    });
+
+    it('returns undefined if the provider has not been set yet', async () => {
+      await withController(({ messenger }) => {
+        const ethQuery = messenger.call('NetworkController:getEthQuery');
+
+        expect(ethQuery).toBeUndefined();
+      });
+    });
+  });
+
   describe('rollbackToPreviousProvider', () => {
     for (const { networkType } of INFURA_NETWORKS) {
       describe(`if the previous provider configuration had a type of "${networkType}"`, () => {
@@ -6328,7 +6370,10 @@ function buildMessenger() {
 function buildNetworkControllerMessenger(messenger = buildMessenger()) {
   return messenger.getRestricted({
     name: 'NetworkController',
-    allowedActions: ['NetworkController:getProviderConfig'],
+    allowedActions: [
+      'NetworkController:getProviderConfig',
+      'NetworkController:getEthQuery',
+    ],
     allowedEvents: [
       'NetworkController:networkDidChange',
       'NetworkController:networkWillChange',
diff --git a/app/scripts/controllers/network/network-controller.ts b/app/scripts/controllers/network/network-controller.ts
index 431c49d9d..a4417a8e2 100644
--- a/app/scripts/controllers/network/network-controller.ts
+++ b/app/scripts/controllers/network/network-controller.ts
@@ -135,7 +135,14 @@ export type NetworkControllerGetProviderConfigAction = {
   handler: () => ProviderConfiguration;
 };
 
-export type NetworkControllerAction = NetworkControllerGetProviderConfigAction;
+export type NetworkControllerGetEthQueryAction = {
+  type: `NetworkController:getEthQuery`;
+  handler: () => EthQuery | undefined;
+};
+
+export type NetworkControllerAction =
+  | NetworkControllerGetProviderConfigAction
+  | NetworkControllerGetEthQueryAction;
 
 /**
  * The messenger that the NetworkController uses to publish events.
@@ -443,6 +450,8 @@ export class NetworkController extends EventEmitter {
 
   #blockTrackerProxy: SwappableProxy<PollingBlockTracker> | null;
 
+  #ethQuery: EthQuery | undefined;
+
   #infuraProjectId: NetworkControllerOptions['infuraProjectId'];
 
   #trackMetaMetricsEvent: NetworkControllerOptions['trackMetaMetricsEvent'];
@@ -486,9 +495,13 @@ export class NetworkController extends EventEmitter {
     }
     this.#infuraProjectId = infuraProjectId;
     this.#trackMetaMetricsEvent = trackMetaMetricsEvent;
+
     this.#messenger.registerActionHandler(`${name}:getProviderConfig`, () => {
       return this.store.getState().providerConfig;
     });
+    this.#messenger.registerActionHandler(`${name}:getEthQuery`, () => {
+      return this.#ethQuery;
+    });
   }
 
   /**
@@ -918,6 +931,18 @@ export class NetworkController extends EventEmitter {
     }
   }
 
+  /**
+   * Creates a new instance of EthQuery that wraps the current provider and
+   * saves it for future usage.
+   */
+  #registerProvider() {
+    const { provider } = this.getProviderAndBlockTracker();
+
+    if (provider) {
+      this.#ethQuery = new EthQuery(provider);
+    }
+  }
+
   /**
    * Creates a network client (a stack of middleware along with a provider and
    * block tracker) to talk to an Infura-supported network.
@@ -941,7 +966,7 @@ export class NetworkController extends EventEmitter {
       infuraProjectId,
       type: NetworkClientType.Infura,
     });
-    this.#setProviderAndBlockTracker({ provider, blockTracker });
+    this.#updateProvider(provider, blockTracker);
   }
 
   /**
@@ -958,7 +983,27 @@ export class NetworkController extends EventEmitter {
       rpcUrl,
       type: NetworkClientType.Custom,
     });
-    this.#setProviderAndBlockTracker({ provider, blockTracker });
+    this.#updateProvider(provider, blockTracker);
+  }
+
+  /**
+   * Given a provider and a block tracker, updates any proxies pointing to
+   * these objects that have been previously set, or initializes any proxies
+   * that have not been previously set, then creates an instance of EthQuery
+   * that wraps the provider.
+   *
+   * @param provider - The provider.
+   * @param blockTracker - The block tracker.
+   */
+  #updateProvider(
+    provider: SafeEventEmitterProvider,
+    blockTracker: PollingBlockTracker,
+  ) {
+    this.#setProviderAndBlockTracker({
+      provider,
+      blockTracker,
+    });
+    this.#registerProvider();
   }
 
   /**