Merge pull request #3 from epilot-dev/feat.support-client-handling-of…

…-large-responses

feat: support client handling of large responses

packages/large-response-middleware/README.md

@@ -23,6 +23,8 @@ 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`.
 
+When the client doesn't provide the large response MIME type, and prefers to deal with the large response as a bad request instead of an HTTP 500, the client can send the `Handle-Large-Response: true` header. The Lambda will rewrite the original response with a custom message and HTTP status code 413 (Payload Too Large). This enables the client to detect a large response and handle it accordingly, by calling the API with a more strict filtering criteria.
+
 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:

packages/large-response-middleware/babel.config.js

@@ -0,0 +1,5 @@
+module.exports = {
+  presets: [
+    ['@babel/preset-env', { targets: { node: 'current' } }],
+  ],
+};

packages/large-response-middleware/docs/architecture-1.plantuml

@@ -9,11 +9,17 @@ Client -> "API Gateway": Request data with/without header
 "API Gateway" -> Lambda: Invoke Lambda
 
 alt If "Accept: application:large-response.vnd+json" is present
+    Lambda -> Lambda: Logs info event
     Lambda -> "S3 Bucket": Save large response
     "S3 Bucket" -> Lambda: Return S3 ref
     Lambda -> "API Gateway": Return S3 ref `{ $payload_ref: "s3://..." }`
     "API Gateway" -> Client: Return $payload_ref
+else If "Handle-Large-Request: true" is present
+    Lambda -> Lambda: Logs info event
+    Lambda -> Lambda: Skips S3 content dump
+    Lambda -> Client: Response 413 (Payload Too Large)
 else If request header is not present
+    Lambda -> Lambda: Logs error event
     Lambda -> "API Gateway": Response 500 (Internal Server Error)
     "API Gateway" -> Client: Response 413 (Payload Too Large)
 end

packages/large-response-middleware/jest.config.ts

@@ -14,6 +14,11 @@ const config: Config.InitialOptions = {
   maxWorkers: 1, // run tests sequentially
   moduleDirectories: ['node_modules', 'src'],
   modulePaths: ['node_modules'],
+  transform: {
+    '^.+\\.js$': 'babel-jest',
+  },
+  transformIgnorePatterns: ['/node_modules/(?!(yn)/)'],
+  extensionsToTreatAsEsm: ['.ts'],
 };
 
 export default config;

packages/large-response-middleware/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@epilot/large-response-middleware",
-  "version": "0.0.11",
+  "version": "0.0.15",
   "license": "MIT",
   "repository": {
     "type": "git",
@@ -40,14 +40,14 @@
   },
   "size-limit": [
     {
-      "limit": "2 kB",
+      "limit": "3 kB",
       "path": "lib/index.js"
     }
   ],
   "devDependencies": {
     "@babel/cli": "^7.23.4",
-    "@babel/core": "^7.23.6",
-    "@babel/preset-env": "^7.23.6",
+    "@babel/core": "^7.25.2",
+    "@babel/preset-env": "^7.25.4",
     "@babel/preset-typescript": "^7.23.3",
     "@dazn/lambda-powertools-logger": "^1.28.1",
     "@middy/core": "^2.5.7",
@@ -62,6 +62,7 @@
     "@types/node": "^20.10.5",
     "aws-lambda": "^1.0.7",
     "aws-sdk": "^2.816.0",
+    "babel-jest": "^29.7.0",
     "cross-env": "^7.0.3",
     "jest": "^29.7.0",
     "jest-junit": "^16.0.0",
@@ -82,6 +83,7 @@
     "aws-sdk": "^2.816.0"
   },
   "dependencies": {
-    "core-js": "^3.34.0"
+    "core-js": "^3.34.0",
+    "yn": "^5.0.0"
   }
 }

packages/large-response-middleware/src/index.test.ts

@@ -5,7 +5,12 @@ import * as Lambda from 'aws-lambda';
 import { getOrgIdFromContext } from './__tests__/util';
 
 import * as middleware from './';
-import { LARGE_RESPONSE_MIME_TYPE, withLargeResponseHandler } from './';
+import {
+  LARGE_RESPONSE_HANDLED_INFO,
+  LARGE_RESPONSE_MIME_TYPE,
+  LARGE_RESPONSE_USER_INFO,
+  withLargeResponseHandler,
+} from './';
 
 const uploadFileSpy = jest.spyOn(middleware, 'uploadFile').mockResolvedValue({
   filename: 'red-redington/2023-12-13/la-caballa',
@@ -63,18 +68,15 @@ describe('withLargeResponseHandler', () => {
       },
     } as any);
 
-    expect(LogWarnSpy).toHaveBeenCalledWith(
-      "Large response detected. 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.",
-      {
-        contentLength: 1572872,
-        event: { requestContext: {} },
-        request: {},
-        response_size_mb: '1.50',
-        $payload_ref: expect.stringMatching(
-          /http:\/\/localhost:4566\/the-bucket-list\/red-redington\/\d+-\d+-\d+\/la-caballa/,
-        ),
-      },
-    );
+    expect(LogWarnSpy).toHaveBeenCalledWith(`Large response detected. ${LARGE_RESPONSE_USER_INFO}`, {
+      contentLength: 1572872,
+      event: { requestContext: {} },
+      request: {},
+      response_size_mb: '1.50',
+      $payload_ref: expect.stringMatching(
+        /http:\/\/localhost:4566\/the-bucket-list\/red-redington\/\d+-\d+-\d+\/la-caballa/,
+      ),
+    });
   });
 
   it('should log ERROR with "Large response detected (limit exceeded)" when content length is over ERROR threshold', async () => {
@@ -101,18 +103,15 @@ describe('withLargeResponseHandler', () => {
 
     await middleware.after(requestResponseContext);
 
-    expect(LogErrorSpy).toHaveBeenCalledWith(
-      "Large response detected (limit exceeded). 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.",
-      {
-        contentLength: 1939873,
-        event: { requestContext: {} },
-        request: {},
-        response_size_mb: '1.85',
-        $payload_ref: expect.stringMatching(
-          /http:\/\/localhost:4566\/the-bucket-list\/red-redington\/\d+-\d+-\d+\/la-caballa/,
-        ),
-      },
-    );
+    expect(LogErrorSpy).toHaveBeenCalledWith(`Large response detected (limit exceeded). ${LARGE_RESPONSE_USER_INFO}`, {
+      contentLength: 1939873,
+      event: { requestContext: {} },
+      request: {},
+      response_size_mb: '1.85',
+      $payload_ref: expect.stringMatching(
+        /http:\/\/localhost:4566\/the-bucket-list\/red-redington\/\d+-\d+-\d+\/la-caballa/,
+      ),
+    });
 
     expect(uploadFileSpy).toHaveBeenCalledWith({
       bucket: 'the-bucket-list',
@@ -146,9 +145,7 @@ describe('withLargeResponseHandler', () => {
 
     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(JSON.parse(requestResponseContext.response?.body)?.message).toBe(LARGE_RESPONSE_USER_INFO);
     expect(requestResponseContext?.response?.statusCode).toBe(413);
   });
 
@@ -267,4 +264,61 @@ describe('withLargeResponseHandler', () => {
       });
     });
   });
+
+  describe('when request header "X-Handle-Large-Response:true" is given', () => {
+    it('should return 413 and not log ERROR with "Large response detected (limit exceeded)" 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 LogErrorSpy = jest.spyOn(Log, 'error');
+      const content = Buffer.alloc(1024 * 1024, 'a').toString();
+      const requestResponseContext = {
+        event: {
+          requestContext: {
+            requestId: 'request-id-123',
+            authorizer: {
+              lambda: {
+                organizationId: 'red-redington',
+              },
+            },
+          } as any,
+          headers: {
+            'Handle-Large-Response': 'true',
+          },
+        } as Partial<Lambda.APIGatewayProxyEventV2>,
+        response: {
+          headers: {
+            random: Buffer.alloc(0.85 * 1024 * 1024, 'a').toString(), // 0.85MB
+          },
+          body: content,
+        },
+      } as any;
+
+      await middleware.after(requestResponseContext);
+
+      expect(LogErrorSpy).not.toHaveBeenCalled();
+      expect(uploadFileSpy).not.toHaveBeenCalled();
+
+      const parsedBody = JSON.parse(requestResponseContext.response.body);
+
+      expect(requestResponseContext.response).toMatchObject({
+        isBase64Encoded: false,
+        statusCode: 413,
+        headers: {
+          random: requestResponseContext.response.headers.random,
+          'content-type': 'application/large-response.vnd+json',
+        },
+      });
+      expect(parsedBody).toMatchObject({
+        meta: {
+          content_length_mb: '1.85',
+        },
+        message: LARGE_RESPONSE_HANDLED_INFO,
+      });
+    });
+  });
 });

packages/large-response-middleware/src/index.ts

@@ -3,6 +3,7 @@ import path from 'path';
 import Log from '@dazn/lambda-powertools-logger';
 import middy from '@middy/core';
 import type { APIGatewayProxyEventV2, APIGatewayProxyStructuredResultV2 } from 'aws-lambda';
+import yn from 'yn';
 
 import { getS3Client } from './s3/s3-client';
 
@@ -16,7 +17,9 @@ const TO_MB_FACTOR = 1_048_576.0;
  */
 export const LIMIT_REQUEST_SIZE_MB = 6.0;
 export const LARGE_RESPONSE_MIME_TYPE = 'application/large-response.vnd+json';
-const LARGE_RESPONSE_USER_INFO = `Call the API with the HTTP header 'Accept: ${LARGE_RESPONSE_MIME_TYPE}' to receive the payload through an S3 ref and avoid HTTP 500 errors.`;
+export const HANDLE_LARGE_RESPONSE_HEADER = 'handle-large-response';
+export const LARGE_RESPONSE_USER_INFO = `Call the API with the HTTP header 'Accept: ${LARGE_RESPONSE_MIME_TYPE}' to receive the payload through an S3 ref and avoid 413 errors or '${HANDLE_LARGE_RESPONSE_HEADER}: true' to acknowledge you can handle the 413.`;
+export const LARGE_RESPONSE_HANDLED_INFO = `'${HANDLE_LARGE_RESPONSE_HEADER}: true' received means client can handle this event. The response is too large and can't be returned to the client.`;
 
 export type FileUploadContext = {
   bucket: string;
@@ -64,10 +67,13 @@ export const withLargeResponseHandler = ({
         const sizeLimitInMB = (_sizeLimitInMB ?? LIMIT_REQUEST_SIZE_MB) * 1.0;
         const thresholdWarnInMB = (thresholdWarn ?? 0.0) * 1.0 * sizeLimitInMB;
         const thresholdErrorInMB = (thresholdError ?? 0.0) * 1.0 * sizeLimitInMB;
+        const clientCanHandleLargeResponseBadRequest = Object.entries(requestHeaders).find(
+          ([header, v]) => header.toLowerCase() === HANDLE_LARGE_RESPONSE_HEADER && yn(v),
+        );
 
         let $payload_ref = null;
 
-        if (contentLengthMB > thresholdWarnInMB) {
+        if (contentLengthMB > thresholdWarnInMB && !clientCanHandleLargeResponseBadRequest) {
           const { url } = await safeUploadLargeResponse({
             groupId: String(groupId),
             contentType: 'application/json',
@@ -97,6 +103,27 @@ export const withLargeResponseHandler = ({
               response_size_mb: contentLengthMB.toFixed(2),
               $payload_ref,
             });
+          } else if (clientCanHandleLargeResponseBadRequest) {
+            response.isBase64Encoded = false;
+            response.statusCode = 413;
+
+            response.body = JSON.stringify({
+              meta: {
+                content_length_mb: contentLengthMB.toFixed(2),
+              },
+              message: getCustomErrorMessage(customErrorMessage || LARGE_RESPONSE_HANDLED_INFO, event),
+            });
+
+            response.headers = { ...response.headers, ['content-type']: LARGE_RESPONSE_MIME_TYPE };
+            Log.info(
+              `Large response detected (limit exceeded). Client signaled that it can handle large responses via 413. Rewriting response with { metadata, message } `,
+              {
+                contentLength: aproxContentLengthBytes,
+                event,
+                request: event.requestContext,
+                response_size_mb: contentLengthMB.toFixed(2),
+              },
+            );
           } else {
             Log.error(`Large response detected (limit exceeded). ${LARGE_RESPONSE_USER_INFO}`, {
               contentLength: aproxContentLengthBytes,