@enbox/dwn-server 0.0.2 → 0.0.3

package/src/connection/socket-connection.ts

@@ -1,54 +1,63 @@
-import type { Dwn, GenericMessage, MessageEvent } from "@enbox/dwn-sdk-js";
-import { DwnMethodName } from "@enbox/dwn-sdk-js";
+import { DwnMethodName } from '@enbox/dwn-sdk-js';
+import type { RequestContext } from '../lib/json-rpc-router.js';
+import type { ServerWebSocket } from 'bun';
+import type { Dwn, GenericMessage, MessageEvent } from '@enbox/dwn-sdk-js';
+import type { JsonRpcErrorResponse, JsonRpcId, JsonRpcRequest, JsonRpcResponse, JsonRpcSubscription } from '../lib/json-rpc.js';
+
+import type { WsData } from '../http-api.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, JsonRpcSubscription } from "../lib/json-rpc.js";
-
-import { requestCounter } from "../metrics.js";
-import { jsonRpcRouter } from "../json-rpc-api.js";
-import { JsonRpcErrorCodes, createJsonRpcErrorResponse, createJsonRpcSuccessResponse } from "../lib/json-rpc.js";
-import { DwnServerError, DwnServerErrorCode } from "../dwn-error.js";
+import { jsonRpcRouter } from '../json-rpc-api.js';
+import { requestCounter } from '../metrics.js';
+import {
+  createJsonRpcErrorResponse,
+  createJsonRpcSuccessResponse,
+  JsonRpcErrorCodes,
+} from '../lib/json-rpc.js';
+import { DwnServerError, DwnServerErrorCode } from '../dwn-error.js';
 
 const HEARTBEAT_INTERVAL = 30_000;
 
 /**
  * SocketConnection handles a WebSocket connection to a DWN using JSON RPC.
  * It also manages references to the long running RPC subscriptions for the connection.
+ *
+ * With Bun's native WebSocket, the message/close/error events are dispatched by the
+ * Bun.serve() websocket handlers in http-api.ts, which delegate to the public `message()`
+ * and `close()` methods on this class.
  */
 export class SocketConnection {
-  private heartbeatInterval: NodeJS.Timer;
+  private heartbeatInterval: ReturnType<typeof setInterval>;
   private subscriptions: Map<JsonRpcId, JsonRpcSubscription> = new Map();
   private isAlive: boolean;
 
   constructor(
-    private socket: WebSocket,
+    private socket: ServerWebSocket<WsData>,
     private dwn: Dwn,
-    private onClose?: () => void
+    private onCloseCallback?: () => void
   ){
-    socket.on('message', this.message.bind(this));
-    socket.on('close', this.close.bind(this));
-    socket.on('error', this.error.bind(this));
-    socket.on('pong', this.pong.bind(this));
-
-    // Sometimes connections between client <-> server can get borked in such a way that
-    // leaves both unaware of the borkage. ping messages can be used as a means to verify
-    // that the remote endpoint is still responsive. Server will ping each socket every 30s
-    // if a pong hasn't received from a socket by the next ping, the server will terminate
-    // the socket connection
+    // Bun handles ping/pong automatically at the protocol level, but we still
+    // want an application-level heartbeat to detect dead connections.
     this.isAlive = true;
     this.heartbeatInterval = setInterval(() => {
       if (this.isAlive === false) {
         this.close();
+        return;
       }
       this.isAlive = false;
       this.socket.ping();
     }, HEARTBEAT_INTERVAL);
   }
 
+  /**
+   * Called when a pong is received (triggered by Bun's built-in ping/pong handling).
+   */
+  pong(): void {
+    this.isAlive = true;
+  }
+
   /**
    * Checks to see if the incoming `JsonRpcId` is already in use for a subscription.
    */
@@ -65,7 +74,7 @@ export class SocketConnection {
       throw new DwnServerError(
         DwnServerErrorCode.ConnectionSubscriptionJsonRpcIdExists,
         `the subscription with id ${subscription.id} already exists`
-      )
+      );
     }
 
     this.subscriptions.set(subscription.id, subscription);
@@ -81,7 +90,7 @@ export class SocketConnection {
       throw new DwnServerError(
         DwnServerErrorCode.ConnectionSubscriptionJsonRpcIdNotFound,
         `the subscription with id ${id} was not found`
-      )
+      );
     }
 
     const connection = this.subscriptions.get(id);
@@ -95,10 +104,7 @@ export class SocketConnection {
   async close(): Promise<void> {
     clearInterval(this.heartbeatInterval);
 
-    // clean up all socket event listeners
-    this.socket.removeAllListeners();
-
-    const closePromises = [];
+    const closePromises: Promise<void>[] = [];
     for (const [id, subscription] of this.subscriptions) {
       closePromises.push(subscription.close());
       this.subscriptions.delete(id);
@@ -111,52 +117,45 @@ export class SocketConnection {
     this.socket.close();
 
     // if there was a close handler passed call it after the connection has been closed
-    if (this.onClose !== undefined) {
-      this.onClose();
+    if (this.onCloseCallback !== undefined) {
+      this.onCloseCallback();
     }
   }
 
-  /**
-   * Pong messages are automatically sent in response to ping messages as required by
-   * the websocket spec. So, no need to send explicit pongs.
-   */
-  private pong(): void {
-    this.isAlive = true;
-  }
-
   /**
    * Log the error and close the connection.
    */
-  private async error(error:Error): Promise<void>{
+  async error(error: Error): Promise<void> {
     log.error(`SocketConnection error, terminating connection`, error);
-    this.socket.terminate();
+    this.socket.close();
     await this.close();
   }
 
   /**
    * Handles a `JSON RPC 2.0` encoded message.
+   * This is called by Bun's websocket message handler via http-api.ts.
    */
-  private async message(dataBuffer: Buffer): Promise<void> {
+  async message(dataBuffer: Buffer): Promise<void> {
     const requestData = dataBuffer.toString();
     if (!requestData) {
       return this.send(createJsonRpcErrorResponse(
         uuidv4(),
         JsonRpcErrorCodes.BadRequest,
         'request payload required.'
-      ))
+      ));
     }
 
     let jsonRequest: JsonRpcRequest;
     try {
       jsonRequest = JSON.parse(requestData);
-    } catch(error) {
+    } catch (error) {
       const errorResponse = createJsonRpcErrorResponse(
         uuidv4(),
         JsonRpcErrorCodes.BadRequest,
         (error as Error).message
       );
       return this.send(errorResponse);
-    };
+    }
 
     const requestContext = await this.buildRequestContext(jsonRequest);
     const { jsonRpcResponse } = await jsonRpcRouter.handle(jsonRequest, requestContext);
@@ -164,15 +163,15 @@ export class SocketConnection {
       requestCounter.inc({ method: jsonRequest.method, error: 1 });
     } else {
       requestCounter.inc({
-        method: jsonRequest.method,
-        status: jsonRpcResponse?.result?.reply?.status?.code || 0,
+        method : jsonRequest.method,
+        status : jsonRpcResponse?.result?.reply?.status?.code || 0,
       });
     }
     this.send(jsonRpcResponse);
   }
 
   /**
-   * Sends a JSON encoded Buffer through the Websocket.
+   * Sends a JSON encoded string through the WebSocket.
    */
   private send(response: JsonRpcResponse | JsonRpcErrorResponse): void {
     this.socket.send(JSON.stringify(response));
@@ -180,20 +179,16 @@ export class SocketConnection {
 
   /**
    * 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: MessageEvent) => void {
     return (event) => {
       const response = createJsonRpcSuccessResponse(id, { event });
       this.send(response);
-    }
+    };
   }
 
   /**
    * Builds a `RequestContext` object to use with the `JSON RPC API`.
-   *
-   * Adds a `subscriptionHandler` for `Subscribe` messages.
    */
   private async buildRequestContext(request: JsonRpcRequest): Promise<RequestContext> {
     const { params, method, subscription } = request;
@@ -202,7 +197,7 @@ export class SocketConnection {
       transport        : 'ws',
       dwn              : this.dwn,
       socketConnection : this,
-    }
+    };
 
     // methods that expect a long-running subscription begin with `rpc.subscribe.`
     if (method.startsWith('rpc.subscribe.') && subscription) {
@@ -210,9 +205,9 @@ export class SocketConnection {
       if (message?.descriptor.method === DwnMethodName.Subscribe) {
         const handlerFunc = this.createSubscriptionHandler(subscription.id);
         requestContext.subscriptionRequest = {
-          id: subscription.id,
-          subscriptionHandler: (message): void => handlerFunc(message),
-        }
+          id                  : subscription.id,
+          subscriptionHandler : (message): void => handlerFunc(message),
+        };
       }
     }
 

package/src/dwn-error.ts

@@ -16,8 +16,8 @@ export class DwnServerError extends Error {
    */
   public toJSON(): { code: string, message: string } {
     return {
-      code: this.code,
-      message: this.message,
+      code    : this.code,
+      message : this.message,
     };
   }
 }

package/src/dwn-server.ts

@@ -1,20 +1,21 @@
 import type { DidResolver } from '@enbox/dids';
+import type { DwnServerConfig } from './config.js';
 import type { EventStream } from '@enbox/dwn-sdk-js';
 import type { ProcessHandlers } from './process-handlers.js';
-import type { Server } from 'http';
-import type { WebSocketServer } from 'ws';
-import type { DwnServerConfig } from './config.js';
+import type { Server } from 'bun';
+
+import type { WsData } from './http-api.js';
 
 import log from 'loglevel';
 import prefix from 'loglevel-plugin-prefix';
+import { Dwn, EventEmitterStream } from '@enbox/dwn-sdk-js';
+
 import { config as defaultConfig } from './config.js';
 import { getDwnConfig } from './storage.js';
-import { HttpServerShutdownHandler } from './lib/http-server-shutdown-handler.js';
 import { HttpApi } from './http-api.js';
 import { PluginLoader } from './plugin-loader.js';
 import { RegistrationManager } from './registration/registration-manager.js';
 import { WsApi } from './ws-api.js';
-import { Dwn, EventEmitterStream } from '@enbox/dwn-sdk-js';
 import { removeProcessHandlers, setProcessHandlers } from './process-handlers.js';
 
 /**
@@ -43,7 +44,7 @@ enum DwnServerState {
 export class DwnServer {
   serverState = DwnServerState.Stopped;
   processHandlers: ProcessHandlers;
-  
+
   /**
    * A custom DID resolver to use in the DWN.
    * Mainly for testing purposes. Ignored if `dwn` is provided.
@@ -51,7 +52,6 @@ export class DwnServer {
   didResolver?: DidResolver;
   dwn?: Dwn;
   config: DwnServerConfig;
-  #httpServerShutdownHandler: HttpServerShutdownHandler;
   #httpApi: HttpApi;
   #wsApi: WsApi;
 
@@ -93,10 +93,10 @@ export class DwnServer {
     if (!this.dwn) {
       // undefined registrationStoreUrl is used as a signal that there is no need for tenant registration, DWN is open for all.
       registrationManager = await RegistrationManager.create({
-        registrationStoreUrl: this.config.registrationStoreUrl,
-        termsOfServiceFilePath: this.config.termsOfServiceFilePath,
-        proofOfWorkChallengeNonceSeed: this.config.registrationProofOfWorkSeed,
-        proofOfWorkInitialMaximumAllowedHash: this.config.registrationProofOfWorkInitialMaxHash,
+        registrationStoreUrl                 : this.config.registrationStoreUrl,
+        termsOfServiceFilePath               : this.config.termsOfServiceFilePath,
+        proofOfWorkChallengeNonceSeed        : this.config.registrationProofOfWorkSeed,
+        proofOfWorkInitialMaximumAllowedHash : this.config.registrationProofOfWorkInitialMaxHash,
       });
 
       let eventStream: EventStream | undefined;
@@ -111,10 +111,10 @@ export class DwnServer {
       }
 
       const dwnConfig = await getDwnConfig(this.config, {
-        didResolver: this.didResolver,
-        tenantGate: registrationManager,
+        didResolver : this.didResolver,
+        tenantGate  : registrationManager,
         eventStream,
-      })
+      });
       this.dwn = await Dwn.create(dwnConfig);
     }
 
@@ -123,12 +123,8 @@ export class DwnServer {
     await this.#httpApi.start(this.config.port);
     log.info(`HttpServer listening on port ${this.config.port}`);
 
-    this.#httpServerShutdownHandler = new HttpServerShutdownHandler(
-      this.#httpApi.server,
-    );
-
     if (this.config.webSocketSupport) {
-      this.#wsApi = new WsApi(this.#httpApi.server, this.dwn);
+      this.#wsApi = new WsApi(this.#httpApi, this.dwn);
       this.#wsApi.start();
       log.info('WebSocketServer ready...');
     }
@@ -143,32 +139,23 @@ export class DwnServer {
     }
 
     await this.dwn.close();
-    await this.#httpApi.close();
 
     // close WebSocket server if it was initialized
     if (this.#wsApi !== undefined) {
       await this.#wsApi.close();
     }
 
-    await new Promise<void>((resolve) => {
-      this.#httpServerShutdownHandler.stop(() => {
-        resolve();
-      });
-    });
+    await this.#httpApi.close();
 
     removeProcessHandlers(this.processHandlers);
 
     this.serverState = DwnServerState.Stopped;
   }
 
-  get httpServer(): Server {
+  get httpServer(): Server<WsData> {
     return this.#httpApi.server;
   }
 
-  get wsServer(): WebSocketServer | undefined {
-    return this.#wsApi?.server;
-  }
-
   /**
    * Gets the RegistrationManager for testing purposes.
    */