Merge pull request #2 from epilot-dev/chore/force-custom-error-message

chore: introduce custom error/message when header not present
epilot-dev · Feb 6, 2024 · 0fa1c76 · 0fa1c76
2 parents ae1d8d5 + f1edefa
commit 0fa1c76
Show file tree

Hide file tree

Showing 3 changed files with 113 additions and 9 deletions.
diff --git a/packages/large-response-middleware/README.md b/packages/large-response-middleware/README.md
@@ -7,7 +7,7 @@ Enables Lambdas to return responses larger than 6MB by offloading the content to
 - This implementation currently provides support for API Gateway with Lambda Proxy Integration only.
 - There are plans to extend this work as described here [#issue-1](https://github.com/epilot-dev/aws-lambda-utility-middlewares/issues/1)
 
-----
+---
 
 <p align="center">
   <img src="https://raw.githubusercontent.com/epilot-dev/aws-lambda-utility-middlewares/main/packages/large-response-middleware/docs/out/architecture-1/Architecture%20-%20Sequence%20Diagram.svg" />
@@ -23,18 +23,19 @@ When a client can handle a Large Response, it must send a request with the HTTP
 
 If the client provides the large response MIME type, the Lambda will not log an error using `Log.error`. Instead, it will rewrite the original response with a reference to the offloaded large payload. Furthermore, the rewritten response will include the HTTP header `Content-Type` with the value `application/large-response.vnd+json`.
 
-If the client does not provide the large response MIME type, the Lambda will log an error with `Log.error`, and the response will fail due to an excessively large response body.
+If the client does not provide the large response MIME type, the Lambda will log an error with `Log.error` and rewrite the original response with a custom message (can be configured) and HTTP status code 413 (Payload Too Large).
 
 ### Middleware Configuration:
 
 Supported Parameters:
 
-| Parameter       | Type                | Description                                                                  |
-| --------------- | ------------------- | ---------------------------------------------------------------------------- |
-| thresholdWarn   | `number`            | Warning threshold level (percentage of `sizeLimitInMB`), e.g: 0.80           |
-| thresholdError  | `number`            | Error threshold level (percentage of `sizeLimitInMB`), e.g: 0.90             |
-| sizeLimitInMB   | `number`            | Maximum allowed size limit in MB, e.g 6                                      |
-| outputBucket    | `string`            | Identifier or name of the output S3 bucket                                   |
+| Parameter | Type | Description |
+| --- | --- | --- |
+| thresholdWarn | `number` | Warning threshold level (percentage of `sizeLimitInMB`), e.g: 0.80 |
+| thresholdError | `number` | Error threshold level (percentage of `sizeLimitInMB`), e.g: 0.90 |
+| sizeLimitInMB | `number` | Maximum allowed size limit in MB, e.g 6 |
+| outputBucket | `string` | Identifier or name of the output S3 bucket |
+| customErrorMessage | `string \| (event:APIGatewayProxyEventV2) => string ` | Custom error message to be returned when the response is too large and the client does not support large responses (no accept header) |
 | groupRequestsBy | `function - mapper` | Function to group requests, based on API Gateway event V2. Defaults to 'all' |
 
 Example Usage:

diff --git a/packages/large-response-middleware/src/index.test.ts b/packages/large-response-middleware/src/index.test.ts
@@ -4,8 +4,8 @@ import * as Lambda from 'aws-lambda';
 
 import { getOrgIdFromContext } from './__tests__/util';
 
-import { LARGE_RESPONSE_MIME_TYPE, withLargeResponseHandler } from './';
 import * as middleware from './';
+import { LARGE_RESPONSE_MIME_TYPE, withLargeResponseHandler } from './';
 
 const uploadFileSpy = jest.spyOn(middleware, 'uploadFile').mockResolvedValue({
   filename: 'red-redington/2023-12-13/la-caballa',
@@ -123,6 +123,94 @@ describe('withLargeResponseHandler', () => {
     });
   });
 
+  it('should overwrite response with default ERROR message + correct status when content length is over ERROR threshold', async () => {
+    const middleware = withLargeResponseHandler({
+      thresholdWarn: 0.5,
+      thresholdError: 0.9,
+      sizeLimitInMB: 1,
+      outputBucket: 'the-bucket-list',
+      groupRequestsBy: getOrgIdFromContext,
+    });
+    const content = Buffer.alloc(1024 * 1024, 'a').toString();
+    const requestResponseContext = {
+      event: {
+        requestContext: {},
+      },
+      response: {
+        headers: {
+          random: Buffer.alloc(0.85 * 1024 * 1024, 'a').toString(), // 0.85MB
+        },
+        body: content,
+      },
+    } as any;
+
+    await middleware.after(requestResponseContext);
+
+    expect(JSON.parse(requestResponseContext.response?.body)?.message).toBe(
+      "Call the API with the HTTP header 'Accept: application/large-response.vnd+json' to receive the payload through an S3 ref and avoid HTTP 500 errors.",
+    );
+    expect(requestResponseContext?.response?.statusCode).toBe(413);
+  });
+
+  it('should overwrite response with ERROR message (string) + correct status when content length is over ERROR threshold', async () => {
+    const middleware = withLargeResponseHandler({
+      thresholdWarn: 0.5,
+      thresholdError: 0.9,
+      customErrorMessage: 'Custom error message',
+      sizeLimitInMB: 1,
+      outputBucket: 'the-bucket-list',
+      groupRequestsBy: getOrgIdFromContext,
+    });
+    const content = Buffer.alloc(1024 * 1024, 'a').toString();
+    const requestResponseContext = {
+      event: {
+        requestContext: {},
+      },
+      response: {
+        headers: {
+          random: Buffer.alloc(0.85 * 1024 * 1024, 'a').toString(), // 0.85MB
+        },
+        body: content,
+      },
+    } as any;
+
+    await middleware.after(requestResponseContext);
+
+    expect(JSON.parse(requestResponseContext.response?.body)?.message).toBe('Custom error message');
+    expect(requestResponseContext?.response?.statusCode).toBe(413);
+  });
+
+  it('should overwrite response with custom ERROR message (callback function) + correct status when content length is over ERROR threshold', async () => {
+    const middleware = withLargeResponseHandler({
+      thresholdWarn: 0.5,
+      thresholdError: 0.9,
+      customErrorMessage: (event: Lambda.APIGatewayProxyEventV2) =>
+        `Custom error message for ${event.requestContext?.requestId}`,
+      sizeLimitInMB: 1,
+      outputBucket: 'the-bucket-list',
+      groupRequestsBy: getOrgIdFromContext,
+    });
+    const content = Buffer.alloc(1024 * 1024, 'a').toString();
+    const requestResponseContext = {
+      event: {
+        requestContext: {
+          requestId: 'request-id-123',
+        },
+      },
+      response: {
+        headers: {
+          random: Buffer.alloc(0.85 * 1024 * 1024, 'a').toString(), // 0.85MB
+        },
+        body: content,
+      },
+    } as any;
+
+    await middleware.after(requestResponseContext);
+
+    expect(JSON.parse(requestResponseContext.response?.body)?.message).toBe('Custom error message for request-id-123');
+    expect(requestResponseContext?.response?.statusCode).toBe(413);
+  });
+
   describe('when request header "Accept":"application/large-response.vnd+json" is given', () => {
     it('should not log ERROR with "Large response detected (limit exceeded)" when content length is over ERROR threshold', async () => {
       const middleware = withLargeResponseHandler({

diff --git a/packages/large-response-middleware/src/index.ts b/packages/large-response-middleware/src/index.ts
@@ -26,17 +26,21 @@ export type FileUploadContext = {
   fileName: string;
 };
 
+export type CustomErrorMessage = string | ((event: APIGatewayProxyEventV2) => string);
+
 export const withLargeResponseHandler = ({
   thresholdWarn,
   thresholdError,
   sizeLimitInMB: _sizeLimitInMB,
   outputBucket,
+  customErrorMessage,
   groupRequestsBy,
 }: {
   thresholdWarn: number;
   thresholdError: number;
   sizeLimitInMB: number;
   outputBucket: string;
+  customErrorMessage?: CustomErrorMessage;
   groupRequestsBy?: (event: APIGatewayProxyEventV2) => string;
 }) => {
   return {
@@ -101,6 +105,11 @@ export const withLargeResponseHandler = ({
               response_size_mb: contentLengthMB.toFixed(2),
               $payload_ref,
             });
+            response.isBase64Encoded = false;
+            response.statusCode = 413;
+            response.body = JSON.stringify({
+              message: getCustomErrorMessage(customErrorMessage, event),
+            });
           }
         } else if (contentLengthMB > thresholdWarnInMB) {
           Log.warn(`Large response detected. ${LARGE_RESPONSE_USER_INFO}`, {
@@ -189,3 +198,9 @@ function getFormattedDate() {
 
   return date.toISOString().split('T')[0];
 }
+
+function getCustomErrorMessage(customErrorMessage: CustomErrorMessage | undefined, event: APIGatewayProxyEventV2) {
+  return typeof customErrorMessage === 'function'
+    ? customErrorMessage(event)
+    : customErrorMessage ?? LARGE_RESPONSE_USER_INFO;
+}