Refactored the logic to remove checks for content types

Author: sdangol

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

@@ -88,20 +88,14 @@ export const SAFE_CHARS = "-._~()'!*:@,;=+&$";
 
 export const UNSAFE_CHARS = '%<> \\[\\]{}|^';
 
-/**
- * Match for compressible content type.
- */
-export const COMPRESSIBLE_CONTENT_TYPE_REGEX = {
-  COMMON: /^\s*application\/json(?:[;\s]|$)/i,
-  OCCASIONAL:
-    /^\s*(?:text\/(?!event-stream(?:[;\s]|$))[^;\s]+|application\/(?:xml|javascript)|[^;\s]+?\+(?:json|text|xml|yaml))(?:[;\s]|$)/i,
-  RARE: /^\s*(?:application\/(?:xml-dtd|ecmascript|dart|postscript|rtf|tar|toml|vnd\.dart|vnd\.ms-fontobject|vnd\.ms-opentype|wasm|x-httpd-php|x-javascript|x-ns-proxy-autoconfig|x-sh|x-tar|x-virtualbox-hdd|x-virtualbox-ova|x-virtualbox-ovf|x-virtualbox-vbox|x-virtualbox-vdi|x-virtualbox-vhd|x-virtualbox-vmdk|x-www-form-urlencoded)|font\/(?:otf|ttf)|image\/(?:bmp|vnd\.adobe\.photoshop|vnd\.microsoft\.icon|vnd\.ms-dds|x-icon|x-ms-bmp)|message\/rfc822|model\/gltf-binary|x-shader\/x-fragment|x-shader\/x-vertex)(?:[;\s]|$)/i,
-};
+export const DEFAULT_COMPRESSION_RESPONSE_THRESHOLD = 1024;
 
 export const CACHE_CONTROL_NO_TRANSFORM_REGEX =
   /(?:^|,)\s*?no-transform\s*?(?:,|$)/i;
 
 export const COMPRESSION_ENCODING_TYPES = {
   GZIP: 'gzip',
   DEFLATE: 'deflate',
+  IDENTITY: 'identity',
+  ANY: '*',
 } as const;

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

@@ -1,5 +1,6 @@
 import type { APIGatewayProxyEvent, APIGatewayProxyResult } from 'aws-lambda';
-import type { HandlerResponse } from '../types/rest.js';
+import type { CompressionOptions, HandlerResponse } from '../types/rest.js';
+import { COMPRESSION_ENCODING_TYPES } from './constants.js';
 import { isAPIGatewayProxyResult } from './utils.js';
 
 /**
@@ -89,11 +90,35 @@ export const webResponseToProxyResult = async (
     }
   }
 
+  // Check if response contains compressed/binary content
+  const contentEncoding = response.headers.get(
+    'content-encoding'
+  ) as CompressionOptions['encoding'];
+  let body: string;
+  let isBase64Encoded = false;
+
+  if (
+    contentEncoding &&
+    [
+      COMPRESSION_ENCODING_TYPES.GZIP,
+      COMPRESSION_ENCODING_TYPES.DEFLATE,
+    ].includes(contentEncoding)
+  ) {
+    // For compressed content, get as buffer and encode to base64
+    const buffer = await response.arrayBuffer();
+    body = Buffer.from(buffer).toString('base64');
+    isBase64Encoded = true;
+  } else {
+    // For text content, use text()
+    body = await response.text();
+    isBase64Encoded = false;
+  }
+
   const result: APIGatewayProxyResult = {
     statusCode: response.status,
     headers,
-    body: await response.text(),
-    isBase64Encoded: false,
+    body,
+    isBase64Encoded,
   };
 
   if (Object.keys(multiValueHeaders).length > 0) {

packages/event-handler/src/rest/middleware/compress.ts

@@ -2,8 +2,8 @@ import type { Middleware } from '../../types/index.js';
 import type { CompressionOptions } from '../../types/rest.js';
 import {
   CACHE_CONTROL_NO_TRANSFORM_REGEX,
-  COMPRESSIBLE_CONTENT_TYPE_REGEX,
   COMPRESSION_ENCODING_TYPES,
+  DEFAULT_COMPRESSION_RESPONSE_THRESHOLD,
 } from '../constants.js';
 
 /**
@@ -25,8 +25,8 @@ import {
  *
  * @example
  * ```typescript
- * import { Router } from '@aws-lambda-powertools/event-handler';
- * import { compress } from '@aws-lambda-powertools/event-handler/rest/middleware';
+ * import { Router } from '@aws-lambda-powertools/event-handler/experimental-rest';
+ * import { compress } from '@aws-lambda-powertools/event-handler/experimental-rest/middleware';
  *
  * const app = new Router();
  *
@@ -41,8 +41,8 @@ import {
  *
  * @example
  * ```typescript
- * import { Router } from '@aws-lambda-powertools/event-handler';
- * import { compress } from '@aws-lambda-powertools/event-handler/rest/middleware';
+ * import { Router } from '@aws-lambda-powertools/event-handler/experimental-rest';
+ * import { compress } from '@aws-lambda-powertools/event-handler/experimental-rest/middleware';
  *
  * const app = new Router();
  *
@@ -62,59 +62,59 @@ import {
  */
 
 const compress = (options?: CompressionOptions): Middleware => {
-  const threshold = options?.threshold ?? 1024;
+  const preferredEncoding =
+    options?.encoding ?? COMPRESSION_ENCODING_TYPES.GZIP;
+  const threshold =
+    options?.threshold ?? DEFAULT_COMPRESSION_RESPONSE_THRESHOLD;
 
   return async (_, reqCtx, next) => {
     await next();
 
-    const contentLength = reqCtx.res.headers.get('content-length');
-    const isEncodedOrChunked =
-      reqCtx.res.headers.has('content-encoding') ||
-      reqCtx.res.headers.has('transfer-encoding');
-
-    // Check if response should be compressed
     if (
-      isEncodedOrChunked ||
-      reqCtx.request.method === 'HEAD' ||
-      (contentLength && Number(contentLength) < threshold) ||
-      !shouldCompress(reqCtx.res) ||
-      !shouldTransform(reqCtx.res) ||
-      !reqCtx.res.body
+      !shouldCompress(reqCtx.request, reqCtx.res, preferredEncoding, threshold)
     ) {
       return;
     }
 
-    const acceptedEncoding = reqCtx.request.headers.get('accept-encoding');
-    const encoding =
-      options?.encoding ??
-      Object.values(COMPRESSION_ENCODING_TYPES).find((encoding) =>
-        acceptedEncoding?.includes(encoding)
-      ) ??
-      COMPRESSION_ENCODING_TYPES.GZIP;
-
     // Compress the response
-    const stream = new CompressionStream(encoding);
+    const stream = new CompressionStream(preferredEncoding);
     reqCtx.res = new Response(reqCtx.res.body.pipeThrough(stream), reqCtx.res);
     reqCtx.res.headers.delete('content-length');
-    reqCtx.res.headers.set('content-encoding', encoding);
+    reqCtx.res.headers.set('content-encoding', preferredEncoding);
   };
 };
 
-const shouldCompress = (res: Response) => {
-  const type = res.headers.get('content-type');
-  return (
-    type &&
-    (COMPRESSIBLE_CONTENT_TYPE_REGEX.COMMON.test(type) ||
-      COMPRESSIBLE_CONTENT_TYPE_REGEX.OCCASIONAL.test(type) ||
-      COMPRESSIBLE_CONTENT_TYPE_REGEX.RARE.test(type))
-  );
-};
+const shouldCompress = (
+  request: Request,
+  response: Response,
+  preferredEncoding: NonNullable<CompressionOptions['encoding']>,
+  threshold: NonNullable<CompressionOptions['threshold']>
+): response is Response & { body: NonNullable<Response['body']> } => {
+  const acceptedEncoding =
+    request.headers.get('accept-encoding') ?? COMPRESSION_ENCODING_TYPES.ANY;
+  const contentLength = response.headers.get('content-length');
+  const cacheControl = response.headers.get('cache-control');
+
+  const isEncodedOrChunked =
+    response.headers.has('content-encoding') ||
+    response.headers.has('transfer-encoding');
+
+  const shouldEncode =
+    !acceptedEncoding.includes(COMPRESSION_ENCODING_TYPES.IDENTITY) &&
+    (acceptedEncoding.includes(preferredEncoding) ||
+      acceptedEncoding.includes(COMPRESSION_ENCODING_TYPES.ANY));
 
-const shouldTransform = (res: Response) => {
-  const cacheControl = res.headers.get('cache-control');
-  // Don't compress for Cache-Control: no-transform
-  // https://tools.ietf.org/html/rfc7234#section-5.2.2.4
-  return !cacheControl || !CACHE_CONTROL_NO_TRANSFORM_REGEX.test(cacheControl);
+  if (
+    !shouldEncode ||
+    isEncodedOrChunked ||
+    request.method === 'HEAD' ||
+    (contentLength && Number(contentLength) < threshold) ||
+    (cacheControl && CACHE_CONTROL_NO_TRANSFORM_REGEX.test(cacheControl)) ||
+    !response.body
+  ) {
+    return false;
+  }
+  return true;
 };
 
 export { compress };

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

@@ -116,7 +116,7 @@ type ValidationResult = {
 };
 
 type CompressionOptions = {
-  encoding?: (typeof COMPRESSION_ENCODING_TYPES)[keyof typeof COMPRESSION_ENCODING_TYPES];
+  encoding?: 'gzip' | 'deflate';
   threshold?: number;
 };
 

packages/event-handler/tests/unit/rest/middleware/compress.test.ts

@@ -1,3 +1,4 @@
+import { gzipSync } from 'node:zlib';
 import context from '@aws-lambda-powertools/testing-utils/context';
 import { Router } from 'src/rest/Router.js';
 import { beforeEach, describe, expect, it } from 'vitest';
@@ -31,6 +32,9 @@ describe('Compress Middleware', () => {
     // Assess
     expect(result.headers?.['content-encoding']).toBe('gzip');
     expect(result.headers?.['content-length']).toBeUndefined();
+    expect(result.body).toEqual(
+      gzipSync(JSON.stringify(body)).toString('base64')
+    );
   });
 
   it('skips compression when content is below threshold', async () => {
@@ -98,50 +102,6 @@ describe('Compress Middleware', () => {
     expect(result.headers?.['content-encoding']).toEqual('gzip');
   });
 
-  it.each([
-    'image/jpeg',
-    'image/png',
-    'image/gif',
-    'audio/mpeg',
-    'audio/mp4',
-    'audio/ogg',
-    'video/mp4',
-    'video/mpeg',
-    'video/webm',
-    'application/zip',
-    'application/gzip',
-    'application/x-gzip',
-    'application/octet-stream',
-    'application/pdf',
-    'application/msword',
-    'text/event-stream',
-  ])(
-    'skips compression for non-compressible content types',
-    async (contentType) => {
-      // Prepare
-      const application = new Router();
-      application.get(
-        '/test',
-        [
-          compress(),
-          createSettingHeadersMiddleware({
-            'content-length': '2000',
-            'content-type': contentType,
-          }),
-        ],
-        async () => {
-          return body;
-        }
-      );
-
-      // Act
-      const result = await application.resolve(event, context);
-
-      // Assess
-      expect(result.headers?.['content-encoding']).toBeUndefined();
-    }
-  );
-
   it('skips compression when cache-control no-transform is set', async () => {
     // Prepare
     const application = new Router();
@@ -191,10 +151,10 @@ describe('Compress Middleware', () => {
     expect(result.headers?.['content-encoding']).toBe('deflate');
   });
 
-  it('infers encoding from Accept-Encoding header', async () => {
+  it('does not compress if Accept-Encoding is set to identity', async () => {
     // Prepare
     const deflateCompressionEvent = createTestEvent('/test', 'GET', {
-      'Accept-Encoding': 'deflate',
+      'Accept-Encoding': 'identity',
     });
     app.get('/test', async () => {
       return body;
@@ -204,6 +164,6 @@ describe('Compress Middleware', () => {
     const result = await app.resolve(deflateCompressionEvent, context);
 
     // Assess
-    expect(result.headers?.['content-encoding']).toBe('deflate');
+    expect(result.headers?.['content-encoding']).not.toBeDefined;
   });
 });