Updated middleware execution to provide new functionality (#91)

* Updated middleware execution to provide new functionality

* Add missing headers to the changeset

* Remove context modification from result middleware

* Update from review

* Added new README sections and changelog message

* Spelling

* Prettier

Author: BlenderDude

.changeset/giant-rats-rhyme.md

@@ -0,0 +1,36 @@
+---
+'@as-integrations/aws-lambda': minor
+---
+
+## Short circuit middleware execution
+
+You can now opt to return a Lambda result object directly from the middleware. This will cancel the middleware chain, bypass GraphQL request processing, and immediately return the Lambda result.
+
+Example
+
+```ts
+export const handler = startServerAndCreateLambdaHandler(
+  server,
+  handlers.createAPIGatewayProxyEventV2RequestHandler(),
+  {
+    context: async () => {
+      return {};
+    },
+    middleware: [
+      async (event) => {
+        const psk = Buffer.from('SuperSecretPSK');
+        const token = Buffer.from(event.headers['X-Auth-Token']);
+        if (
+          psk.byteLength !== token.byteLength ||
+          crypto.timingSafeEqual(psk, token)
+        ) {
+          return {
+            statusCode: '403',
+            body: 'Forbidden',
+          };
+        }
+      },
+    ],
+  },
+);
+```

README.md

@@ -22,7 +22,7 @@ import {
 // The GraphQL schema
 const typeDefs = `#graphql
   type Query {
-    hello: String
+    hello: String!
   }
 `;
 
@@ -45,6 +45,59 @@ export default startServerAndCreateLambdaHandler(
 );
 ```
 
+## Context
+
+As with all Apollo Server 4 integrations, the context resolution is done in the integration. For the Lambda integration, it will look like the following:
+
+```ts
+import { ApolloServer } from '@apollo/server';
+import {
+  startServerAndCreateLambdaHandler,
+  handlers,
+} from '@as-integrations/aws-lambda';
+
+type ContextValue = {
+  isAuthenticated: boolean;
+};
+
+// The GraphQL schema
+const typeDefs = `#graphql
+  type Query {
+    hello: String!
+    isAuthenticated: Boolean!
+  }
+`;
+
+// Set up Apollo Server
+const server = new ApolloServer<ContextValue>({
+  typeDefs,
+  resolvers: {
+    Query: {
+      hello: () => 'world',
+      isAuthenticated: (root, args, context) => {
+        // For context typing to be valid one of the following must be implemented
+        //   1. `resolvers` defined inline in the server config (not particularly scalable, but works)
+        //   2. Add the type in the resolver function. ex. `(root, args, context: ContextValue)`
+        //   3. Propagate the type from an outside definition like GraphQL Codegen
+        return context.isAuthenticated;
+      },
+    },
+  },
+});
+
+export default startServerAndCreateLambdaHandler(
+  server,
+  handlers.createAPIGatewayProxyEventV2RequestHandler({
+    context: async ({ event }) => {
+      // Do some parsing on the event (parse JWT, cookie, auth header, etc.)
+      return {
+        isAuthenticated: true,
+      };
+    },
+  }),
+);
+```
+
 ## Middleware
 
 For mutating the event before passing off to `@apollo/server` or mutating the result right before returning, middleware can be utilized.
@@ -83,6 +136,8 @@ export default startServerAndCreateLambdaHandler(
 );
 ```
 
+### Middleware Typing
+
 If you want to define strictly typed middleware outside of the middleware array, the easiest way would be to extract your request handler into a variable and utilize the `typeof` keyword from Typescript. You could also manually use the `RequestHandler` type and fill in the event and result values yourself.
 
 ```typescript
@@ -133,6 +188,47 @@ export default startServerAndCreateLambdaHandler(server, requestHandler, {
 });
 ```
 
+### Middleware Short Circuit
+
+In some situations, a middleware function might require the execution end before reaching Apollo Server. This might be a global auth guard or session token lookup.
+
+To achieve this, the request middleware function accepts `ResultType` or `Promise<ResultType>` as a return type. Should middleware resolve to such a value, that result is returned and no further execution occurs.
+
+```typescript
+import {
+  startServerAndCreateLambdaHandler,
+  middleware,
+  handlers,
+} from '@as-integrations/aws-lambda';
+import type {
+  APIGatewayProxyEventV2,
+  APIGatewayProxyStructuredResultV2,
+} from 'aws-lambda';
+import { server } from './server';
+
+const requestHandler = handlers.createAPIGatewayProxyEventV2RequestHandler();
+
+// Utilizing typeof
+const sessionMiddleware: middleware.MiddlewareFn<typeof requestHandler> = (
+  event,
+) => {
+  // ... check session
+  if (!event.headers['X-Session-Key']) {
+    // If header doesn't exist, return early
+    return {
+      statusCode: 401
+      body: 'Unauthorized'
+    }
+  }
+};
+
+export default startServerAndCreateLambdaHandler(server, requestHandler, {
+  middleware: [
+    sessionMiddleware,
+  ],
+});
+```
+
 ## Event Extensions
 
 Each of the provided request handler factories has a generic for you to pass a manually extended event type if you have custom authorizers, or if the event type you need has a generic you must pass yourself. For example, here is a request that allows access to the lambda authorizer:

cspell-dict.txt

@@ -12,4 +12,5 @@ unawaited
 vendia
 withrequired
 typeof
-instanceof
+instanceof
+codegen

src/__tests__/middleware.test.ts

@@ -1,35 +1,54 @@
 import { ApolloServer } from '@apollo/server';
 import type { APIGatewayProxyEventV2 } from 'aws-lambda';
 import { handlers, startServerAndCreateLambdaHandler } from '..';
+import gql from 'graphql-tag';
+import { type DocumentNode, print } from 'graphql';
 
-const event: APIGatewayProxyEventV2 = {
-  version: '2',
-  headers: {
-    'content-type': 'application/json',
-  },
-  isBase64Encoded: false,
-  rawQueryString: '',
-  requestContext: {
-    http: {
-      method: 'POST',
+function createEvent(doc: DocumentNode): APIGatewayProxyEventV2 {
+  return {
+    version: '2',
+    headers: {
+      'content-type': 'application/json',
     },
-    // Other requestContext properties omitted for brevity
-  } as any,
-  rawPath: '/',
-  routeKey: '/',
-  body: '{"operationName": null, "variables": null, "query": "{ hello }"}',
-};
+    isBase64Encoded: false,
+    rawQueryString: '',
+    requestContext: {
+      http: {
+        method: 'POST',
+      },
+      // Other requestContext properties omitted for brevity
+    } as any,
+    rawPath: '/',
+    routeKey: '/',
+    body: JSON.stringify({
+      query: print(doc),
+    }),
+  };
+}
 
 const typeDefs = `#graphql
   type Query {
     hello: String
   }
+  type Mutation {
+    mutateContext: String
+  }
 `;
 
 const resolvers = {
   Query: {
     hello: () => 'world',
   },
+  Mutation: {
+    mutateContext: async (
+      _root: any,
+      _args: any,
+      context: { foo: string | null },
+    ) => {
+      context.foo = 'bar';
+      return 'ok';
+    },
+  },
 };
 
 const server = new ApolloServer({
@@ -39,13 +58,23 @@ const server = new ApolloServer({
 
 describe('Request mutation', () => {
   it('updates incoming event headers', async () => {
+    const event = createEvent(gql`
+      query {
+        hello
+      }
+    `);
     const headerAdditions = {
       'x-injected-header': 'foo',
     };
     const lambdaHandler = startServerAndCreateLambdaHandler(
       server,
       handlers.createAPIGatewayProxyEventV2RequestHandler(),
       {
+        context: async () => {
+          return {
+            foo: null,
+          };
+        },
         middleware: [
           async (event) => {
             Object.assign(event.headers, headerAdditions);
@@ -58,15 +87,52 @@ describe('Request mutation', () => {
       expect(event.headers[key]).toBe(value);
     }
   });
+  it('returns early if middleware returns a result', async () => {
+    const event = createEvent(gql`
+      query {
+        hello
+      }
+    `);
+    const lambdaHandler = startServerAndCreateLambdaHandler(
+      server,
+      handlers.createAPIGatewayProxyEventV2RequestHandler(),
+      {
+        context: async () => {
+          return {
+            foo: null,
+          };
+        },
+        middleware: [
+          async () => {
+            return {
+              statusCode: 418,
+            };
+          },
+        ],
+      },
+    );
+    const result = await lambdaHandler(event, {} as any, () => {})!;
+    expect(result.statusCode).toBe(418);
+  });
 });
 
 describe('Response mutation', () => {
   it('adds cookie values to emitted result', async () => {
+    const event = createEvent(gql`
+      query {
+        hello
+      }
+    `);
     const cookieValue = 'foo=bar';
     const lambdaHandler = startServerAndCreateLambdaHandler(
       server,
       handlers.createAPIGatewayProxyEventV2RequestHandler(),
       {
+        context: async () => {
+          return {
+            foo: null,
+          };
+        },
         middleware: [
           async () => {
             return async (result) => {

src/lambdaHandler.ts

@@ -75,17 +75,30 @@ export function startServerAndCreateLambdaHandler<
       [];
     try {
       for (const middlewareFn of options?.middleware ?? []) {
-        const resultCallback = await middlewareFn(event);
-        if (resultCallback) {
-          resultMiddlewareFns.push(resultCallback);
+        const middlewareReturnValue = await middlewareFn(event);
+        // If the middleware returns an object, we assume it's a LambdaResponse
+        if (
+          typeof middlewareReturnValue === 'object' &&
+          middlewareReturnValue !== null
+        ) {
+          return middlewareReturnValue;
+        }
+        // If the middleware returns a function, we assume it's a result callback
+        if (middlewareReturnValue) {
+          resultMiddlewareFns.push(middlewareReturnValue);
         }
       }
 
       const httpGraphQLRequest = handler.fromEvent(event);
 
       const response = await server.executeHTTPGraphQLRequest({
         httpGraphQLRequest,
-        context: () => contextFunction({ event, context }),
+        context: () => {
+          return contextFunction({
+            event,
+            context,
+          });
+        },
       });
 
       const result = handler.toSuccessResult(response);

src/middleware.ts

@@ -4,7 +4,7 @@ export type LambdaResponse<ResultType> = (result: ResultType) => Promise<void>;
 
 export type LambdaRequest<EventType, ResultType> = (
   event: EventType,
-) => Promise<LambdaResponse<ResultType> | void>;
+) => Promise<LambdaResponse<ResultType> | ResultType | void>;
 
 export type MiddlewareFn<RH extends RequestHandler<any, any>> =
   RH extends RequestHandler<infer EventType, infer ResultType>