Simplify

Author: FrederikBolding

packages/examples/packages/network-access/snap.manifest.json

@@ -7,7 +7,7 @@
     "url": "https://github.com/MetaMask/snaps.git"
   },
   "source": {
-    "shasum": "lQOEArGTw4fxV2GxGPHqPhG2RQlGQLQvnZXTDk85qLQ=",
+    "shasum": "2HlA1Jt3eX8cW5/Nudr49NmNjTlR7iw9CgGVP+1pGPs=",
     "location": {
       "npm": {
         "filePath": "dist/bundle.js",

packages/examples/packages/network-access/src/index.ts

@@ -19,21 +19,84 @@ const WEBSOCKET_URL = 'ws://localhost:8545';
  *
  * @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.
+ *
+ * @returns Null.
+ * @throws If the WebSocket connection fails to open.
+ */
+async function subscribe() {
+  const id = await snap.request({
+    method: 'snap_openWebSocket',
+    params: {
+      url: WEBSOCKET_URL,
+    },
+  });
+
+  const message = JSON.stringify({
+    jsonrpc: '2.0',
+    id: 1,
+    method: 'eth_subscribe',
+    params: ['newHeads'],
+  });
+
+  return snap.request({
+    method: 'snap_sendWebSocketMessage',
+    params: { id, message },
+  });
+}
+
+/**
+ * Close a WebSocket connection to a local Ethereum node, if it exists.
+ *
+ * @returns Null.
+ */
+async function unsubscribe() {
+  const sockets = await snap.request({
+    method: 'snap_getWebSockets',
+  });
+
+  if (sockets.length === 0) {
+    return null;
+  }
+
+  await snap.request({
+    method: 'snap_setState',
+    params: {
+      key: 'blockNumber',
+      value: null,
+      encrypted: false,
+    },
+  });
+
+  const socket = sockets[0];
+  return 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.
+ * - `getBlockNumber`: Get the latest block number from the WebSocket connection,
+ * stored in the Snap state.
  *
  * @param params - The request parameters.
  * @param params.request - The JSON-RPC request object.
@@ -50,42 +113,11 @@ export const onRpcRequest: OnRpcRequestHandler = async ({ request }) => {
       return await getJson(params.url);
     }
 
-    case 'startWebSocket': {
-      const id = await snap.request({
-        method: 'snap_openWebSocket',
-        params: {
-          url: WEBSOCKET_URL,
-        },
-      });
-
-      const message = JSON.stringify({
-        jsonrpc: '2.0',
-        id: 1,
-        method: 'eth_subscribe',
-        params: ['newHeads'],
-      });
-
-      return snap.request({
-        method: 'snap_sendWebSocketMessage',
-        params: { id, message },
-      });
-    }
-
-    case 'stopWebSocket': {
-      const sockets = await snap.request({
-        method: 'snap_getWebSockets',
-      });
+    case 'startWebSocket':
+      return subscribe();
 
-      if (sockets.length === 0) {
-        return null;
-      }
-
-      const socket = sockets[0];
-      return snap.request({
-        method: 'snap_closeWebSocket',
-        params: { id: socket.id },
-      });
-    }
+    case 'stopWebSocket':
+      return unsubscribe();
 
     case 'getBlockNumber': {
       return snap.request({