fix(event-handler): threshold limit for compression not respected when content-length not set

Author: svozza

packages/event-handler/src/http/Router.ts

@@ -24,6 +24,7 @@ import type {
   ErrorHandler,
   ErrorResolveOptions,
   HttpMethod,
+  HttpResolveOptions,
   HttpRouteOptions,
   HttpRouterOptions,
   Middleware,
@@ -221,7 +222,7 @@ class Router {
   async #resolve(
     event: unknown,
     context: Context,
-    options?: ResolveOptions
+    options?: HttpResolveOptions
   ): Promise<RequestContext> {
     if (
       !isAPIGatewayProxyEventV1(event) &&
@@ -248,7 +249,12 @@ class Router {
           event,
           context,
           req: new Request('https://invalid'),
-          res: new Response('', { status: HttpStatusCodes.METHOD_NOT_ALLOWED }),
+          res: new Response(null, {
+            status: HttpStatusCodes.METHOD_NOT_ALLOWED,
+            ...(options?.isHttpStreaming && {
+              headers: { 'transfer-encoding': 'chunked' },
+            }),
+          }),
           params: {},
           responseType,
         };
@@ -262,7 +268,12 @@ class Router {
       req,
       // this response should be overwritten by the handler, if it isn't
       // it means something went wrong with the middleware chain
-      res: new Response('', { status: HttpStatusCodes.INTERNAL_SERVER_ERROR }),
+      res: new Response('', {
+        status: HttpStatusCodes.INTERNAL_SERVER_ERROR,
+        ...(options?.isHttpStreaming && {
+          headers: { 'transfer-encoding': 'chunked' },
+        }),
+      }),
       params: {},
       responseType,
     };
@@ -393,7 +404,10 @@ class Router {
     context: Context,
     options: ResolveStreamOptions
   ): Promise<void> {
-    const reqCtx = await this.#resolve(event, context, options);
+    const reqCtx = await this.#resolve(event, context, {
+      ...options,
+      isHttpStreaming: true,
+    });
     await this.#streamHandlerResponse(reqCtx, options.responseStream);
   }
 

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

@@ -70,6 +70,14 @@ const compress = (options?: CompressionOptions): Middleware => {
   return async ({ reqCtx, next }) => {
     await next();
 
+    const transferEncoding = reqCtx.res.headers.get('transfer-encoding');
+    if (transferEncoding === 'chunked') return;
+
+    if (reqCtx.res.body !== null && !reqCtx.res.headers.has('content-length')) {
+      const body = await reqCtx.res.clone().arrayBuffer();
+      reqCtx.res.headers.set('content-length', body.byteLength.toString());
+    }
+
     if (!shouldCompress(reqCtx.req, reqCtx.res, preferredEncoding, threshold)) {
       return;
     }

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

@@ -34,6 +34,8 @@ type RequestContext = {
   isBase64Encoded?: boolean;
 };
 
+type HttpResolveOptions = ResolveOptions & { isHttpStreaming?: boolean };
+
 type ErrorResolveOptions = RequestContext & ResolveOptions;
 
 type ErrorHandler<T extends Error = Error> = (
@@ -263,6 +265,7 @@ export type {
   ErrorHandler,
   ErrorResolveOptions,
   HandlerResponse,
+  HttpResolveOptions,
   HttpStatusCode,
   HttpMethod,
   Middleware,

packages/event-handler/tests/unit/http/Router/streaming.test.ts

@@ -357,4 +357,19 @@ describe.each([
     expect(result.statusCode).toBe(200);
     expect(JSON.parse(result.body)).toEqual({ message: 'duplex stream body' });
   });
+
+  it('handles invalid HTTP method', async () => {
+    // Prepare
+    const app = new Router();
+    const handler = streamify(app);
+    const responseStream = new ResponseStream();
+    const event = createEvent('/test', 'TRACE');
+
+    // Act
+    const result = await handler(event, responseStream, context);
+
+    // Assess
+    expect(result.statusCode).toBe(405);
+    expect(result.body).toBe('');
+  });
 });

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

@@ -1,9 +1,14 @@
-import { gzipSync } from 'node:zlib';
+import { deflateSync, gzipSync } from 'node:zlib';
 import context from '@aws-lambda-powertools/testing-utils/context';
 import { beforeEach, describe, expect, it } from 'vitest';
+import { streamify } from '../../../../src/http/index.js';
 import { compress } from '../../../../src/http/middleware/index.js';
 import { Router } from '../../../../src/http/Router.js';
-import { createSettingHeadersMiddleware, createTestEvent } from '../helpers.js';
+import {
+  createSettingHeadersMiddleware,
+  createTestEvent,
+  ResponseStream,
+} from '../helpers.js';
 
 describe('Compress Middleware', () => {
   const event = createTestEvent('/test', 'GET');
@@ -62,6 +67,24 @@ describe('Compress Middleware', () => {
     expect(result.isBase64Encoded).toBe(false);
   });
 
+  it('skips compression when content is below threshold and content-length is not set', async () => {
+    // Prepare
+    const application = new Router();
+    const smallBody = { message: 'Small' };
+
+    application.use(compress({ threshold: 100 }));
+    application.get('/test', () => {
+      return smallBody;
+    });
+
+    // Act
+    const result = await application.resolve(event, context);
+
+    // Assess
+    expect(result.headers?.['content-encoding']).toBeUndefined();
+    expect(result.isBase64Encoded).toBe(false);
+  });
+
   it('skips compression for HEAD requests', async () => {
     // Prepare
     const headEvent = createTestEvent('/test', 'HEAD');
@@ -104,6 +127,7 @@ describe('Compress Middleware', () => {
     // Assess
     expect(result.headers?.['content-encoding']).toEqual('gzip');
     expect(result.isBase64Encoded).toBe(true);
+    expect(result.body).toBe(gzipSync(JSON.stringify(body)).toString('base64'));
   });
 
   it('skips compression when cache-control no-transform is set', async () => {
@@ -155,6 +179,9 @@ describe('Compress Middleware', () => {
     // Assess
     expect(result.headers?.['content-encoding']).toBe('deflate');
     expect(result.isBase64Encoded).toBe(true);
+    expect(result.body).toBe(
+      deflateSync(JSON.stringify(body)).toString('base64')
+    );
   });
 
   it('does not compress if Accept-Encoding is set to identity', async () => {
@@ -173,4 +200,26 @@ describe('Compress Middleware', () => {
     expect(result.headers?.['content-encoding']).toBeUndefined();
     expect(result.isBase64Encoded).toBe(false);
   });
+
+  it('skips compression and content-length in streaming mode', async () => {
+    // Prepare
+    const application = new Router();
+    application.use(compress());
+    application.get('/test', () => body);
+
+    const handler = streamify(application);
+    const responseStream = new ResponseStream();
+
+    // Act
+    const result = await handler(event, responseStream, context);
+
+    // Assess
+    expect(result.statusCode).toBe(200);
+    expect(result.headers['content-encoding']).toBeUndefined();
+    expect(result.headers['content-length']).toBeUndefined();
+    expect(result.body).toBe(JSON.stringify(body));
+    expect(result.body).not.toBe(
+      gzipSync(JSON.stringify(body)).toString('base64')
+    );
+  });
 });