From 5652d07908443dd5990f6ced24ae12395750304a Mon Sep 17 00:00:00 2001
From: Alex <adonesky@gmail.com>
Date: Fri, 16 Feb 2024 14:20:44 -0600
Subject: [PATCH] update domains state such that when a networkClient is
 deleted any domains that were pointed at that networkClient are repointed to
 the globally selected network

---
 .../src/SelectedNetworkController.ts          | 22 ++++++++++
 .../tests/SelectedNetworkController.test.ts   | 40 ++++++++++++++++++-
 2 files changed, 61 insertions(+), 1 deletion(-)

diff --git a/packages/selected-network-controller/src/SelectedNetworkController.ts b/packages/selected-network-controller/src/SelectedNetworkController.ts
index f51497bba18..20cccbbbd2f 100644
--- a/packages/selected-network-controller/src/SelectedNetworkController.ts
+++ b/packages/selected-network-controller/src/SelectedNetworkController.ts
@@ -135,6 +135,28 @@ export class SelectedNetworkController extends BaseController<
       state,
     });
     this.#registerMessageHandlers();
+
+    this.messagingSystem.subscribe(
+      'NetworkController:stateChange',
+      ({ selectedNetworkClientId }, patches) => {
+        patches.forEach(({ op, path }) => {
+          // if a network is removed, update the networkClientId for all domains that were using it to the selected network
+          if (op === 'remove' && path[0] === 'networkConfigurations') {
+            const removedNetworkClientId = path[1] as NetworkClientId;
+            Object.entries(this.state.domains).forEach(
+              ([domain, networkClientIdForDomain]) => {
+                if (networkClientIdForDomain === removedNetworkClientId) {
+                  this.setNetworkClientIdForDomain(
+                    domain,
+                    selectedNetworkClientId,
+                  );
+                }
+              },
+            );
+          }
+        });
+      },
+    );
   }
 
   #registerMessageHandlers(): void {
diff --git a/packages/selected-network-controller/tests/SelectedNetworkController.test.ts b/packages/selected-network-controller/tests/SelectedNetworkController.test.ts
index 61d41ae4f0a..9a202639e10 100644
--- a/packages/selected-network-controller/tests/SelectedNetworkController.test.ts
+++ b/packages/selected-network-controller/tests/SelectedNetworkController.test.ts
@@ -118,7 +118,10 @@ const setup = ({
   });
   const messenger = buildMessenger();
   const selectedNetworkControllerMessenger =
-    buildSelectedNetworkControllerMessenger({ messenger, hasPermissions });
+    buildSelectedNetworkControllerMessenger({
+      messenger,
+      hasPermissions,
+    });
   const controller = new SelectedNetworkController({
     messenger: selectedNetworkControllerMessenger,
     state,
@@ -158,6 +161,41 @@ describe('SelectedNetworkController', () => {
     });
   });
 
+  describe('It updates domain state when the network controller state changes', () => {
+    describe('when a networkClient is deleted from the network controller state', () => {
+      it('updates the networkClientId for domains which were previously set to the deleted networkClientId', () => {
+        const { controller, messenger } = setup({
+          state: {
+            perDomainNetwork: true,
+            domains: {
+              metamask: 'goerli',
+              'example.com': 'test-network-client-id',
+              'test.com': 'test-network-client-id',
+            },
+          },
+        });
+
+        messenger.publish(
+          'NetworkController:stateChange',
+          {
+            providerConfig: { chainId: '0x5', ticker: 'ETH', type: 'goerli' },
+            selectedNetworkClientId: 'goerli',
+            networkConfigurations: {},
+            networksMetadata: {},
+          },
+          [
+            {
+              op: 'remove',
+              path: ['networkConfigurations', 'test-network-client-id'],
+            },
+          ],
+        );
+        expect(controller.state.domains['example.com']).toBe('goerli');
+        expect(controller.state.domains['test.com']).toBe('goerli');
+      });
+    });
+  });
+
   describe('setNetworkClientIdForDomain', () => {
     afterEach(() => {
       jest.clearAllMocks();