feat: Add WebSocket logic to network access Snap (#3458)

FrederikBolding · web-flow · commit 8bf1e1bc9c31 · 2025-06-17T13:16:54.000Z
Expands the network access Snap to support opening a WebSocket connection and subscribing to block updates from a local Ethereum node. Additionally fixes some mistakes made in the original WebSocket PR. Closes #3456
diff --git a/packages/examples/packages/network-access/README.md b/packages/examples/packages/network-access/README.md
@@ -1,11 +1,12 @@
 # `@metamask/network-access-example-snap`
 
 This snap demonstrates how to use the `endowment:network-access` permission to
-get access to the `fetch` function from a snap.
+get access to the `fetch` function from a Snap. It also demonstrates how to use
+WebSockets.
 
 ## Snap manifest
 
-> **Note**: Using `fetch` requires the `endowment:network-access`
+> **Note**: Using `fetch` or WebSockets requires the `endowment:network-access`
 > permissions. Refer to [the documentation](https://docs.metamask.io/snaps/reference/permissions/#endowmentnetwork-access)
 > for more information.
 
@@ -29,6 +30,11 @@ JSON-RPC methods:
 
 - `fetch` - Fetch a JSON document from the optional `url`, and return the
   fetched data.
+- `startWebSocket` - Open a WebSocket connection to a local Ethereum node
+  and subscribe to block updates.
+- `stopWebSocket` - Close a WebSocket connection, if one exists.
+- `getState` - Get the state of the Snap, including the block number and whether
+  the WebSocket connection is active.
 
 For more information, you can refer to
 [the end-to-end tests](./src/index.test.ts).
diff --git a/packages/examples/packages/network-access/snap.manifest.json b/packages/examples/packages/network-access/snap.manifest.json
@@ -7,7 +7,7 @@
     "url": "https://github.com/MetaMask/snaps.git"
   },
   "source": {
-    "shasum": "jIA8HIUJKUlasU0G8ZppBteN6gVgSaBMddSWhYFhFNk=",
+    "shasum": "rrThY3oXtGgw6Jm85xyEQ0mdHgxmVC7nfb2CFQlwrsY=",
     "location": {
       "npm": {
         "filePath": "dist/bundle.js",
@@ -21,7 +21,8 @@
       "dapps": true,
       "snaps": false
     },
-    "endowment:network-access": {}
+    "endowment:network-access": {},
+    "snap_manageState": {}
   },
   "platformVersion": "8.0.0",
   "manifestVersion": "0.1"
diff --git a/packages/examples/packages/network-access/src/index.ts b/packages/examples/packages/network-access/src/index.ts
@@ -1,10 +1,13 @@
 import {
   MethodNotFoundError,
   type OnRpcRequestHandler,
+  type OnWebSocketEventHandler,
 } from '@metamask/snaps-sdk';
 import { assert } from '@metamask/utils';
 
-import type { FetchParams } from './types';
+import type { UrlParams } from './types';
+
+const DEFAULT_WEBSOCKET_URL = 'ws://localhost:8545';
 
 /**
  * Fetch a JSON file from the provided URL. This uses the standard `fetch`
@@ -16,21 +19,78 @@ import type { FetchParams } from './types';
  *
  * @param url - The URL to fetch the data from. This function assumes that the
  * provided URL is a JSON document.
- * @returns There response as JSON.
+ * @returns The response as JSON.
  * @throws If the provided URL is not a JSON document.
  */
 async function getJson(url: string) {
   const response = await fetch(url);
   return await response.json();
 }
 
+/**
+ * Open a WebSocket connection to a local Ethereum node and subscribe to
+ * block updates.
+ *
+ * @param url - The URL of the WebSocket connection.
+ * @returns Null.
+ * @throws If the WebSocket connection fails to open.
+ */
+async function subscribe(url: string = DEFAULT_WEBSOCKET_URL) {
+  const id = await snap.request({
+    method: 'snap_openWebSocket',
+    params: {
+      url,
+    },
+  });
+
+  const message = JSON.stringify({
+    jsonrpc: '2.0',
+    id: 1,
+    method: 'eth_subscribe',
+    params: ['newHeads'],
+  });
+
+  return await snap.request({
+    method: 'snap_sendWebSocketMessage',
+    params: { id, message },
+  });
+}
+
+/**
+ * Close a WebSocket connection to a local Ethereum node, if it exists.
+ *
+ * @param url - The URL of the WebSocket connection.
+ * @returns Null.
+ */
+async function unsubscribe(url: string = DEFAULT_WEBSOCKET_URL) {
+  const sockets = await snap.request({
+    method: 'snap_getWebSockets',
+  });
+
+  const socket = sockets.find((socketInfo) => socketInfo.url === url);
+
+  if (!socket) {
+    return null;
+  }
+
+  return await snap.request({
+    method: 'snap_closeWebSocket',
+    params: { id: socket.id },
+  });
+}
+
 /**
  * Handle incoming JSON-RPC requests from the dapp, sent through the
- * `wallet_invokeSnap` method. This handler handles a single method:
+ * `wallet_invokeSnap` method. This handler handles four methods:
  *
  * - `fetch`: Fetch a JSON document from the provided URL. This demonstrates
  * that a snap can make network requests through the `fetch` function. Note that
  * the `endowment:network-access` permission must be enabled for this to work.
+ * - `startWebSocket`: Open a WebSocket connection to a local Ethereum node
+ * and subscribe to block updates.
+ * - `closeWebSocket`: Close a WebSocket connection, if one exists.
+ * - `getState`: Get the state of the Snap, including the block number and whether
+ * the WebSocket connection is active.
  *
  * @param params - The request parameters.
  * @param params.request - The JSON-RPC request object.
@@ -40,14 +100,77 @@ async function getJson(url: string) {
  * @see https://docs.metamask.io/snaps/reference/permissions/#endowmentnetwork-access
  */
 export const onRpcRequest: OnRpcRequestHandler = async ({ request }) => {
+  const params = request.params as UrlParams | undefined;
+  const url = params?.url;
+
   switch (request.method) {
     case 'fetch': {
-      const params = request.params as FetchParams | undefined;
-      assert(params?.url, 'Required url parameter was not specified.');
-      return await getJson(params.url);
+      assert(url, 'Required url parameter was not specified.');
+      return await getJson(url);
+    }
+
+    case 'startWebSocket':
+      return subscribe(url);
+
+    case 'stopWebSocket':
+      return unsubscribe(url);
+
+    case 'getState': {
+      const state = await snap.request({
+        method: 'snap_getState',
+        params: { encrypted: false },
+      });
+      return state ?? { blockNumber: null, open: false };
     }
 
     default:
       throw new MethodNotFoundError({ method: request.method });
   }
 };
+
+/**
+ * Handle incoming WebSocket events sent by a client.
+ *
+ * @param params - The request parameters.
+ * @param params.event - The WebSocket event.
+ * @returns Nothing.
+ */
+export const onWebSocketEvent: OnWebSocketEventHandler = async ({ event }) => {
+  const { origin } = event;
+
+  if (event.type === 'open') {
+    await snap.request({
+      method: 'snap_setState',
+      params: {
+        value: { blockNumber: null, origin, open: true },
+        encrypted: false,
+      },
+    });
+    return;
+  }
+
+  if (event.type === 'close') {
+    await snap.request({
+      method: 'snap_setState',
+      params: {
+        value: { blockNumber: null, origin: null, open: false },
+        encrypted: false,
+      },
+    });
+    return;
+  }
+
+  assert(event.data.type === 'text');
+  const json = JSON.parse(event.data.message);
+
+  if (!json.params?.result) {
+    return;
+  }
+
+  const blockNumber = parseInt(json.params.result.number.slice(2), 16);
+
+  await snap.request({
+    method: 'snap_setState',
+    params: { key: 'blockNumber', value: blockNumber, encrypted: false },
+  });
+};
diff --git a/packages/examples/packages/network-access/src/types.ts b/packages/examples/packages/network-access/src/types.ts
@@ -1,7 +1,7 @@
 /**
- * The params for the `fetch` JSON-RPC method.
+ * Generic parameters containing an URL.
  */
-export type FetchParams = {
+export type UrlParams = {
   /**
    * The URL to fetch.
    */
diff --git a/packages/test-snaps/src/features/snaps/network-access/NetworkAccess.tsx b/packages/test-snaps/src/features/snaps/network-access/NetworkAccess.tsx
@@ -1,7 +1,7 @@
 import { logError } from '@metamask/snaps-utils';
 import type { ChangeEvent, FunctionComponent } from 'react';
 import { useState } from 'react';
-import { Button, Form } from 'react-bootstrap';
+import { Button, ButtonGroup, Form } from 'react-bootstrap';
 
 import {
   NETWORK_ACCESS_PORT,
@@ -20,14 +20,37 @@ export const NetworkAccess: FunctionComponent = () => {
     setUrl(event.target.value);
   };
 
-  const handleSubmit = () => {
+  const snapId = getSnapId(NETWORK_ACCESS_SNAP_ID, NETWORK_ACCESS_PORT);
+
+  const handleFetch = () => {
     invokeSnap({
-      snapId: getSnapId(NETWORK_ACCESS_SNAP_ID, NETWORK_ACCESS_PORT),
+      snapId,
       method: 'fetch',
       params: { url },
     }).catch(logError);
   };
 
+  const handleStartWebSocket = () => {
+    invokeSnap({
+      snapId,
+      method: 'startWebSocket',
+    }).catch(logError);
+  };
+
+  const handleStopWebSocket = () => {
+    invokeSnap({
+      snapId,
+      method: 'stopWebSocket',
+    }).catch(logError);
+  };
+
+  const handleGetState = () => {
+    invokeSnap({
+      snapId,
+      method: 'getState',
+    }).catch(logError);
+  };
+
   return (
     <Snap
       name="Network Access Snap"
@@ -49,10 +72,39 @@ export const NetworkAccess: FunctionComponent = () => {
         id="sendNetworkAccessTest"
         className="mb-3"
         disabled={isLoading}
-        onClick={handleSubmit}
+        onClick={handleFetch}
       >
         Fetch
       </Button>
+
+      <ButtonGroup className="mb-3">
+        <Button
+          variant="primary"
+          id="startWebSocket"
+          disabled={isLoading}
+          onClick={handleStartWebSocket}
+        >
+          Start WebSocket
+        </Button>
+        <Button
+          variant="primary"
+          id="stopWebSocket"
+          disabled={isLoading}
+          onClick={handleStopWebSocket}
+        >
+          Stop WebSocket
+        </Button>
+
+        <Button
+          variant="primary"
+          id="getWebSocketState"
+          disabled={isLoading}
+          onClick={handleGetState}
+        >
+          Get State
+        </Button>
+      </ButtonGroup>
+
       <Result>
         <span id="networkAccessResult">
           {JSON.stringify(data, null, 2)}
