From 05d60f7a6424075d9884e2e5e410d8a22a557ae8 Mon Sep 17 00:00:00 2001 From: Maarten Zuidhoorn Date: Tue, 24 Mar 2026 15:44:51 +0100 Subject: [PATCH 1/3] Refactor `WebSocketService` to use messenger exposed methods --- .../WebSocketService-method-action-types.ts | 63 ++++++++++++++ .../src/websocket/WebSocketService.ts | 87 +++++-------------- .../snaps-controllers/src/websocket/index.ts | 14 ++- 3 files changed, 99 insertions(+), 65 deletions(-) create mode 100644 packages/snaps-controllers/src/websocket/WebSocketService-method-action-types.ts diff --git a/packages/snaps-controllers/src/websocket/WebSocketService-method-action-types.ts b/packages/snaps-controllers/src/websocket/WebSocketService-method-action-types.ts new file mode 100644 index 0000000000..1dff4bae87 --- /dev/null +++ b/packages/snaps-controllers/src/websocket/WebSocketService-method-action-types.ts @@ -0,0 +1,63 @@ +/** + * This file is auto generated by `scripts/generate-method-action-types.ts`. + * Do not edit manually. + */ + +import type { WebSocketService } from './WebSocketService'; + +/** + * Open a WebSocket connection. + * + * @param snapId - The Snap ID. + * @param url - The URL for the WebSocket connection. + * @param protocols - An optional parameter for protocols. + * @returns The identifier for the opened connection. + * @throws If the connection fails. + */ +export type WebSocketServiceOpenAction = { + type: `WebSocketService:open`; + handler: WebSocketService['open']; +}; + +/** + * Close a given WebSocket connection. + * + * @param snapId - The Snap ID. + * @param id - The identifier for the WebSocket connection. + */ +export type WebSocketServiceCloseAction = { + type: `WebSocketService:close`; + handler: WebSocketService['close']; +}; + +/** + * Send a message from a given Snap ID to a WebSocket connection. + * + * @param snapId - The Snap ID. + * @param id - The identifier for the WebSocket connection. + * @param data - The message to send. + */ +export type WebSocketServiceSendMessageAction = { + type: `WebSocketService:sendMessage`; + handler: WebSocketService['sendMessage']; +}; + +/** + * Get a list of all open WebSocket connections for a Snap ID. + * + * @param snapId - The Snap ID. + * @returns A list of WebSocket connections. + */ +export type WebSocketServiceGetAllAction = { + type: `WebSocketService:getAll`; + handler: WebSocketService['getAll']; +}; + +/** + * Union of all WebSocketService action types. + */ +export type WebSocketServiceMethodActions = + | WebSocketServiceOpenAction + | WebSocketServiceCloseAction + | WebSocketServiceSendMessageAction + | WebSocketServiceGetAllAction; diff --git a/packages/snaps-controllers/src/websocket/WebSocketService.ts b/packages/snaps-controllers/src/websocket/WebSocketService.ts index 0b077e9ca0..9c83195526 100644 --- a/packages/snaps-controllers/src/websocket/WebSocketService.ts +++ b/packages/snaps-controllers/src/websocket/WebSocketService.ts @@ -1,14 +1,11 @@ import type { Messenger } from '@metamask/messenger'; import { rpcErrors } from '@metamask/rpc-errors'; -import type { - GetWebSocketsResult, - SnapId, - WebSocketEvent, -} from '@metamask/snaps-sdk'; +import type { SnapId, WebSocketEvent } from '@metamask/snaps-sdk'; import { HandlerType, isEqual, logError } from '@metamask/snaps-utils'; import { assert, createDeferredPromise } from '@metamask/utils'; import { nanoid } from 'nanoid'; +import type { WebSocketServiceMethodActions } from './WebSocketService-method-action-types'; import type { SnapControllerHandleRequestAction, SnapControllerSnapInstalledEvent, @@ -19,51 +16,26 @@ import { METAMASK_ORIGIN } from '../snaps'; const serviceName = 'WebSocketService'; -export type WebSocketServiceOpenAction = { - type: `${typeof serviceName}:open`; - handler: ( - snapId: SnapId, - url: string, - protocols?: string[], - ) => Promise; -}; - -export type WebSocketServiceCloseAction = { - type: `${typeof serviceName}:close`; - handler: (snapId: SnapId, id: string) => void; -}; - -export type WebSocketServiceSendMessageAction = { - type: `${typeof serviceName}:sendMessage`; - handler: ( - snapId: SnapId, - id: string, - data: string | number[], - ) => Promise; -}; - -export type WebSocketServiceGetAllAction = { - type: `${typeof serviceName}:getAll`; - handler: (snapId: SnapId) => GetWebSocketsResult; -}; +const MESSENGER_EXPOSED_METHODS = [ + 'open', + 'close', + 'sendMessage', + 'getAll', +] as const; -export type WebSocketServiceActions = - | WebSocketServiceOpenAction - | WebSocketServiceCloseAction - | WebSocketServiceSendMessageAction - | WebSocketServiceGetAllAction; +export type WebSocketServiceActions = WebSocketServiceMethodActions; -export type WebSocketServiceAllowedActions = SnapControllerHandleRequestAction; +type AllowedActions = SnapControllerHandleRequestAction; -export type WebSocketServiceEvents = +type AllowedEvents = | SnapControllerSnapUninstalledEvent | SnapControllerSnapUpdatedEvent | SnapControllerSnapInstalledEvent; export type WebSocketServiceMessenger = Messenger< 'WebSocketService', - WebSocketServiceActions | WebSocketServiceAllowedActions, - WebSocketServiceEvents + WebSocketServiceActions | AllowedActions, + AllowedEvents >; type WebSocketServiceArgs = { @@ -93,22 +65,9 @@ export class WebSocketService { this.#messenger = messenger; this.#sockets = new Map(); - this.#messenger.registerActionHandler( - `${serviceName}:open`, - async (...args) => this.#open(...args), - ); - - this.#messenger.registerActionHandler(`${serviceName}:close`, (...args) => - this.#close(...args), - ); - - this.#messenger.registerActionHandler( - `${serviceName}:sendMessage`, - async (...args) => this.#sendMessage(...args), - ); - - this.#messenger.registerActionHandler(`${serviceName}:getAll`, (...args) => - this.#getAll(...args), + this.#messenger.registerMethodActionHandlers( + this, + MESSENGER_EXPOSED_METHODS, ); this.#messenger.subscribe('SnapController:snapUpdated', (snap) => { @@ -154,7 +113,7 @@ export class WebSocketService { * @returns True if a matching connection already exists, otherwise false. */ #exists(snapId: SnapId, url: string, protocols: string[]) { - return this.#getAll(snapId).some( + return this.getAll(snapId).some( (socket) => socket.url === url && isEqual(socket.protocols, protocols), ); } @@ -190,7 +149,7 @@ export class WebSocketService { * @returns The identifier for the opened connection. * @throws If the connection fails. */ - async #open(snapId: SnapId, url: string, protocols: string[] = []) { + async open(snapId: SnapId, url: string, protocols: string[] = []) { assert( !this.#exists(snapId, url, protocols), `An open WebSocket connection to ${url} already exists.`, @@ -280,7 +239,7 @@ export class WebSocketService { * @param snapId - The Snap ID. * @param id - The identifier for the WebSocket connection. */ - #close(snapId: SnapId, id: string) { + close(snapId: SnapId, id: string) { const { socket } = this.#get(snapId, id); socket.close(); @@ -292,8 +251,8 @@ export class WebSocketService { * @param snapId - The Snap ID. */ #closeAll(snapId: SnapId) { - for (const socket of this.#getAll(snapId)) { - this.#close(snapId, socket.id); + for (const socket of this.getAll(snapId)) { + this.close(snapId, socket.id); } } @@ -304,7 +263,7 @@ export class WebSocketService { * @param id - The identifier for the WebSocket connection. * @param data - The message to send. */ - async #sendMessage(snapId: SnapId, id: string, data: string | number[]) { + async sendMessage(snapId: SnapId, id: string, data: string | number[]) { const { socket, openPromise } = this.#get(snapId, id); await openPromise; @@ -320,7 +279,7 @@ export class WebSocketService { * @param snapId - The Snap ID. * @returns A list of WebSocket connections. */ - #getAll(snapId: SnapId) { + getAll(snapId: SnapId) { return [...this.#sockets.values()] .filter((socket) => socket.snapId === snapId) .map((socket) => ({ diff --git a/packages/snaps-controllers/src/websocket/index.ts b/packages/snaps-controllers/src/websocket/index.ts index f0a947eb67..35e9be0ac8 100644 --- a/packages/snaps-controllers/src/websocket/index.ts +++ b/packages/snaps-controllers/src/websocket/index.ts @@ -1 +1,13 @@ -export * from './WebSocketService'; +export type { + WebSocketServiceActions, + WebSocketServiceEvents, + WebSocketServiceMessenger, +} from './WebSocketService'; +export { WebSocketService } from './WebSocketService'; +export type { + WebSocketServiceCloseAction, + WebSocketServiceGetAllAction, + WebSocketServiceMethodActions, + WebSocketServiceOpenAction, + WebSocketServiceSendMessageAction, +} from './WebSocketService-method-action-types'; From 61ef968950ca2a4d62fa065594e1514ad3806b42 Mon Sep 17 00:00:00 2001 From: Maarten Zuidhoorn Date: Tue, 24 Mar 2026 15:47:45 +0100 Subject: [PATCH 2/3] Remove export --- packages/snaps-controllers/src/websocket/index.ts | 1 - 1 file changed, 1 deletion(-) diff --git a/packages/snaps-controllers/src/websocket/index.ts b/packages/snaps-controllers/src/websocket/index.ts index 35e9be0ac8..b9f6a13d3a 100644 --- a/packages/snaps-controllers/src/websocket/index.ts +++ b/packages/snaps-controllers/src/websocket/index.ts @@ -1,6 +1,5 @@ export type { WebSocketServiceActions, - WebSocketServiceEvents, WebSocketServiceMessenger, } from './WebSocketService'; export { WebSocketService } from './WebSocketService'; From 44b0f399df1ea86bf6b4bd229005ed5d67db2c4a Mon Sep 17 00:00:00 2001 From: Maarten Zuidhoorn Date: Tue, 24 Mar 2026 15:47:57 +0100 Subject: [PATCH 3/3] Remove another export --- packages/snaps-controllers/src/websocket/index.ts | 1 - 1 file changed, 1 deletion(-) diff --git a/packages/snaps-controllers/src/websocket/index.ts b/packages/snaps-controllers/src/websocket/index.ts index b9f6a13d3a..1d535dc632 100644 --- a/packages/snaps-controllers/src/websocket/index.ts +++ b/packages/snaps-controllers/src/websocket/index.ts @@ -6,7 +6,6 @@ export { WebSocketService } from './WebSocketService'; export type { WebSocketServiceCloseAction, WebSocketServiceGetAllAction, - WebSocketServiceMethodActions, WebSocketServiceOpenAction, WebSocketServiceSendMessageAction, } from './WebSocketService-method-action-types';