Repository: michalkvasnicak/aws-lambda-graphql
Pull request #70 (2020-02-13 16:59:00)

Pass event and lambdaContext to onConnect function
For a proper authorization flow the context variable passed into the lambda functions are important because they contain the response of the authorization function. The easiest way to solve this with this library is attaching auth context to the connection database entry through the onConnect function. Apparently in the current version this is not possible because the lambda context is not passed into the onConnect function. This pull request fixes this by adding the lambda context as additional variables to the onConnect function.

nenti: > Could you please add a mention about this to `Unreleased` in `CHANGELOG.md` ?

Done.

Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index d59bd9d..f80e5ce 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -9,6 +9,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### [Unreleased](https://github.com/michalkvasnicak/aws-lambda-graphql/compare/aws-lambda-graphql@1.0.0-alpha.3...HEAD)
 
+#### Added
+
+- Added lambda event and context to onConnect function payload to enable adding custom authorizer logic to the connection flow. [@nenti](https://github.com/nenti) see [#70](https://github.com/michalkvasnicak/aws-lambda-graphql/pull/70)
+
 ### [v1.0.0-alpha.3](https://github.com/michalkvasnicak/aws-lambda-graphql/compare/aws-lambda-graphql@1.0.0-alpha.3...HEAD)
 
 #### Fixed
diff --git a/packages/aws-lambda-graphql/README.md b/packages/aws-lambda-graphql/README.md
index 54cf39d..2e007f7 100644
--- a/packages/aws-lambda-graphql/README.md
+++ b/packages/aws-lambda-graphql/README.md
@@ -48,7 +48,7 @@ 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): Promise<boolean|object> | object | boolean`** (`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
   - **`onOperation(message: OperationRequest, params: ExecutionParams, connection: IConnection): Promise<ExecutionParams>|ExecutionParams`** (`optional`)
   - **`onOperationComplete(connection: IConnection, operationId: string): void`** (`optional`)
   - **`onDisconnect(connection: IConnection): void`** (`optional`)
diff --git a/packages/aws-lambda-graphql/src/Server.ts b/packages/aws-lambda-graphql/src/Server.ts
index 4e030cc..581316a 100644
--- a/packages/aws-lambda-graphql/src/Server.ts
+++ b/packages/aws-lambda-graphql/src/Server.ts
@@ -70,6 +70,8 @@ export interface ServerConfig<
     onConnect?: (
       messagePayload: { [key: string]: any } | undefined | null,
       connection: IConnection,
+      event: APIGatewayWebSocketEvent,
+      context: LambdaContext,
     ) =>
       | Promise<boolean | { [key: string]: any }>
       | boolean
@@ -293,7 +295,12 @@ export class Server<
 
               if (onConnect) {
                 try {
-                  const result = await onConnect(operation.payload, connection);
+                  const result = await onConnect(
+                    operation.payload,
+                    connection,
+                    event,
+                    lambdaContext,
+                  );
 
                   if (result === false) {
                     throw new Error('Prohibited connection!');
diff --git a/packages/aws-lambda-graphql/src/formatMessage.ts b/packages/aws-lambda-graphql/src/formatMessage.ts
index 9875d49..d10f3ad 100644
--- a/packages/aws-lambda-graphql/src/formatMessage.ts
+++ b/packages/aws-lambda-graphql/src/formatMessage.ts
@@ -5,6 +5,7 @@ import {
   GQLData,
   GQLComplete,
   GQLStopOperation,
+  GQLConnectionInit,
 } from './protocol';
 
 type AllowedProtocolEvents =
@@ -13,6 +14,7 @@ type AllowedProtocolEvents =
   | GQLErrorEvent
   | GQLData
   | GQLComplete
+  | GQLConnectionInit
   | GQLStopOperation;
 
 export function formatMessage(event: AllowedProtocolEvents): string {


Test patch:
diff --git a/packages/aws-lambda-graphql/src/__tests__/Server.test.ts b/packages/aws-lambda-graphql/src/__tests__/Server.test.ts
index 7b87d54..1f9173e 100644
--- a/packages/aws-lambda-graphql/src/__tests__/Server.test.ts
+++ b/packages/aws-lambda-graphql/src/__tests__/Server.test.ts
@@ -408,24 +408,27 @@ describe('Server', () => {
           {},
         );
 
-        await handlerWithOnConnect(
-          {
-            body: formatMessage({
-              payload: { key: 'value2' },
-              type: CLIENT_EVENT_TYPES.GQL_CONNECTION_INIT,
-            }),
-            requestContext: {
-              connectionId: '1',
-              domainName: 'domain',
-              routeKey: '$default',
-              stage: 'stage',
-            } as any,
+        const event = {
+          body: formatMessage({
+            payload: { key: 'value2' },
+            type: CLIENT_EVENT_TYPES.GQL_CONNECTION_INIT,
+          }),
+          requestContext: {
+            connectionId: '1',
+            domainName: 'domain',
+            routeKey: '$default',
+            stage: 'stage',
           } as any,
-          {} as any,
-        );
+        } as any;
+        await handlerWithOnConnect(event, {} as any);
 
         expect(onConnect).toHaveBeenCalledTimes(1);
-        expect(onConnect).toHaveBeenCalledWith({ key: 'value2' }, {});
+        expect(onConnect).toHaveBeenCalledWith(
+          { key: 'value2' },
+          {},
+          event,
+          {},
+        );
 
         expect(connectionManager.setConnectionData).toHaveBeenCalledTimes(1);
         expect(connectionManager.setConnectionData).toHaveBeenCalledWith(
@@ -503,22 +506,23 @@ describe('Server', () => {
           {},
         );
 
+        // @ts-ignore
+        const testEvent = {
+          body: formatMessage({
+            payload: {},
+            type: CLIENT_EVENT_TYPES.GQL_CONNECTION_INIT,
+          }),
+          requestContext: {
+            connectionId: '1',
+            domainName: 'domain',
+            routeKey: '$default',
+            stage: 'stage',
+          } as any,
+        } as any;
+        const testContext = { testContextVariable: true } as any;
+
         await expect(
-          handlerWithOnConnect(
-            {
-              body: formatMessage({
-                payload: {},
-                type: CLIENT_EVENT_TYPES.GQL_CONNECTION_INIT,
-              }),
-              requestContext: {
-                connectionId: '1',
-                domainName: 'domain',
-                routeKey: '$default',
-                stage: 'stage',
-              } as any,
-            } as any,
-            {} as any,
-          ),
+          handlerWithOnConnect(testEvent, testContext),
         ).resolves.toEqual(
           expect.objectContaining({
             body: formatMessage({
@@ -530,7 +534,7 @@ describe('Server', () => {
         );
 
         expect(onConnect).toHaveBeenCalledTimes(1);
-        expect(onConnect).toHaveBeenCalledWith({}, {});
+        expect(onConnect).toHaveBeenCalledWith({}, {}, testEvent, testContext);
 
         expect(connectionManager.setConnectionData).toHaveBeenCalledTimes(0);