Use selected network controller for Snaps

packages/selected-network-controller/src/SelectedNetworkController.ts

@@ -26,12 +26,6 @@ const stateMetadata = {
 
 const getDefaultState = () => ({ domains: {} });
 
-// npm and local are currently the only valid prefixes for snap domains
-// TODO: eventually we maybe want to pull this in from snaps-utils to ensure it stays in sync
-// For now it seems like overkill to add a dependency for this one constant
-// https://github.com/MetaMask/snaps/blob/2beee7803bfe9e540788a3558b546b9f55dc3cb4/packages/snaps-utils/src/types.ts#L120
-const snapsPrefixes = ['npm:', 'local:'] as const;
-
 export type Domain = string;
 
 export const METAMASK_DOMAIN = 'metamask' as const;
@@ -357,10 +351,6 @@ export class SelectedNetworkController extends BaseController<
       );
     }
 
-    if (snapsPrefixes.some((prefix) => domain.startsWith(prefix))) {
-      return;
-    }
-
     if (!this.#domainHasPermissions(domain)) {
       throw new Error(
         'NetworkClientId for domain cannot be called with a domain that has not yet been granted permissions',
@@ -386,11 +376,8 @@ export class SelectedNetworkController extends BaseController<
    * @returns The proxy and block tracker proxies.
    */
   getProviderAndBlockTracker(domain: Domain): NetworkProxy {
-    // If the domain is 'metamask' or a snap, return the NetworkController's globally selected network client proxy
-    if (
-      domain === METAMASK_DOMAIN ||
-      snapsPrefixes.some((prefix) => domain.startsWith(prefix))
-    ) {
+    // If the domain is 'metamask', return the NetworkController's globally selected network client proxy
+    if (domain === METAMASK_DOMAIN) {
       const networkClient = this.messagingSystem.call(
         'NetworkController:getSelectedNetworkClient',
       );

packages/selected-network-controller/tests/SelectedNetworkController.test.ts

@@ -525,33 +525,46 @@ describe('SelectedNetworkController', () => {
       });
 
       describe('when the requesting domain is a snap (starts with "npm:" or "local:"', () => {
-        it('skips setting the networkClientId for the passed in domain', () => {
+        it('sets the networkClientId for the passed in snap ID', () => {
           const { controller, mockHasPermissions } = setup({
             state: { domains: {} },
             useRequestQueuePreference: true,
           });
           mockHasPermissions.mockReturnValue(true);
-          const snapDomainOne = 'npm:@metamask/bip32-example-snap';
-          const snapDomainTwo = 'local:@metamask/bip32-example-snap';
-          const nonSnapDomain = 'example.com';
+          const domain = 'npm:foo-snap';
           const networkClientId = 'network1';
+          controller.setNetworkClientIdForDomain(domain, networkClientId);
+          expect(controller.state.domains[domain]).toBe(networkClientId);
+        });
 
+        it('updates the provider and block tracker proxy when they already exist for the snap ID', () => {
+          const { controller, mockProviderProxy, mockHasPermissions } = setup({
+            state: { domains: {} },
+            useRequestQueuePreference: true,
+          });
+          mockHasPermissions.mockReturnValue(true);
+          const initialNetworkClientId = '123';
+
+          // creates the proxy for the new domain
           controller.setNetworkClientIdForDomain(
-            nonSnapDomain,
-            networkClientId,
-          );
-          controller.setNetworkClientIdForDomain(
-            snapDomainOne,
-            networkClientId,
+            'npm:foo-snap',
+            initialNetworkClientId,
           );
+          const newNetworkClientId = 'abc';
+
+          expect(mockProviderProxy.setTarget).toHaveBeenCalledTimes(1);
+
+          // calls setTarget on the proxy
           controller.setNetworkClientIdForDomain(
-            snapDomainTwo,
-            networkClientId,
+            'npm:foo-snap',
+            newNetworkClientId,
           );
 
-          expect(controller.state.domains).toStrictEqual({
-            [nonSnapDomain]: networkClientId,
-          });
+          expect(mockProviderProxy.setTarget).toHaveBeenNthCalledWith(
+            2,
+            expect.objectContaining({ request: expect.any(Function) }),
+          );
+          expect(mockProviderProxy.setTarget).toHaveBeenCalledTimes(2);
         });
       });
 
@@ -774,23 +787,22 @@ describe('SelectedNetworkController', () => {
 
     // TODO - improve these tests by using a full NetworkController and doing more robust behavioral testing
     describe('when the domain is a snap (starts with "npm:" or "local:")', () => {
-      it('returns a proxied globally selected networkClient and does not create a new proxy in the domainProxyMap', () => {
-        const { controller, domainProxyMap, messenger } = setup({
+      it('calls to NetworkController:getSelectedNetworkClient and creates a new proxy provider and block tracker with the proxied globally selected network client', () => {
+        const { controller, messenger } = setup({
           state: {
             domains: {},
           },
-          useRequestQueuePreference: true,
+          useRequestQueuePreference: false,
         });
         jest.spyOn(messenger, 'call');
-        const snapDomain = 'npm:@metamask/bip32-example-snap';
-
-        const result = controller.getProviderAndBlockTracker(snapDomain);
 
-        expect(domainProxyMap.get(snapDomain)).toBeUndefined();
+        const result = controller.getProviderAndBlockTracker('npm:foo-snap');
+        expect(result).toBeDefined();
+        // unfortunately checking which networkController method is called is the best
+        // proxy (no pun intended) for checking that the correct instance of the networkClient is used
         expect(messenger.call).toHaveBeenCalledWith(
           'NetworkController:getSelectedNetworkClient',
         );
-        expect(result).toBeDefined();
       });
 
       it('throws an error if the globally selected network client is not initialized', () => {