Small improvements (#440)

- Clears associated recovery callbacks when closing a connection (see
[this
convo](#431 (comment)))
- Simplify some code by embedding the network field on the
`WebSocketConnection` entity
- Can now manually close connections from the test dapp:

![image](https://github.com/user-attachments/assets/9c90d100-03a3-49cf-bdfb-127be4bc3e77)

Author: xavier-brochard

packages/site/src/components/Handlers/WebSockets.tsx

@@ -23,6 +23,18 @@ export const WebSockets = () => {
     });
   };
 
+  const closeAllConnections = async () => {
+    showToasterForResponse(
+      { result: 'ok' },
+      {
+        title: 'Closed all WebSocket connections',
+      },
+    );
+    await invokeSnap({
+      method: TestDappRpcRequestMethod.TestCloseAllConnections,
+    });
+  };
+
   const subscribeToAccount = async () => {
     await invokeSnap({
       method: TestDappRpcRequestMethod.TestSubscribeToAccount,
@@ -52,6 +64,13 @@ export const WebSockets = () => {
           >
             Setup All Connections
           </Button>
+          <Button
+            variant="outline"
+            onClick={closeAllConnections}
+            marginRight="1"
+          >
+            Close All Connections
+          </Button>
         </Flex>
         <Flex>
           <Button

packages/snap/snap.manifest.json

@@ -7,7 +7,7 @@
     "url": "https://github.com/MetaMask/snap-solana-wallet.git"
   },
   "source": {
-    "shasum": "g/jfpUCvmwu6piUwA21tuyxrmaZZnZl3JdLeVJysYFM=",
+    "shasum": "EBrVuE+uXKdUKPSGja1momxWXqItfVGppjHtfMaDo+I=",
     "location": {
       "npm": {
         "filePath": "dist/bundle.js",

packages/snap/src/core/handlers/onRpcRequest/index.ts

@@ -18,4 +18,8 @@ export const handlers: Record<RpcRequestMethod, OnRpcRequestHandler> = {
     await eventEmitter.emitSync('onTestSubscribeToAccount');
     return null;
   },
+  [TestDappRpcRequestMethod.TestCloseAllConnections as any]: async () => {
+    await eventEmitter.emitSync('onTestCloseAllConnections');
+    return null;
+  },
 };

packages/snap/src/core/handlers/onRpcRequest/types.ts

@@ -10,4 +10,5 @@ export enum RpcRequestMethod {
 export enum TestDappRpcRequestMethod {
   TestSetupAllConnections = 'testSetupAllConnections',
   TestSubscribeToAccount = 'testSubscribeToAccount',
+  TestCloseAllConnections = 'testCloseAllConnections',
 }

packages/snap/src/core/services/subscriptions/WebSocketConnectionRepository.test.ts

@@ -1,10 +1,24 @@
+import { Network } from '../../constants/solana';
+import type { ConfigProvider } from '../config';
 import { WebSocketConnectionRepository } from './WebSocketConnectionRepository';
 
+const mockNetwork = Network.Mainnet;
+const mockWebSocketUrl = 'ws://localhost:8080';
+const mockConnectionId = 'some-connection-id';
+
 describe('WebSocketConnectionRepository', () => {
   let repository: WebSocketConnectionRepository;
+  let mockConfigProvider: ConfigProvider;
 
   beforeEach(() => {
-    repository = new WebSocketConnectionRepository();
+    mockConfigProvider = {
+      getNetworkBy: jest.fn().mockReturnValue({
+        caip2Id: mockNetwork,
+        webSocketUrl: mockWebSocketUrl,
+      }),
+    } as unknown as ConfigProvider;
+
+    repository = new WebSocketConnectionRepository(mockConfigProvider);
 
     const snap = {
       request: jest.fn(),
@@ -24,11 +38,11 @@ describe('WebSocketConnectionRepository', () => {
     it('returns all connections opened in the extension', async () => {
       jest
         .spyOn(snap, 'request')
-        .mockResolvedValue([{ id: '1', url: 'ws://localhost:8080' }]);
+        .mockResolvedValue([{ id: mockConnectionId, url: mockWebSocketUrl }]);
 
       const connections = await repository.getAll();
       expect(connections).toStrictEqual([
-        { id: '1', url: 'ws://localhost:8080' },
+        { id: mockConnectionId, url: mockWebSocketUrl, network: mockNetwork },
       ]);
     });
   });
@@ -45,53 +59,70 @@ describe('WebSocketConnectionRepository', () => {
     it('returns the connection when it exists', async () => {
       jest
         .spyOn(snap, 'request')
-        .mockResolvedValue([{ id: '1', url: 'ws://localhost:8080' }]);
+        .mockResolvedValue([{ id: mockConnectionId, url: mockWebSocketUrl }]);
 
-      const connection = await repository.getById('1');
+      const connection = await repository.getById(mockConnectionId);
 
-      expect(connection).toStrictEqual({ id: '1', url: 'ws://localhost:8080' });
+      expect(connection).toStrictEqual({
+        id: mockConnectionId,
+        url: mockWebSocketUrl,
+        network: mockNetwork,
+      });
     });
   });
 
   describe('findByUrl', () => {
     it('returns null when the connection does not exist', async () => {
       jest.spyOn(snap, 'request').mockResolvedValue([]);
 
-      const connection = await repository.findByUrl('ws://localhost:8080');
+      const connection = await repository.findByNetwork(mockNetwork);
 
       expect(connection).toBeNull();
     });
 
     it('returns the connection when it exists', async () => {
       jest
         .spyOn(snap, 'request')
-        .mockResolvedValue([{ id: '1', url: 'ws://localhost:8080' }]);
+        .mockResolvedValue([{ id: mockConnectionId, url: mockWebSocketUrl }]);
 
-      const connection = await repository.findByUrl('ws://localhost:8080');
+      const connection = await repository.findByNetwork(mockNetwork);
 
-      expect(connection).toStrictEqual({ id: '1', url: 'ws://localhost:8080' });
+      expect(connection).toStrictEqual({
+        id: mockConnectionId,
+        url: mockWebSocketUrl,
+        network: mockNetwork,
+      });
     });
   });
 
   describe('save', () => {
     it('saves the connection', async () => {
-      jest.spyOn(snap, 'request').mockResolvedValue('1');
+      jest.spyOn(snap, 'request').mockResolvedValue(mockConnectionId);
 
-      const connectionId = await repository.save('ws://localhost:8080');
+      const connection = await repository.save({
+        network: mockNetwork,
+        url: mockWebSocketUrl,
+        protocols: [],
+      });
 
-      expect(connectionId).toBe('1');
+      expect(connection).toStrictEqual({
+        id: mockConnectionId,
+        network: mockNetwork,
+        url: mockWebSocketUrl,
+        protocols: [],
+      });
     });
   });
 
   describe('delete', () => {
     it('deletes the connection', async () => {
       jest.spyOn(snap, 'request').mockResolvedValue(null);
 
-      await repository.delete('1');
+      await repository.delete(mockConnectionId);
 
       expect(snap.request).toHaveBeenCalledWith({
         method: 'snap_closeWebSocket',
-        params: { id: '1' },
+        params: { id: mockConnectionId },
       });
     });
   });

packages/snap/src/core/services/subscriptions/WebSocketConnectionRepository.ts

@@ -1,4 +1,6 @@
 import type { WebSocketConnection } from '../../../entities';
+import type { Network } from '../../constants/solana';
+import type { ConfigProvider } from '../config';
 
 /**
  * Repository that is treating the Snap's WebSocket storage as a persistent data store, where:
@@ -9,43 +11,76 @@ import type { WebSocketConnection } from '../../../entities';
  * It also tracks bidirectional mappings between the connection ID and the URL to perform fast lookups.
  */
 export class WebSocketConnectionRepository {
+  readonly #configProvider: ConfigProvider;
+
+  constructor(configProvider: ConfigProvider) {
+    this.#configProvider = configProvider;
+  }
+
+  /**
+   * Gets all connections.
+   * @returns All connections.
+   */
   async getAll(): Promise<WebSocketConnection[]> {
-    return snap.request({
+    const snapConnections = await snap.request({
       method: 'snap_getWebSockets',
     });
+
+    // Enhance the connections with the network
+    return snapConnections.map((connection) => ({
+      ...connection,
+      network: this.#findNetworkByWebSocketUrl(connection.url),
+    }));
   }
 
+  /**
+   * Gets the connection for the specified ID.
+   * @param id - The ID of the connection to get.
+   * @returns The connection, or null if no connection exists for the ID.
+   */
   async getById(id: string): Promise<WebSocketConnection | null> {
     const existingConnections = await this.getAll();
     return (
       existingConnections.find((connection) => connection.id === id) ?? null
     );
   }
 
-  async findByUrl(url: string): Promise<WebSocketConnection | null> {
+  /**
+   * Finds the connection for the specified network.
+   * @param network - The network to find the connection for.
+   * @returns The connection, or null if no connection exists for the network.
+   */
+  async findByNetwork(network: Network): Promise<WebSocketConnection | null> {
     const existingConnections = await this.getAll();
-
     return (
-      existingConnections.find((connection) => connection.url === url) ?? null
+      existingConnections.find(
+        (connection) => connection.network === network,
+      ) ?? null
     );
   }
 
   /**
    * Creates a new connection to the specified URL.
-   * @param url - The URL of the connection.
-   * @param protocols - The protocols of the connection.
+   * @param connection - The connection to create, without the `id` field.
    * @returns The connection ID.
    */
-  async save(url: string, protocols?: string[]) {
-    const connectionId = await snap.request({
+  async save(
+    connection: Omit<WebSocketConnection, 'id'>,
+  ): Promise<WebSocketConnection> {
+    const { url, protocols } = connection;
+
+    const id = await snap.request({
       method: 'snap_openWebSocket',
       params: {
         url,
         ...(protocols ? { protocols } : {}),
       },
     });
 
-    return connectionId;
+    return {
+      ...connection,
+      id,
+    };
   }
 
   /**
@@ -58,4 +93,22 @@ export class WebSocketConnectionRepository {
       params: { id },
     });
   }
+
+  /**
+   * Gets the network for the specified connection ID.
+   * @param webSocketUrl - The WebSocket URL to get the network for.
+   * @returns The network.
+   */
+  #findNetworkByWebSocketUrl(webSocketUrl: string): Network {
+    const network = this.#configProvider.getNetworkBy(
+      'webSocketUrl',
+      webSocketUrl,
+    );
+
+    if (!network) {
+      throw new Error(`No network found for WebSocket URL: ${webSocketUrl}`);
+    }
+
+    return network.caip2Id;
+  }
 }