WebSocket Subscriptions via JRPC

src/connection/connection-manager.ts

@@ -23,13 +23,12 @@ export class InMemoryConnectionManager implements ConnectionManager {
   constructor(private dwn: Dwn, private connections: Map<WebSocket, SocketConnection> = new Map()) {}
 
   async connect(socket: WebSocket): Promise<void> {
-    const connection = new SocketConnection(socket, this.dwn);
-    this.connections.set(socket, connection);
-    // attach to the socket's close handler to clean up this connection.
-    socket.on('close', () => {
-      // the connection internally already cleans itself up upon a socket close event, we just ned to remove it from our set.
+    const connection = new SocketConnection(socket, this.dwn, () => {
+      // this is the onClose handler to clean up any closed connections.
       this.connections.delete(socket);
     });
+
+    this.connections.set(socket, connection);
   }
 
   async closeAll(): Promise<void> {

src/connection/socket-connection.ts

@@ -9,7 +9,7 @@ import type { RequestContext } from "../lib/json-rpc-router.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";
+import { jsonRpcRouter } from "../json-rpc-api.js";
 import { JsonRpcErrorCodes, createJsonRpcErrorResponse, createJsonRpcSuccessResponse } from "../lib/json-rpc.js";
 import { DwnServerError, DwnServerErrorCode } from "../dwn-error.js";
 
@@ -26,7 +26,8 @@ export class SocketConnection {
 
   constructor(
     private socket: WebSocket,
-    private dwn: Dwn
+    private dwn: Dwn,
+    private onClose?: () => void
   ){
     socket.on('message', this.message.bind(this));
     socket.on('close', this.close.bind(this));
@@ -93,6 +94,7 @@ export class SocketConnection {
    */
   async close(): Promise<void> {
     clearInterval(this.heartbeatInterval);
+
     // clean up all socket event listeners
     this.socket.removeAllListeners();
 
@@ -107,6 +109,11 @@ export class SocketConnection {
 
     // close the socket.
     this.socket.close();
+
+    // if there was a close handler passed call it after the connection has been closed
+    if (this.onClose !== undefined) {
+      this.onClose();
+    }
   }
 
   /**
@@ -152,7 +159,7 @@ export class SocketConnection {
     };
 
     const requestContext = await this.buildRequestContext(jsonRequest);
-    const { jsonRpcResponse } = await jsonRpcApi.handle(jsonRequest, requestContext);
+    const { jsonRpcResponse } = await jsonRpcRouter.handle(jsonRequest, requestContext);
     if (jsonRpcResponse.error) {
       requestCounter.inc({ method: jsonRequest.method, error: 1 });
     } else {
@@ -189,20 +196,21 @@ export class SocketConnection {
    * Adds a `subscriptionHandler` for `Subscribe` messages.
    */
   private async buildRequestContext(request: JsonRpcRequest): Promise<RequestContext> {
-    const { params, method, subscribe } = request;
+    const { params, method, subscription } = request;
 
     const requestContext: RequestContext = {
       transport        : 'ws',
       dwn              : this.dwn,
       socketConnection : this,
     }
 
-    if (method.startsWith('rpc.subscribe.') && subscribe) {
+    // methods that expect a long-running subscription begin with `rpc.subscribe.`
+    if (method.startsWith('rpc.subscribe.') && subscription) {
       const { message } = params as { message?: GenericMessage };
       if (message?.descriptor.method === DwnMethodName.Subscribe) {
-        const handlerFunc = this.createSubscriptionHandler(subscribe.id);
+        const handlerFunc = this.createSubscriptionHandler(subscription.id);
         requestContext.subscriptionRequest = {
-          id: subscribe.id,
+          id: subscription.id,
           subscriptionHandler: (message): void => handlerFunc(message),
         }
       }

src/http-api.ts

@@ -17,7 +17,7 @@ import { createJsonRpcErrorResponse, JsonRpcErrorCodes } from './lib/json-rpc.js
 import type { DwnServerConfig } from './config.js';
 import { config } from './config.js';
 import { type DwnServerError } from './dwn-error.js';
-import { jsonRpcApi } from './json-rpc-api.js';
+import { jsonRpcRouter } from './json-rpc-api.js';
 import { requestCounter, responseHistogram } from './metrics.js';
 import type { RegistrationManager } from './registration/registration-manager.js';
 
@@ -149,7 +149,7 @@ export class HttpApi {
         transport  : 'http',
         dataStream : requestDataStream,
       };
-      const { jsonRpcResponse, dataStream: responseDataStream } = await jsonRpcApi.handle(dwnRpcRequest, requestContext as RequestContext);
+      const { jsonRpcResponse, dataStream: responseDataStream } = await jsonRpcRouter.handle(dwnRpcRequest, requestContext as RequestContext);
 
       // If the handler catches a thrown exception and returns a JSON RPC InternalError, return the equivalent
       // HTTP 500 Internal Server Error with the response.

src/index.ts

@@ -1,6 +1,6 @@
 export { DwnServerConfig } from './config.js';
 export { DwnServer, DwnServerOptions } from './dwn-server.js';
 export { HttpApi } from './http-api.js';
-export { jsonRpcApi } from './json-rpc-api.js';
+export { jsonRpcRouter } from './json-rpc-api.js';
 export { EStoreType, BackendTypes, StoreType } from './storage.js';
 export { WsApi } from './ws-api.js';

src/json-rpc-api.ts

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

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

@@ -62,6 +62,8 @@ export const handleDwnProcessMessage: JsonRpcHandler = async (
       return { jsonRpcResponse };
     }
 
+    // if this is a subscription request, we first check if the connection has a subscription with this Id
+    // we do this ahead of time to prevent opening a subscription on the dwn only to close it after attempting to add it to the subscription manager
     if (subscriptionRequest !== undefined && socketConnection?.hasSubscription(subscriptionRequest.id)) {
       const jsonRpcResponse = createJsonRpcErrorResponse(
         requestId,
@@ -84,9 +86,9 @@ export const handleDwnProcessMessage: JsonRpcHandler = async (
       delete reply.record.data; // not serializable via JSON
     }
 
-    // Subscribe messages return a close function to facilitate closing the subscription
     if (subscriptionRequest && reply.subscription) {
       const { close } = reply.subscription;
+      // Subscribe messages return a close function to facilitate closing the 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 = {
@@ -111,8 +113,8 @@ export const handleDwnProcessMessage: JsonRpcHandler = async (
       e.message,
     );
 
-    // log the error response
-    log.error('handleDwnProcessMessage error', jsonRpcResponse);
+    // log the unhandled error response
+    log.error('handleDwnProcessMessage error', jsonRpcResponse, e);
     return { jsonRpcResponse } as HandlerResponse;
   }
 };

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

@@ -16,27 +16,27 @@ import {
 /**
  * Closes a subscription tied to a specific `SocketConnection`.
  *
- * @param dwnRequest must include JsonRpcId of the subscription request within the `params`.
+ * @param jsonRpcRequest must include JsonRpcId of the subscription request within a `subscription object`.
  * @param context must include the associated `SocketConnection`.
  *
  */
 export const handleSubscriptionsClose: JsonRpcHandler = async (
-  dwnRequest,
+  jsonRpcRequest,
   context,
 ) => {
-  const requestId = dwnRequest.id ?? uuidv4();
+  const requestId = jsonRpcRequest.id ?? uuidv4();
   if (context.socketConnection === undefined) {
     const jsonRpcResponse = createJsonRpcErrorResponse(requestId, JsonRpcErrorCodes.InvalidRequest, 'socket connection does not exist');
     return { jsonRpcResponse };
   }
 
-  if (dwnRequest.subscribe === undefined) {
+  if (jsonRpcRequest.subscription === undefined) {
     const jsonRpcResponse = createJsonRpcErrorResponse(requestId, JsonRpcErrorCodes.InvalidRequest, 'subscribe options do not exist');
     return { jsonRpcResponse };
   }
 
   const { socketConnection } = context;
-  const { id } = dwnRequest.subscribe as { id: JsonRpcId };
+  const { id } = jsonRpcRequest.subscription as { id: JsonRpcId };
 
   let jsonRpcResponse:JsonRpcResponse;
   try {

src/json-rpc-socket.ts

@@ -3,7 +3,7 @@ import { v4 as uuidv4 } from 'uuid';
 import WebSocket from 'ws';
 
 import type { JsonRpcId, JsonRpcRequest, JsonRpcResponse } from "./lib/json-rpc.js";
-import { createJsonRpcSubscribeRequest } from "./lib/json-rpc.js";
+import { createJsonRpcSubscriptionRequest } from "./lib/json-rpc.js";
 
 // These were arbitrarily chosen, but can be modified via connect options
 const CONNECT_TIMEOUT = 3_000;
@@ -18,6 +18,8 @@ export interface JsonRpcSocketOptions {
   onclose?: () => void;
   /** optional socket error handler */
   onerror?: (error?: any) => void;
+  /** optional already connected socket to inject */
+  socket?: WebSocket;
 }
 
 /**
@@ -30,21 +32,22 @@ export class JsonRpcSocket {
     const { connectTimeout = CONNECT_TIMEOUT, responseTimeout = RESPONSE_TIMEOUT, onclose, onerror } = options;
 
     const socket = new WebSocket(url);
-    if (onclose === undefined) {
+
+    socket.onclose = onclose;
+    socket.onerror = onerror;
+
+    if (!socket.onclose) {
       socket.onclose = ():void => {
         log.info(`JSON RPC Socket close ${url}`);
       }
     }
 
-    if (onerror === undefined) {
+    if (!socket.onerror) {
       socket.onerror = (error?: any):void => {
         log.error(`JSON RPC Socket error ${url}`, error);
       }
     }
 
-    socket.onclose = onclose;
-    socket.onerror = onerror;
-
     return new Promise<JsonRpcSocket>((resolve, reject) => {
       socket.on('open', () => {
         resolve(new JsonRpcSocket(socket, responseTimeout));
@@ -98,33 +101,33 @@ export class JsonRpcSocket {
       throw new Error('subscribe rpc requests must include the `rpc.subscribe` prefix');
     }
 
-    if (!request.subscribe) {
+    if (!request.subscription) {
       throw new Error('subscribe rpc requests must include subscribe options');
     }
 
-    const subscriptionId = request.subscribe.id;
-    const messageHandler = (event: { data: any }):void => {
+    const subscriptionId = request.subscription.id;
+    const socketEventListener = (event: { data: any }):void => {
       const jsonRpcResponse = JSON.parse(event.data.toString()) as 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.socket.removeEventListener('message', socketEventListener);
           this.closeSubscription(subscriptionId);
         }
         listener(jsonRpcResponse);
       }
     };
-    this.socket.addEventListener('message', messageHandler);
+    this.socket.addEventListener('message', socketEventListener);
 
     const response = await this.request(request);
     if (response.error) {
-      this.socket.removeEventListener('message', messageHandler);
+      this.socket.removeEventListener('message', socketEventListener);
       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);
+      this.socket.removeEventListener('message', socketEventListener);
       await this.closeSubscription(subscriptionId);
     }
 
@@ -136,7 +139,7 @@ export class JsonRpcSocket {
 
   private closeSubscription(id: JsonRpcId): Promise<JsonRpcResponse> {
     const requestId = uuidv4();
-    const request = createJsonRpcSubscribeRequest(requestId, 'rpc.subscribe.close', {}, id);
+    const request = createJsonRpcSubscriptionRequest(requestId, 'rpc.subscribe.close', {}, id);
     return this.request(request);
   }
 

src/lib/json-rpc-router.ts

@@ -8,6 +8,7 @@ import type { SocketConnection } from '../connection/socket-connection.js';
 export type RequestContext = {
   transport: 'http' | 'ws';
   dwn: Dwn;
+  /** the socket connection associated with this request if over sockets */
   socketConnection?: SocketConnection;
   subscriptionRequest?: {
     /** The JsonRpcId of the subscription handler */

src/lib/json-rpc.ts

@@ -1,5 +1,3 @@
-import { v4 as uuidv4 } from 'uuid';
-
 export type JsonRpcId = string | number | null;
 export type JsonRpcParams = any;
 export type JsonRpcVersion = '2.0';
@@ -9,8 +7,8 @@ export interface JsonRpcRequest {
   id?: JsonRpcId;
   method: string;
   params?: JsonRpcParams;
-  /** JSON RPC Subscribe Extension Parameters */
-  subscribe?: { 
+  /** JSON RPC Subscription Extension Parameters */
+  subscription?: {
     id: JsonRpcId
   };
 }
@@ -87,7 +85,7 @@ export const createJsonRpcNotification = (
   };
 };
 
-export const createJsonRpcSubscribeRequest = (
+export const createJsonRpcSubscriptionRequest = (
   id: JsonRpcId,
   method: string,
   params?: JsonRpcParams,
@@ -98,8 +96,8 @@ export const createJsonRpcSubscribeRequest = (
     id,
     method,
     params,
-    subscribe: {
-      id: subscriptionId ?? uuidv4(),
+    subscription: {
+      id: subscriptionId,
     }
   }
 }