feat(event-handler): implement mechanism to manipulate response in middleware

Author: svozza

packages/event-handler/src/rest/BaseRouter.ts

@@ -9,6 +9,7 @@ import type {
   ErrorConstructor,
   ErrorHandler,
   ErrorResolveOptions,
+  HandlerOptions,
   HttpMethod,
   Middleware,
   Path,
@@ -19,6 +20,7 @@ import type {
 import { HttpErrorCodes, HttpVerbs } from './constants.js';
 import {
   handlerResultToProxyResult,
+  handlerResultToResponse,
   proxyEventToWebRequest,
   responseToProxyResult,
 } from './converters.js';
@@ -209,6 +211,15 @@ abstract class BaseRouter {
 
     const request = proxyEventToWebRequest(event);
 
+    const handlerOptions: HandlerOptions = {
+      event,
+      context,
+      request,
+      // this response should be overwritten by the handler, if it isn't
+      // it means somthing went wrong with the middleware chain
+      res: new Response('', { status: 500 }),
+    };
+
     try {
       const path = new URL(request.url).pathname as Path;
 
@@ -223,34 +234,38 @@ abstract class BaseRouter {
           ? route.handler.bind(options.scope)
           : route.handler;
 
+      const handlerMiddleware: Middleware = async (params, options, next) => {
+        const handlerResult = await handler(params, options);
+        const handlerResponse = handlerResultToResponse(
+          handlerResult,
+          options.res.headers
+        );
+
+        options.res = handlerResponse;
+
+        await next();
+      };
+
       const middleware = composeMiddleware([
         ...this.middleware,
         ...route.middleware,
+        handlerMiddleware,
       ]);
 
-      const result = await middleware(
+      const middlewareResult = await middleware(
         route.params,
-        {
-          event,
-          context,
-          request,
-        },
-        () => handler(route.params, { event, context, request })
+        handlerOptions,
+        () => Promise.resolve()
       );
 
-      // In practice this we never happen because the final 'middleware' is
-      // the handler function that allways returns HandlerResponse. However, the
-      // type signature of of NextFunction includes undefined so we need this for
-      // the TS compiler
-      if (result === undefined) throw new InternalServerError();
+      // middleware result takes precedence to allow short-circuiting
+      const result = middlewareResult ?? handlerOptions.res;
 
-      return await handlerResultToProxyResult(result);
+      return handlerResultToProxyResult(result);
     } catch (error) {
       this.logger.debug(`There was an error processing the request: ${error}`);
       const result = await this.handleError(error as Error, {
-        request,
-        event,
-        context,
+        ...handlerOptions,
         scope: options?.scope,
       });
       return await responseToProxyResult(result);
@@ -281,13 +296,10 @@ abstract class BaseRouter {
     const handler = this.errorHandlerRegistry.resolve(error);
     if (handler !== null) {
       try {
-        const body = await handler.apply(options.scope ?? this, [
+        const { scope, ...handlerOptions } = options;
+        const body = await handler.apply(scope ?? this, [
           error,
-          {
-            request: options.request,
-            event: options.event,
-            context: options.context,
-          },
+          handlerOptions,
         ]);
         return new Response(JSON.stringify(body), {
           status: body.statusCode,

packages/event-handler/src/rest/converters.ts

@@ -89,13 +89,62 @@ export const responseToProxyResult = async (
     }
   }
 
-  return {
+  const result: APIGatewayProxyResult = {
     statusCode: response.status,
     headers,
-    multiValueHeaders,
     body: await response.text(),
     isBase64Encoded: false,
   };
+
+  if (Object.keys(multiValueHeaders).length > 0) {
+    result.multiValueHeaders = multiValueHeaders;
+  }
+
+  return result;
+};
+
+/**
+ * Converts a handler response to a Web API Response object.
+ * Handles APIGatewayProxyResult, Response objects, and plain objects.
+ *
+ * @param response - The handler response (APIGatewayProxyResult, Response, or plain object)
+ * @param headers - Optional headers to be included in the response
+ * @returns A Web API Response object
+ */
+export const handlerResultToResponse = (
+  response: HandlerResponse,
+  headers: Headers = new Headers()
+): Response => {
+  if (response instanceof Response) {
+    return response;
+  }
+  if (isAPIGatewayProxyResult(response)) {
+    const responseHeaders = new Headers();
+
+    for (const [key, value] of Object.entries(response.headers ?? {})) {
+      if (value != null) {
+        responseHeaders.set(key, String(value));
+      }
+    }
+
+    for (const [key, values] of Object.entries(
+      response.multiValueHeaders ?? {}
+    )) {
+      for (const value of values ?? []) {
+        responseHeaders.append(key, String(value));
+      }
+    }
+
+    return new Response(response.body, {
+      status: response.statusCode,
+      headers: responseHeaders,
+    });
+  }
+  headers.set('Content-Type', 'application/json');
+  return new Response(JSON.stringify(response), {
+    status: 200,
+    headers,
+  });
 };
 
 /**
@@ -117,7 +166,7 @@ export const handlerResultToProxyResult = async (
   return {
     statusCode: 200,
     body: JSON.stringify(response),
-    headers: { 'Content-Type': 'application/json' },
+    headers: { 'content-type': 'application/json' },
     isBase64Encoded: false,
   };
 };

packages/event-handler/src/rest/utils.ts

@@ -2,11 +2,11 @@ import { isRecord, isString } from '@aws-lambda-powertools/commons/typeutils';
 import type { APIGatewayProxyEvent, APIGatewayProxyResult } from 'aws-lambda';
 import type {
   CompiledRoute,
+  HandlerOptions,
   HandlerResponse,
   HttpMethod,
   Middleware,
   Path,
-  RequestOptions,
   ValidationResult,
 } from '../types/rest.js';
 import {
@@ -146,7 +146,7 @@ export const isAPIGatewayProxyResult = (
 export const composeMiddleware = (middleware: Middleware[]): Middleware => {
   return async (
     params: Record<string, string>,
-    options: RequestOptions,
+    options: HandlerOptions,
     next: () => Promise<HandlerResponse | void>
   ): Promise<HandlerResponse | void> => {
     let index = -1;

packages/event-handler/src/types/rest.ts

@@ -2,7 +2,11 @@ import type {
   GenericLogger,
   JSONObject,
 } from '@aws-lambda-powertools/commons/types';
-import type { APIGatewayProxyEvent, Context } from 'aws-lambda';
+import type {
+  APIGatewayProxyEvent,
+  APIGatewayProxyResult,
+  Context,
+} from 'aws-lambda';
 import type { BaseRouter } from '../rest/BaseRouter.js';
 import type { HttpErrorCodes, HttpVerbs } from '../rest/constants.js';
 import type { Route } from '../rest/Route.js';
@@ -14,17 +18,18 @@ type ErrorResponse = {
   message: string;
 };
 
-type RequestOptions = {
+type HandlerOptions = {
   request: Request;
   event: APIGatewayProxyEvent;
   context: Context;
+  res: Response;
 };
 
-type ErrorResolveOptions = RequestOptions & ResolveOptions;
+type ErrorResolveOptions = HandlerOptions & ResolveOptions;
 
 type ErrorHandler<T extends Error = Error> = (
   error: T,
-  options: RequestOptions
+  options: HandlerOptions
 ) => Promise<ErrorResponse>;
 
 interface ErrorConstructor<T extends Error = Error> {
@@ -58,7 +63,7 @@ type HandlerResponse = Response | JSONObject;
 type RouteHandler<
   TParams = Record<string, unknown>,
   TReturn = HandlerResponse,
-> = (args: TParams, options: RequestOptions) => Promise<TReturn>;
+> = (args: TParams, options: HandlerOptions) => Promise<TReturn>;
 
 type HttpMethod = keyof typeof HttpVerbs;
 
@@ -83,7 +88,7 @@ type NextFunction = () => Promise<HandlerResponse | void>;
 
 type Middleware = (
   params: Record<string, string>,
-  options: RequestOptions,
+  options: HandlerOptions,
   next: NextFunction
 ) => Promise<void | HandlerResponse>;
 
@@ -123,7 +128,7 @@ export type {
   HttpMethod,
   Middleware,
   Path,
-  RequestOptions,
+  HandlerOptions,
   RouterOptions,
   RouteHandler,
   RouteOptions,