Add websocket handler

package.json

@@ -1,7 +1,7 @@
 {
   "name": "serverless-simple-middleware",
   "description": "Simple middleware to translate the interface of lambda's handler to request => response",
-  "version": "0.0.74",
+  "version": "0.0.75",
   "main": "dist/index.js",
   "types": "dist/index.d.ts",
   "author": "VoyagerX",

src/middleware/buildWebSocket.ts

@@ -0,0 +1,199 @@
+import type {
+  APIGatewayProxyWebsocketEventV2,
+  APIGatewayProxyWebsocketHandlerV2,
+  Context,
+} from 'aws-lambda';
+import { stringifyError } from '../utils';
+import { getLogger } from '../utils/logger';
+import { HandlerPluginBase } from './base';
+import {
+  WebSocketHandler,
+  WebSocketHandlerAuxBase,
+  WebSocketHandlerRequest,
+  WebSocketHandlerResponse,
+} from './websocketBase';
+
+const logger = getLogger(__filename);
+
+type WebSocketDelegator = () => Promise<WebSocketHandlerResponse | undefined>;
+
+class WebSocketHandlerMiddleware<A extends WebSocketHandlerAuxBase> {
+  public auxPromise: Promise<A>;
+  public plugins: Array<HandlerPluginBase<any>>;
+
+  constructor(plugins: Array<HandlerPluginBase<any>>) {
+    this.plugins = plugins;
+    this.auxPromise = this.createAuxPromise();
+  }
+
+  private createAuxPromise = (): Promise<A> => {
+    return !this.plugins || this.plugins.length === 0
+      ? Promise.resolve({} as A) // tslint:disable-line
+      : Promise.all(
+          this.plugins.map((plugin) => {
+            const maybePromise = plugin.create();
+            return maybePromise instanceof Promise
+              ? maybePromise
+              : Promise.resolve(maybePromise);
+          }),
+        ).then(
+          (auxes) =>
+            auxes.reduce((all, each) => ({ ...all, ...each }), {}) as A,
+        );
+  };
+}
+
+class WebSocketHandlerProxy<A extends WebSocketHandlerAuxBase> {
+  private request: WebSocketHandlerRequest;
+  private aux: A;
+  private result: WebSocketHandlerResponse;
+
+  public constructor(event: APIGatewayProxyWebsocketEventV2, context: Context) {
+    logger.stupid(`WebSocket event`, event);
+    this.request = new WebSocketHandlerRequest(event, context);
+    this.aux = {} as A; // tslint:disable-line
+    this.result = { statusCode: 200 };
+  }
+
+  public call = async (
+    middleware: WebSocketHandlerMiddleware<A>,
+    handler: WebSocketHandler<A>,
+  ): Promise<WebSocketHandlerResponse> => {
+    try {
+      this.aux = await middleware.auxPromise;
+    } catch (error) {
+      logger.error(
+        `Error while initializing plugins' aux: ${stringifyError(error)}`,
+      );
+      return {
+        statusCode: 500,
+        body: JSON.stringify(
+          error instanceof Error ? { error: error.message } : error,
+        ),
+      };
+    }
+
+    const actualHandler = [this.generateHandlerDelegator(handler)];
+    const beginHandlers = middleware.plugins.map((plugin) =>
+      this.generatePluginDelegator(plugin.begin),
+    );
+    const endHandlers = middleware.plugins.map((plugin) =>
+      this.generatePluginDelegator(plugin.end),
+    );
+    const errorHandlers = middleware.plugins.map((plugin) =>
+      this.generatePluginDelegator(plugin.error),
+    );
+
+    const iterate = async (handlers: WebSocketDelegator[]) =>
+      Promise.all(handlers.map((each) => this.safeCall(each, errorHandlers)));
+
+    const results = [
+      ...(await iterate(beginHandlers)),
+      ...(await iterate(actualHandler)),
+      ...(await iterate(endHandlers)),
+    ].filter((x) => x);
+
+    // In test phase, throws any exception if there was.
+    if (process.env.NODE_ENV === 'test') {
+      for (const each of results) {
+        if (each instanceof Error) {
+          logger.error(`Error occurred: ${stringifyError(each)}`);
+          throw each;
+        }
+      }
+    }
+
+    results.forEach((result) =>
+      logger.silly(`WebSocket middleware result: ${JSON.stringify(result)}`),
+    );
+
+    return this.result;
+  };
+
+  private safeCall = async (
+    delegator: WebSocketDelegator,
+    errorHandlers: WebSocketDelegator[],
+  ) => {
+    try {
+      const result = await delegator();
+      return result;
+    } catch (error) {
+      const handled = await this.handleError(error, errorHandlers);
+      return handled;
+    }
+  };
+
+  private generateHandlerDelegator =
+    (handler: WebSocketHandler<A>): WebSocketDelegator =>
+    async () => {
+      const maybePromise = handler({
+        request: this.request,
+        response: undefined, // WebSocket doesn't use response
+        aux: this.aux,
+      });
+      const result =
+        maybePromise instanceof Promise ? await maybePromise : maybePromise;
+      logger.stupid(`WebSocket handler result`, result);
+
+      if (result) {
+        this.result = result;
+      }
+      return result;
+    };
+
+  private generatePluginDelegator =
+    (pluginCallback: (context: any) => any): WebSocketDelegator =>
+    async () => {
+      const maybePromise = pluginCallback({
+        request: this.request,
+        response: undefined, // WebSocket doesn't use response (for HTTP plugin compatibility)
+        aux: this.aux,
+      });
+      const result =
+        maybePromise instanceof Promise ? await maybePromise : maybePromise;
+      logger.stupid(`WebSocket plugin callback result`, result);
+      return result;
+    };
+
+  private handleError = async (
+    error: Error,
+    errorHandlers?: WebSocketDelegator[],
+  ) => {
+    logger.error(error);
+    this.request.lastError = error;
+
+    if (errorHandlers) {
+      for (const handler of errorHandlers) {
+        try {
+          await handler();
+        } catch (ignorable) {
+          logger.error(ignorable);
+        }
+      }
+    }
+
+    this.result = {
+      statusCode: 500,
+      body: JSON.stringify(
+        error instanceof Error ? { error: error.message } : error,
+      ),
+    };
+
+    return error;
+  };
+}
+
+const buildWebSocket = <Aux extends WebSocketHandlerAuxBase>(
+  plugins: Array<HandlerPluginBase<any>>,
+) => {
+  const middleware = new WebSocketHandlerMiddleware<Aux>(plugins);
+  return (handler: WebSocketHandler<Aux>): APIGatewayProxyWebsocketHandlerV2 =>
+    async (event: APIGatewayProxyWebsocketEventV2, context: Context) => {
+      return new WebSocketHandlerProxy<Aux>(event, context).call(
+        middleware,
+        handler,
+      );
+    };
+};
+
+export default buildWebSocket;

src/middleware/index.ts

@@ -1,4 +1,5 @@
 import build from './build';
+import buildWebSocket from './buildWebSocket';
 
 import aws from './aws';
 import logger from './logger';
@@ -7,6 +8,7 @@ import trace from './trace';
 
 export const middleware = {
   build,
+  buildWebSocket,
   aws,
   trace,
   logger,
@@ -19,3 +21,4 @@ export * from './database/index';
 export * from './logger';
 export * from './mysql';
 export * from './trace';
+export * from './websocketBase';

src/middleware/websocketBase.ts

@@ -0,0 +1,83 @@
+import type { APIGatewayProxyWebsocketEventV2, Context } from 'aws-lambda';
+import { getLogger } from '../utils/logger';
+
+const logger = getLogger(__filename);
+
+export interface WebSocketHandlerAuxBase {
+  [key: string]: any;
+}
+
+export class WebSocketHandlerRequest {
+  public event: APIGatewayProxyWebsocketEventV2;
+  public context: Context;
+  public lastError: Error | string | undefined;
+
+  private lazyBody?: any;
+
+  constructor(event: APIGatewayProxyWebsocketEventV2, context: Context) {
+    this.event = event;
+    this.context = context;
+    this.lastError = undefined;
+  }
+
+  get body() {
+    if (!this.event.body) {
+      return {};
+    }
+    if (this.lazyBody === undefined) {
+      try {
+        this.lazyBody = JSON.parse(this.event.body);
+      } catch (error) {
+        logger.error(`Failed to parse WebSocket body: ${error}`);
+        this.lazyBody = {};
+      }
+    }
+    return this.lazyBody || {};
+  }
+
+  get connectionId(): string {
+    return this.event.requestContext.connectionId;
+  }
+
+  get routeKey(): string {
+    return this.event.requestContext.routeKey;
+  }
+
+  get domainName(): string {
+    return this.event.requestContext.domainName;
+  }
+
+  get stage(): string {
+    return this.event.requestContext.stage;
+  }
+
+  // HTTP plugin compatibility methods
+
+  /**
+   * For HTTP plugin compatibility (TracerPlugin uses this).
+   * WebSocket events may have headers in $connect route (HTTP handshake),
+   * but typically don't have headers in other routes.
+   */
+  public header(key: string): string | undefined {
+    const event = this.event as any;
+    if (event.headers) {
+      return event.headers[key.toLowerCase()];
+    }
+    return undefined;
+  }
+}
+
+export interface WebSocketHandlerResponse {
+  statusCode: number;
+  body?: string;
+}
+
+export interface WebSocketHandlerContext<A extends WebSocketHandlerAuxBase> {
+  request: WebSocketHandlerRequest;
+  response: undefined; // For HTTP plugin compatibility (not used in WebSocket handlers)
+  aux: A;
+}
+
+export type WebSocketHandler<A extends WebSocketHandlerAuxBase> = (
+  context: WebSocketHandlerContext<A>,
+) => WebSocketHandlerResponse | Promise<WebSocketHandlerResponse> | undefined;