add onWebsocketConnect handler

Author: jcane86

packages/aws-lambda-graphql/README.md

@@ -48,7 +48,8 @@ All options from Apollo Lambda Server and
 - **onError** (`(err: any) => void`, `optional`) - use to log errors from websocket handler on unknown error
 - **subscriptionManager** (`ISubscriptionManager`, `required`)
 - **subscriptions** (`optional`)
-  - **`onConnect(messagePayload: object, connection: IConnection, event: APIGatewayWebSocketEvent, context: LambdaContext): Promise<boolean|object> | object | boolean`** (`optional`) - Return an object to set a context to your connection object saved in the database e.g. for saving authentication details
+  - **`onWebsocketConnect(connection: IConnection, event: APIGatewayWebSocketEvent, context: LambdaContext): Promise<boolean|object> | object | boolean`** (`optional`) - onWebsocketConnect is called when the Websocket connection is initialized ($connect route).  Return an object to set a context to your connection object saved in the database e.g. for saving authentication details. This is especially useful to get authentication details (API GW authorizers only run in $connect route)
+  - **`onConnect(messagePayload: object, connection: IConnection, event: APIGatewayWebSocketEvent, context: LambdaContext): Promise<boolean|object> | object | boolean`** (`optional`) - onConnect is called when the GraphQL connection is initialized (connection_init message). Return an object to set a context to your connection object saved in the database e.g. for saving authentication details. NOTE: This is not the websocket $connect route, see onWebsocketConnect for the $connect route.
   - **`onOperation(message: OperationRequest, params: ExecutionParams, connection: IConnection): Promise<ExecutionParams>|ExecutionParams`** (`optional`)
   - **`onOperationComplete(connection: IConnection, operationId: string): void`** (`optional`)
   - **`onDisconnect(connection: IConnection): void`** (`optional`)

packages/aws-lambda-graphql/src/Server.ts

@@ -67,6 +67,27 @@ export interface ServerConfig<
       connection: IConnection,
       operationId: string,
     ) => void;
+  /**
+     * onWebsocketConnect is called when the Websocket connection is initialized ($connect route).
+     * Return an object to set a context to your connection object saved in the database e.g. for saving authentication details.
+     * This is especially useful to get authentication details (API GW authorizers only run in $connect route)
+     * 
+     */
+    onWebsocketConnect?: (
+      connection: IConnection,
+      event: APIGatewayWebSocketEvent,
+      context: LambdaContext,
+    ) =>
+      | Promise<boolean | { [key: string]: any }>
+      | boolean
+      | { [key: string]: any };
+    /**
+     * onConnect is called when the GraphQL connection is initialized (connection_init message).
+     * Return an object to set a context to your connection object saved in the database e.g. for saving authentication details.
+     * 
+     * NOTE: This is not the websocket $connect route, see onWebsocketConnect for the $connect route
+     *
+     */
     onConnect?: (
       messagePayload: { [key: string]: any } | undefined | null,
       connection: IConnection,
@@ -227,16 +248,65 @@ export class Server<
         // based on routeKey, do actions
         switch (event.requestContext.routeKey) {
           case '$connect': {
+            const {
+              onWebsocketConnect,
+            } = this.subscriptionOptions || {};
+
             // register connection
             // if error is thrown during registration, connection is rejected
             // we can implement some sort of authorization here
             const endpoint = extractEndpointFromEvent(event);
 
-            await this.connectionManager.registerConnection({
+            const connection = await this.connectionManager.registerConnection({
               endpoint,
               connectionId: event.requestContext.connectionId,
             });
 
+            let newConnectionContext = {}
+
+            if (onWebsocketConnect) {
+              try {
+                const result = await onWebsocketConnect(
+                  connection,
+                  event,
+                  lambdaContext,
+                );
+
+                if (result === false) {
+                  throw new Error('Prohibited connection!');
+                } else if (result !== null && typeof result === 'object') {
+                  newConnectionContext = result
+                }
+              } catch (err) {
+                const errorResponse = formatMessage({
+                  type: SERVER_EVENT_TYPES.GQL_ERROR,
+                  payload: { message: err.message },
+                });
+
+                await this.connectionManager.sendToConnection(
+                  connection,
+                  errorResponse,
+                );
+                await this.connectionManager.closeConnection(connection);
+
+                return {
+                  body: errorResponse,
+                  statusCode: 401,
+                };
+              }
+            }
+
+            // set connection context which will be available during graphql execution
+            const connectionData = {
+              ...connection.data,
+              context: newConnectionContext
+            };
+
+            await this.connectionManager.setConnectionData(
+              connectionData,
+              connection,
+            );
+
             return {
               body: '',
               headers: event.headers?.['Sec-WebSocket-Protocol']?.includes(
@@ -336,7 +406,7 @@ export class Server<
               // set connection context which will be available during graphql execution
               const connectionData = {
                 ...connection.data,
-                context: newConnectionContext,
+                context: {...connection.data.context, ...newConnectionContext},
                 isInitialized: true,
               };