Refactor to use a notification handler function

Author: rekmarks

src/JsonRpcEngine.test.ts

@@ -28,13 +28,13 @@ describe('JsonRpcEngine', () => {
 
   it('handle: returns error for invalid request method', async () => {
     const engine = new JsonRpcEngine();
-    let response: any = await engine.handle({ method: null } as any);
+    let response: any = await engine.handle({ id: 1, method: null } as any);
     expect(response.error.code).toStrictEqual(-32600);
     expect(response.result).toBeUndefined();
 
+    // No response if duck-typed as a notification
     response = await engine.handle({ method: true } as any);
-    expect(response.error.code).toStrictEqual(-32600);
-    expect(response.result).toBeUndefined();
+    expect(response).toBeUndefined();
   });
 
   it('handle: basic middleware test 1', async () => {

src/JsonRpcEngine.ts

@@ -1,6 +1,6 @@
 import SafeEventEmitter from '@metamask/safe-event-emitter';
 import { errorCodes, EthereumRpcError, serializeError } from 'eth-rpc-errors';
-import { isJsonRpcNotification, isJsonRpcRequest } from './utils';
+import { isJsonRpcRequest } from './utils';
 
 type Maybe<T> = Partial<T> | null | undefined;
 
@@ -86,58 +86,32 @@ export type JsonRpcMiddleware<T, U> = (
   end: JsonRpcEngineEndCallback,
 ) => void;
 
-export type JsonRpcNotificationMiddleware<T> = (
-  req: JsonRpcNotification<T>,
-) => Promise<void>;
+export type JsonRpcNotificationHandler<T> = (
+  notification: JsonRpcNotification<T>,
+) => void | Promise<void>;
 
 /**
  * A JSON-RPC request and response processor.
  * Give it a stack of middleware, pass it requests, and get back responses.
  */
 export class JsonRpcEngine extends SafeEventEmitter {
-  private _requestMiddleware: JsonRpcMiddleware<unknown, unknown>[];
+  private _middleware: JsonRpcMiddleware<unknown, unknown>[];
 
-  private _notificationMiddleware: JsonRpcNotificationMiddleware<unknown>[];
+  private readonly _notificationHandler?: JsonRpcNotificationHandler<unknown>;
 
-  constructor() {
+  constructor(notificationHandler?: JsonRpcNotificationHandler<unknown>) {
     super();
-    this._requestMiddleware = [];
-    this._notificationMiddleware = [];
+    this._middleware = [];
+    this._notificationHandler = notificationHandler;
   }
 
   /**
    * Add a middleware function to the engine's request middleware stack.
    *
    * @param middleware - The middleware function to add.
    */
-  pushRequestMiddleware<T, U>(middleware: JsonRpcMiddleware<T, U>): void {
-    this._requestMiddleware.push(
-      middleware as JsonRpcMiddleware<unknown, unknown>,
-    );
-  }
-
-  /**
-   * Alias for `pushRequestMiddleware`, for backwards compatibility.
-   *
-   * @param middleware - The request middleware function to add.
-   */
   push<T, U>(middleware: JsonRpcMiddleware<T, U>): void {
-    this.pushRequestMiddleware(
-      middleware as JsonRpcMiddleware<unknown, unknown>,
-    );
-  }
-
-  /**
-   * Adds a middleware function to the engine's notification middleware stack.
-   *
-   * @param middleware - The notification middleware function to add.
-   */
-  pushNotificationMiddleware<T>(
-    middleware: JsonRpcNotificationMiddleware<T>,
-  ): void {
-    this._requestMiddleware.push(
-      middleware as JsonRpcNotificationMiddleware<unknown>,
-    );
+    this._middleware.push(middleware as JsonRpcMiddleware<unknown, unknown>);
   }
 
   /**
@@ -158,7 +132,7 @@ export class JsonRpcEngine extends SafeEventEmitter {
    * @param callback - An error-first callback that will receive a `void` response.
    */
   handle<T>(
-    request: JsonRpcNotification<T>,
+    notification: JsonRpcNotification<T>,
     callback: (error: unknown, response: void) => void,
   ): void;
 
@@ -228,21 +202,8 @@ export class JsonRpcEngine extends SafeEventEmitter {
   asMiddleware(): JsonRpcMiddleware<unknown, unknown> {
     return async (req, res, next, end) => {
       try {
-        if (isJsonRpcNotification(req)) {
-          // This will invoke the end callback.
-          return await JsonRpcEngine._processNotification(
-            req,
-            this._notificationMiddleware,
-            () => end(),
-          );
-        }
-
         const [middlewareError, isComplete, returnHandlers] =
-          await JsonRpcEngine._runAllRequestMiddleware(
-            req,
-            res,
-            this._requestMiddleware,
-          );
+          await JsonRpcEngine._runAllMiddleware(req, res, this._middleware);
 
         if (isComplete) {
           await JsonRpcEngine._runReturnHandlers(returnHandlers);
@@ -367,7 +328,7 @@ export class JsonRpcEngine extends SafeEventEmitter {
 
       if (isJsonRpcRequest(callerReq)) {
         return callback(error, {
-          id: (callerReq as any).id ?? null,
+          id: callerReq.id ?? null,
           jsonrpc: '2.0',
           error,
         });
@@ -380,12 +341,8 @@ export class JsonRpcEngine extends SafeEventEmitter {
     // We can't use isJsonRpcNotification here because that narrows callerReq to
     // "never" after the if clause for unknown reasons.
     if (!isJsonRpcRequest(callerReq)) {
-      // This will invoke the callback.
-      return await JsonRpcEngine._processNotification(
-        { ...callerReq },
-        this._notificationMiddleware,
-        callback,
-      );
+      await this._notificationHandler?.(callerReq);
+      return callback(null);
     }
 
     let error: JsonRpcEngineCallbackError = null;
@@ -398,7 +355,7 @@ export class JsonRpcEngine extends SafeEventEmitter {
     };
 
     try {
-      await JsonRpcEngine._processRequest(req, res, this._requestMiddleware);
+      await JsonRpcEngine._processRequest(req, res, this._middleware);
     } catch (_error) {
       // A request handler error, a re-thrown middleware error, or something
       // unexpected.
@@ -416,52 +373,22 @@ export class JsonRpcEngine extends SafeEventEmitter {
     return callback(error, res as JsonRpcResponse<unknown>);
   }
 
-  /**
-   * Runs all notification middleware on the specified notification.
-   * Does not throw.
-   *
-   * @param notification - The notification object to process.
-   * @param notificationMiddlewares - The stack of notification middleware
-   * functions.
-   * @param callback - Callback from {@link JsonRpcEngine._handle} or
-   * {@link JsonRpcEngineEndCallback}. Will always be invoked with nothing,
-   * because notifications must not be responded to.
-   */
-  private static async _processNotification(
-    notification: JsonRpcNotification<unknown>,
-    notificationMiddlewares: JsonRpcNotificationMiddleware<unknown>[],
-    callback: (error?: unknown, response?: JsonRpcResponse<unknown>) => void,
-  ): Promise<void> {
-    try {
-      for (const middleware of notificationMiddlewares) {
-        await middleware(notification);
-      }
-    } catch (middlewareError) {
-      console.error(middlewareError);
-    }
-    callback(null);
-  }
-
   /**
    * For the given request and response, runs all middleware and their return
    * handlers, if any, and ensures that internal request processing semantics
    * are satisfied.
    *
    * @param req - The request object.
    * @param res - The response object.
-   * @param requestMiddlewares - The stack of request middleware functions.
+   * @param middlewares - The stack of request middleware functions.
    */
   private static async _processRequest(
     req: JsonRpcRequest<unknown>,
     res: PendingJsonRpcResponse<unknown>,
-    requestMiddlewares: JsonRpcMiddleware<unknown, unknown>[],
+    middlewares: JsonRpcMiddleware<unknown, unknown>[],
   ): Promise<void> {
     const [error, isComplete, returnHandlers] =
-      await JsonRpcEngine._runAllRequestMiddleware(
-        req,
-        res,
-        requestMiddlewares,
-      );
+      await JsonRpcEngine._runAllMiddleware(req, res, middlewares);
 
     // Throw if "end" was not called, or if the response has neither a result
     // nor an error.
@@ -483,15 +410,15 @@ export class JsonRpcEngine extends SafeEventEmitter {
    *
    * @param req - The request object.
    * @param res - The response object.
-   * @param requestMiddlewares - The stack of middleware functions to execute.
+   * @param middlewares - The stack of middleware functions to execute.
    * @returns An array of any error encountered during middleware execution,
    * a boolean indicating whether the request was completed, and an array of
    * middleware-defined return handlers.
    */
-  private static async _runAllRequestMiddleware(
+  private static async _runAllMiddleware(
     req: JsonRpcRequest<unknown>,
     res: PendingJsonRpcResponse<unknown>,
-    requestMiddlewares: JsonRpcMiddleware<unknown, unknown>[],
+    middlewares: JsonRpcMiddleware<unknown, unknown>[],
   ): Promise<
     [
       unknown, // error
@@ -504,8 +431,8 @@ export class JsonRpcEngine extends SafeEventEmitter {
     let isComplete = false;
 
     // Go down stack of middleware, call and collect optional returnHandlers
-    for (const middleware of requestMiddlewares) {
-      [error, isComplete] = await JsonRpcEngine._runRequestMiddleware(
+    for (const middleware of middlewares) {
+      [error, isComplete] = await JsonRpcEngine._runMiddleware(
         req,
         res,
         middleware,
@@ -529,7 +456,7 @@ export class JsonRpcEngine extends SafeEventEmitter {
    * @returns An array of any error encountered during middleware exection,
    * and a boolean indicating whether the request should end.
    */
-  private static _runRequestMiddleware(
+  private static _runMiddleware(
     req: JsonRpcRequest<unknown>,
     res: PendingJsonRpcResponse<unknown>,
     middleware: JsonRpcMiddleware<unknown, unknown>,
@@ -602,7 +529,7 @@ export class JsonRpcEngine extends SafeEventEmitter {
    *
    * @param req - The request object.
    * @param res - The response object.
-   * @param isComplete - Boolean from {@link JsonRpcEngine._runAllRequestMiddleware}
+   * @param isComplete - Boolean from {@link JsonRpcEngine._runAllMiddleware}
    * indicating whether a middleware ended the request.
    */
   private static _checkForCompletion(