WebSocket Subscriptions via JRPC

src/connection/socket-connection.ts

@@ -1,12 +1,12 @@
-import type { Dwn, GenericMessage } from "@tbd54566975/dwn-sdk-js";
+import type { Dwn, GenericMessage, MessageEvent } from "@tbd54566975/dwn-sdk-js";
 import { DwnMethodName } from "@tbd54566975/dwn-sdk-js";
 
 import type { WebSocket } from "ws";
 import log from 'loglevel';
 import { v4 as uuidv4 } from 'uuid';
 
 import type { RequestContext } from "../lib/json-rpc-router.js";
-import type { JsonRpcErrorResponse, JsonRpcId, JsonRpcRequest, JsonRpcResponse } from "../lib/json-rpc.js";
+import type { JsonRpcErrorResponse, JsonRpcId, JsonRpcRequest, JsonRpcResponse, JsonRpcSubscription } from "../lib/json-rpc.js";
 
 import { requestCounter } from "../metrics.js";
 import { jsonRpcApi } from "../json-rpc-api.js";
@@ -15,12 +15,6 @@ import { DwnServerError, DwnServerErrorCode } from "../dwn-error.js";
 
 const HEARTBEAT_INTERVAL = 30_000;
 
-export interface JsonRpcSubscription {
-  /** JSON RPC Id of the Subscription Request */
-  id: JsonRpcId;
-  close: () => Promise<void>;
-}
-
 /**
  * SocketConnection handles a WebSocket connection to a DWN using JSON RPC.
  * It also manages references to the long running RPC subscriptions for the connection.
@@ -54,6 +48,13 @@ export class SocketConnection {
     }, HEARTBEAT_INTERVAL);
   }
 
+  /**
+   * Checks to see if the incoming `JsonRpcId` is already in use for a subscription.
+   */
+  hasSubscription(id: JsonRpcId): boolean {
+    return this.subscriptions.has(id);
+  }
+
   /**
    * Adds a reference for the JSON RPC Subscription to this connection.
    * Used for cleanup if the connection is closed.
@@ -147,7 +148,6 @@ export class SocketConnection {
         JsonRpcErrorCodes.BadRequest,
         (error as Error).message
       );
-
       return this.send(errorResponse);
     };
 
@@ -165,26 +165,20 @@ export class SocketConnection {
   }
 
   /**
-   * Sends a JSON encoded Buffer through the Websocket. Accepts a callback, if none is provided an error logger is used.
+   * Sends a JSON encoded Buffer through the Websocket.
    */
-  private send(response: JsonRpcResponse | JsonRpcErrorResponse, cb?: (error?: Error) => void): void {
-    if (!cb) {
-      cb = (error):void => {
-        if(error) { log.error('socket send error', error, response); }
-      }
-    }
-
-    this.socket.send(Buffer.from(JSON.stringify(response)), cb);
+  private send(response: JsonRpcResponse | JsonRpcErrorResponse): void {
+    this.socket.send(Buffer.from(JSON.stringify(response)));
   }
 
   /**
    * Creates a subscription handler to send messages matching the subscription requested.
    *
    * Wraps the incoming `message` in a `JSON RPC Success Response` using the original subscription`JSON RPC Id` to send through the WebSocket.
    */
-  private createSubscriptionHandler(id: JsonRpcId): (message: GenericMessage) => void {
+  private createSubscriptionHandler(id: JsonRpcId): (message: MessageEvent) => void {
     return (event) => {
-      const response = createJsonRpcSuccessResponse(id, { reply: { event } });
+      const response = createJsonRpcSuccessResponse(id, { event });
       this.send(response);
     }
   }
@@ -195,17 +189,22 @@ export class SocketConnection {
    * Adds a `subscriptionHandler` for `Subscribe` messages.
    */
   private async buildRequestContext(request: JsonRpcRequest): Promise<RequestContext> {
-    const { id, params, method } = request;
+    const { params, method, subscribe } = request;
+
     const requestContext: RequestContext = {
       transport        : 'ws',
       dwn              : this.dwn,
       socketConnection : this,
     }
 
-    if (method === 'dwn.processMessage') {
-      const { message } = params as { message: GenericMessage };
-      if (message.descriptor.method === DwnMethodName.Subscribe) {
-        requestContext.subscriptionHandler = this.createSubscriptionHandler(id).bind(this);
+    if (method.startsWith('rpc.subscribe.') && subscribe) {
+      const { message } = params as { message?: GenericMessage };
+      if (message?.descriptor.method === DwnMethodName.Subscribe) {
+        const handlerFunc = this.createSubscriptionHandler(subscribe.id);
+        requestContext.subscriptionRequest = {
+          id: subscribe.id,
+          subscriptionHandler: (message): void => handlerFunc(message),
+        }
       }
     }
 

src/json-rpc-api.ts

@@ -6,4 +6,6 @@ import { handleSubscriptionsClose } from './json-rpc-handlers/subscription/index
 export const jsonRpcApi = new JsonRpcRouter();
 
 jsonRpcApi.on('dwn.processMessage', handleDwnProcessMessage);
-jsonRpcApi.on('subscription.close', handleSubscriptionsClose);
+jsonRpcApi.on('rpc.subscribe.dwn.processMessage', handleDwnProcessMessage);
+
+jsonRpcApi.on('rpc.subscribe.close', handleSubscriptionsClose);

src/json-rpc-handlers/dwn/process-message.ts

@@ -2,26 +2,27 @@ import type { GenericMessage } from '@tbd54566975/dwn-sdk-js';
 import { DwnInterfaceName, DwnMethodName } from '@tbd54566975/dwn-sdk-js';
 
 import type { Readable as IsomorphicReadable } from 'readable-stream';
+import log from 'loglevel';
 import { v4 as uuidv4 } from 'uuid';
 
+import type { JsonRpcSubscription } from '../../lib/json-rpc.js';
 import type {
   HandlerResponse,
   JsonRpcHandler,
 } from '../../lib/json-rpc-router.js';
 
-import { DwnServerErrorCode } from '../../dwn-error.js';
 import {
   createJsonRpcErrorResponse,
   createJsonRpcSuccessResponse,
   JsonRpcErrorCodes,
 } from '../../lib/json-rpc.js';
-import log from 'loglevel';
+
 
 export const handleDwnProcessMessage: JsonRpcHandler = async (
   dwnRequest,
   context,
 ) => {
-  const { dwn, dataStream, subscriptionHandler, socketConnection, transport } = context;
+  const { dwn, dataStream, subscriptionRequest, socketConnection, transport } = context;
   const { target, message } = dwnRequest.params as { target: string, message: GenericMessage };
   const requestId = dwnRequest.id ?? uuidv4();
 
@@ -41,23 +42,41 @@ export const handleDwnProcessMessage: JsonRpcHandler = async (
       return { jsonRpcResponse };
     }
 
+    // subscribe methods must come with a subscriptionRequest context 
+    if (message.descriptor.method === DwnMethodName.Subscribe && subscriptionRequest === undefined) {
+      const jsonRpcResponse = createJsonRpcErrorResponse(
+        requestId,
+        JsonRpcErrorCodes.InvalidRequest,
+        `subscribe methods must contain a subscriptionRequest context`
+      );
+      return { jsonRpcResponse };
+    }
+
     // Subscribe methods are only supported on 'ws' (WebSockets)
-    if (transport !== 'ws' && message.descriptor.method === DwnMethodName.Subscribe) {
+    if (transport !== 'ws' && subscriptionRequest !== undefined) {
       const jsonRpcResponse = createJsonRpcErrorResponse(
         requestId,
         JsonRpcErrorCodes.InvalidParams,
-        `Subscribe not supported via ${context.transport}`
+        `subscriptions are not supported via ${context.transport}`
+      )
+      return { jsonRpcResponse };
+    }
+
+    if (subscriptionRequest !== undefined && socketConnection?.hasSubscription(subscriptionRequest.id)) {
+      const jsonRpcResponse = createJsonRpcErrorResponse(
+        requestId,
+        JsonRpcErrorCodes.InvalidParams,
+        `the subscribe id: ${subscriptionRequest.id} is in use by an active subscription`
       )
       return { jsonRpcResponse };
     }
 
     const reply = await dwn.processMessage(target, message, {
       dataStream: dataStream as IsomorphicReadable,
-      subscriptionHandler,
+      subscriptionHandler: subscriptionRequest?.subscriptionHandler,
     });
 
-    const { record, subscription } = reply;
-
+    const { record } = reply;
     // RecordsRead messages return record data as a stream to for accommodate large amounts of data
     let recordDataStream: IsomorphicReadable;
     if (record !== undefined && record.data !== undefined) {
@@ -66,29 +85,16 @@ export const handleDwnProcessMessage: JsonRpcHandler = async (
     }
 
     // Subscribe messages return a close function to facilitate closing the subscription
-    if (subscription !== undefined) {
-      const { close } = subscription;
-      try {
-        // adding a reference to the close function for this subscription request to the connection.
-        // this will facilitate closing the subscription later.
-        await socketConnection.addSubscription({ id: requestId, close });
-        delete reply.subscription.close // not serializable via JSON
-      } catch(error) {
-        // close the subscription upon receiving an error here
-        await close();
-        if (error.code === DwnServerErrorCode.ConnectionSubscriptionJsonRpcIdExists) {
-          // a subscription with this request id already exists
-          const errorResponse = createJsonRpcErrorResponse(
-            requestId,
-            JsonRpcErrorCodes.BadRequest,
-            `the request id ${requestId} already has an active subscription`
-          );
-          return { jsonRpcResponse: errorResponse };
-        } else {
-          // will catch as an unknown error below
-          throw new Error('unknown error adding subscription');
-        }
+    if (subscriptionRequest && reply.subscription) {
+      const { close } = reply.subscription;
+      // we add a reference to the close function for this subscription request to the socket connection.
+      // this will facilitate closing the subscription later.
+      const subscriptionReply: JsonRpcSubscription = {
+        id: subscriptionRequest.id,
+        close,
       }
+      await socketConnection.addSubscription(subscriptionReply);
+      delete reply.subscription.close // delete the close method from the reply as it's not JSON serializable and has a held reference.
     }
 
     const jsonRpcResponse = createJsonRpcSuccessResponse(requestId, { reply });
@@ -107,7 +113,6 @@ export const handleDwnProcessMessage: JsonRpcHandler = async (
 
     // log the error response
     log.error('handleDwnProcessMessage error', jsonRpcResponse);
-
     return { jsonRpcResponse } as HandlerResponse;
   }
 };

src/json-rpc-handlers/subscription/close.ts

@@ -25,8 +25,18 @@ export const handleSubscriptionsClose: JsonRpcHandler = async (
   context,
 ) => {
   const requestId = dwnRequest.id ?? uuidv4();
+  if (context.socketConnection === undefined) {
+    const jsonRpcResponse = createJsonRpcErrorResponse(requestId, JsonRpcErrorCodes.InvalidRequest, 'socket connection does not exist');
+    return { jsonRpcResponse };
+  }
+
+  if (dwnRequest.subscribe === undefined) {
+    const jsonRpcResponse = createJsonRpcErrorResponse(requestId, JsonRpcErrorCodes.InvalidRequest, 'subscribe options do not exist');
+    return { jsonRpcResponse };
+  }
+
   const { socketConnection } = context;
-  const { id } = dwnRequest.params as { id: JsonRpcId};
+  const { id } = dwnRequest.subscribe as { id: JsonRpcId };
 
   let jsonRpcResponse:JsonRpcResponse;
   try {

src/json-rpc-socket.ts

@@ -2,7 +2,8 @@ import log from 'loglevel';
 import { v4 as uuidv4 } from 'uuid';
 import WebSocket from 'ws';
 
-import type { JsonRpcRequest, JsonRpcResponse } from "./lib/json-rpc.js";
+import type { JsonRpcId, JsonRpcRequest, JsonRpcResponse } from "./lib/json-rpc.js";
+import { createJsonRpcSubscribeRequest } from "./lib/json-rpc.js";
 
 // These were arbitrarily chosen, but can be modified via connect options
 const CONNECT_TIMEOUT = 3_000;
@@ -72,7 +73,6 @@ export class JsonRpcSocket {
           return resolve(jsonRpsResponse);
         }
       };
-
       // subscribe to the listener before sending the request
       this.socket.addEventListener('message', handleResponse);
       this.send(request);
@@ -89,33 +89,61 @@ export class JsonRpcSocket {
    * Sends a JSON-RPC request through the socket and keeps a listener open to read associated responses as they arrive.
    * Returns a close method to clean up the listener.
    */
-  subscribe(request: JsonRpcRequest, listener: (response: JsonRpcResponse) => void): { close: () => void } {
-    request.id ??= uuidv4();
+  async subscribe(request: JsonRpcRequest, listener: (response: JsonRpcResponse) => void): Promise<{
+    response: JsonRpcResponse;
+    close?: () => Promise<void>;
+   }> {
+
+    if (!request.method.startsWith('rpc.subscribe.')) {
+      throw new Error('subscribe rpc requests must include the `rpc.subscribe` prefix');
+    }
+
+    if (!request.subscribe) {
+      throw new Error('subscribe rpc requests must include subscribe options');
+    }
+
+    const subscriptionId = request.subscribe.id;
     const messageHandler = (event: { data: any }):void => {
       const jsonRpcResponse = JSON.parse(event.data.toString()) as JsonRpcResponse;
-      if (jsonRpcResponse.id === request.id) {
-        // if the incoming response id matches the request id, trigger the listener
-        return listener(jsonRpcResponse);
+      if (jsonRpcResponse.id === subscriptionId) {
+        if (jsonRpcResponse.error !== undefined) {
+          // remove the event listener upon receipt of a JSON RPC Error.
+          this.socket.removeEventListener('message', messageHandler);
+          this.closeSubscription(subscriptionId);
+        }
+        listener(jsonRpcResponse);
       }
     };
-
-    // subscribe to the listener before sending the request
     this.socket.addEventListener('message', messageHandler);
-    this.send(request);
+
+    const response = await this.request(request);
+    if (response.error) {
+      this.socket.removeEventListener('message', messageHandler);
+      return { response }
+    }
+
+    // clean up listener and create a `rpc.subscribe.close` message to use when closing this JSON RPC subscription
+    const close = async (): Promise<void> => {
+      this.socket.removeEventListener('message', messageHandler);
+      await this.closeSubscription(subscriptionId);
+    }
 
     return {
-      close: ():void => {
-        // removes the listener for this particular request
-        this.socket.removeEventListener('message', messageHandler);
-      }
-    };
+      response,
+      close
+    }
+  }
+
+  private closeSubscription(id: JsonRpcId): Promise<JsonRpcResponse> {
+    const requestId = uuidv4();
+    const request = createJsonRpcSubscribeRequest(requestId, 'rpc.subscribe.close', {}, id);
+    return this.request(request);
   }
 
   /**
    * Sends a JSON-RPC request through the socket. You must subscribe to a message listener separately to capture the response.
    */
   send(request: JsonRpcRequest):void {
     this.socket.send(Buffer.from(JSON.stringify(request)));
-    return;
   }
 }

src/lib/json-rpc-router.ts

@@ -1,16 +1,20 @@
-import type { Dwn, GenericMessage } from '@tbd54566975/dwn-sdk-js';
+import type { Dwn, EventSubscriptionHandler } from '@tbd54566975/dwn-sdk-js';
 
 import type { Readable } from 'node:stream';
 
-import type { JsonRpcRequest, JsonRpcResponse } from './json-rpc.js';
+import type { JsonRpcId, JsonRpcRequest, JsonRpcResponse } from './json-rpc.js';
 import type { SocketConnection } from '../connection/socket-connection.js';
 
 export type RequestContext = {
   transport: 'http' | 'ws';
   dwn: Dwn;
   socketConnection?: SocketConnection;
-  /** The `GenericMessage` handler associated with a subscription request, only used in `ws` requests */
-  subscriptionHandler?: (message: GenericMessage) => void;
+  subscriptionRequest?: {
+    /** The JsonRpcId of the subscription handler */
+    id: JsonRpcId; 
+    /** The `MessageEvent` handler associated with a subscription request, only used in `ws` requests */
+    subscriptionHandler: EventSubscriptionHandler;
+  }
   /** The `Readable` stream associated with a `RecordsWrite` request only used in `http` requests */
   dataStream?: Readable;
 };