Add snap_getWebSockets

Author: FrederikBolding

packages/snaps-execution-environments/coverage.json

@@ -2,5 +2,5 @@
   "branches": 90,
   "functions": 94.69,
   "lines": 90.41,
-  "statements": 89.64
+  "statements": 89.63
 }

packages/snaps-rpc-methods/jest.config.js

@@ -10,10 +10,10 @@ module.exports = deepmerge(baseConfig, {
   ],
   coverageThreshold: {
     global: {
-      branches: 95.06,
-      functions: 98.68,
-      lines: 98.83,
-      statements: 98.52,
+      branches: 95.1,
+      functions: 98.69,
+      lines: 98.84,
+      statements: 98.55,
     },
   },
 });

packages/snaps-rpc-methods/src/permitted/getWebSockets.test.ts

@@ -0,0 +1,142 @@
+import { JsonRpcEngine } from '@metamask/json-rpc-engine';
+import { rpcErrors } from '@metamask/rpc-errors';
+import type {
+  GetWebSocketsParams,
+  GetWebSocketsResult,
+} from '@metamask/snaps-sdk';
+import type { JsonRpcRequest, PendingJsonRpcResponse } from '@metamask/utils';
+
+import { getWebSocketsHandler } from './getWebSockets';
+
+describe('snap_getWebSockets', () => {
+  describe('openWebSocketHandler', () => {
+    it('has the expected shape', () => {
+      expect(getWebSocketsHandler).toMatchObject({
+        methodNames: ['snap_getWebSockets'],
+        implementation: expect.any(Function),
+        hookNames: {
+          hasPermission: true,
+          getWebSockets: true,
+        },
+      });
+    });
+  });
+
+  describe('implementation', () => {
+    it('throws if the origin does not have permission', async () => {
+      const { implementation } = getWebSocketsHandler;
+
+      const getWebSockets = jest.fn();
+      const hasPermission = jest.fn().mockReturnValue(false);
+      const hooks = { hasPermission, getWebSockets };
+
+      const engine = new JsonRpcEngine();
+
+      engine.push((request, response, next, end) => {
+        const result = implementation(
+          request as JsonRpcRequest<GetWebSocketsParams>,
+          response as PendingJsonRpcResponse<GetWebSocketsResult>,
+          next,
+          end,
+          hooks,
+        );
+
+        result?.catch(end);
+      });
+
+      const response = await engine.handle({
+        jsonrpc: '2.0',
+        id: 1,
+        method: 'snap_getWebSockets',
+      });
+
+      expect(response).toStrictEqual({
+        error: {
+          code: 4100,
+          message:
+            'The requested account and/or method has not been authorized by the user.',
+          stack: expect.any(String),
+        },
+        id: 1,
+        jsonrpc: '2.0',
+      });
+    });
+
+    it('returns the open WebSocket connections', async () => {
+      const { implementation } = getWebSocketsHandler;
+
+      const getWebSockets = jest
+        .fn()
+        .mockReturnValue([{ id: 'foo', url: 'wss://metamask.io' }]);
+      const hasPermission = jest.fn().mockReturnValue(true);
+      const hooks = { hasPermission, getWebSockets };
+
+      const engine = new JsonRpcEngine();
+
+      engine.push((request, response, next, end) => {
+        const result = implementation(
+          request as JsonRpcRequest<GetWebSocketsParams>,
+          response as PendingJsonRpcResponse<GetWebSocketsResult>,
+          next,
+          end,
+          hooks,
+        );
+
+        result?.catch(end);
+      });
+
+      const response = await engine.handle({
+        jsonrpc: '2.0',
+        id: 1,
+        method: 'snap_getWebSockets',
+      });
+
+      expect(response).toStrictEqual({
+        jsonrpc: '2.0',
+        id: 1,
+        result: [{ id: 'foo', url: 'wss://metamask.io' }],
+      });
+      expect(getWebSockets).toHaveBeenCalledWith();
+    });
+
+    it('returns an error thrown by the hook', async () => {
+      const { implementation } = getWebSocketsHandler;
+
+      const getWebSockets = jest.fn().mockImplementation(() => {
+        throw rpcErrors.internal();
+      });
+      const hasPermission = jest.fn().mockReturnValue(true);
+      const hooks = { hasPermission, getWebSockets };
+
+      const engine = new JsonRpcEngine();
+
+      engine.push((request, response, next, end) => {
+        const result = implementation(
+          request as JsonRpcRequest<GetWebSocketsParams>,
+          response as PendingJsonRpcResponse<GetWebSocketsResult>,
+          next,
+          end,
+          hooks,
+        );
+
+        result?.catch(end);
+      });
+
+      const response = await engine.handle({
+        jsonrpc: '2.0',
+        id: 1,
+        method: 'snap_getWebSockets',
+      });
+
+      expect(response).toStrictEqual({
+        error: {
+          code: -32603,
+          message: 'Internal JSON-RPC error.',
+          stack: expect.any(String),
+        },
+        id: 1,
+        jsonrpc: '2.0',
+      });
+    });
+  });
+});

packages/snaps-rpc-methods/src/permitted/getWebSockets.ts

@@ -0,0 +1,67 @@
+import type { JsonRpcEngineEndCallback } from '@metamask/json-rpc-engine';
+import type { PermittedHandlerExport } from '@metamask/permission-controller';
+import { providerErrors } from '@metamask/rpc-errors';
+import type {
+  GetWebSocketsParams,
+  GetWebSocketsResult,
+  JsonRpcRequest,
+} from '@metamask/snaps-sdk';
+import type { PendingJsonRpcResponse } from '@metamask/utils';
+
+import { SnapEndowments } from '../endowments';
+import type { MethodHooksObject } from '../utils';
+
+const hookNames: MethodHooksObject<GetWebSocketsMethodHooks> = {
+  hasPermission: true,
+  getWebSockets: true,
+};
+
+export type GetWebSocketsMethodHooks = {
+  hasPermission: (permissionName: string) => boolean;
+  getWebSockets: () => GetWebSocketsResult;
+};
+
+/**
+ * Handler for the `snap_getWebSockets` method.
+ */
+export const getWebSocketsHandler: PermittedHandlerExport<
+  GetWebSocketsMethodHooks,
+  GetWebSocketsParams,
+  GetWebSocketsResult
+> = {
+  methodNames: ['snap_getWebSockets'],
+  implementation: getWebSocketsImplementation,
+  hookNames,
+};
+
+/**
+ * The `snap_getWebSockets` method implementation.
+ *
+ * @param _req - The JSON-RPC request object. Not used by this function.
+ * @param res - The JSON-RPC response object.
+ * @param _next - The `json-rpc-engine` "next" callback. Not used by this function.
+ * @param end - The `json-rpc-engine` "end" callback.
+ * @param hooks - The RPC method hooks.
+ * @param hooks.hasPermission - The function to check if a snap has the `endowment:network-access` permission.
+ * @param hooks.getWebSockets - The function to get the connected WebSockets for the origin.
+ * @returns Nothing.
+ */
+function getWebSocketsImplementation(
+  _req: JsonRpcRequest<GetWebSocketsParams>,
+  res: PendingJsonRpcResponse<GetWebSocketsResult>,
+  _next: unknown,
+  end: JsonRpcEngineEndCallback,
+  { hasPermission, getWebSockets }: GetWebSocketsMethodHooks,
+): void {
+  if (!hasPermission(SnapEndowments.NetworkAccess)) {
+    return end(providerErrors.unauthorized());
+  }
+
+  try {
+    res.result = getWebSockets();
+  } catch (error) {
+    return end(error);
+  }
+
+  return end();
+}

packages/snaps-rpc-methods/src/permitted/handlers.ts

@@ -12,6 +12,7 @@ import { getInterfaceContextHandler } from './getInterfaceContext';
 import { getInterfaceStateHandler } from './getInterfaceState';
 import { getSnapsHandler } from './getSnaps';
 import { getStateHandler } from './getState';
+import { getWebSocketsHandler } from './getWebSockets';
 import { invokeKeyringHandler } from './invokeKeyring';
 import { invokeSnapSugarHandler } from './invokeSnapSugar';
 import { listEntropySourcesHandler } from './listEntropySources';
@@ -51,6 +52,7 @@ export const methodHandlers = {
   snap_openWebSocket: openWebSocketHandler,
   snap_closeWebSocket: closeWebSocketHandler,
   snap_sendWebSocketMessage: sendWebSocketMessageHandler,
+  snap_getWebSockets: getWebSocketsHandler,
 };
 /* eslint-enable @typescript-eslint/naming-convention */
 

packages/snaps-rpc-methods/src/permitted/index.ts

@@ -1,5 +1,6 @@
 import type { CancelBackgroundEventMethodHooks } from './cancelBackgroundEvent';
 import type { ClearStateHooks } from './clearState';
+import type { CloseWebSocketMethodHooks } from './closeWebSocket';
 import type { CreateInterfaceMethodHooks } from './createInterface';
 import type { ProviderRequestMethodHooks } from './experimentalProviderRequest';
 import type { GetAllSnapsHooks } from './getAllSnaps';
@@ -9,10 +10,13 @@ import type { GetCurrencyRateMethodHooks } from './getCurrencyRate';
 import type { GetInterfaceStateMethodHooks } from './getInterfaceState';
 import type { GetSnapsHooks } from './getSnaps';
 import type { GetStateHooks } from './getState';
+import type { GetWebSocketsMethodHooks } from './getWebSockets';
 import type { ListEntropySourcesHooks } from './listEntropySources';
+import type { OpenWebSocketMethodHooks } from './openWebSocket';
 import type { RequestSnapsHooks } from './requestSnaps';
 import type { ResolveInterfaceMethodHooks } from './resolveInterface';
 import type { ScheduleBackgroundEventMethodHooks } from './scheduleBackgroundEvent';
+import type { SendWebSocketMessageMethodHooks } from './sendWebSocketMessage';
 import type { SetStateHooks } from './setState';
 import type { UpdateInterfaceMethodHooks } from './updateInterface';
 
@@ -32,7 +36,11 @@ export type PermittedRpcMethodHooks = ClearStateHooks &
   ScheduleBackgroundEventMethodHooks &
   CancelBackgroundEventMethodHooks &
   GetBackgroundEventsMethodHooks &
-  SetStateHooks;
+  SetStateHooks &
+  OpenWebSocketMethodHooks &
+  CloseWebSocketMethodHooks &
+  SendWebSocketMessageMethodHooks &
+  GetWebSocketsMethodHooks;
 
 export * from './handlers';
 export * from './middleware';

packages/snaps-sdk/src/types/methods/get-web-sockets.ts

@@ -0,0 +1,3 @@
+export type GetWebSocketsParams = never;
+
+export type GetWebSocketsResult = { id: string; url: string }[];

packages/snaps-sdk/src/types/methods/index.ts

@@ -33,3 +33,4 @@ export type * from './track-event';
 export type * from './open-web-socket';
 export type * from './close-web-socket';
 export type * from './send-web-socket-message';
+export type * from './get-web-sockets';