Serverless Typescript Lambda Gzip Middleware

gzipMiddleware.ts

import middy from "@middy/core";
import get from "lodash/get";
import set from "lodash/set";
import { gzipSync } from "zlib";
import isString from "lodash/isString";
import { hasProperty } from "@smallable/help-purchase-shared";

export type GzipMiddlewareConfig = Record<string, never>;

type NormalizedResponse = {
  body?: unknown;
  headers: { [header: string]: string };
  statusCode?: number;
};

export const normalizeHttpResponse = (
  response: unknown
): NormalizedResponse => {
  // May require updating to catch other types
  if (response === undefined) {
    return { headers: {} };
  }

  if (hasProperty(response, "body")) {
    return { ...response, headers: get(response, "headers", {}) };
  }

  return {
    body: response,
    statusCode: 200,
    headers: {},
  };
};

export const gzipMiddleware: middy.Middleware<GzipMiddlewareConfig> = () => ({
  after: (
    handler: middy.HandlerLambda<unknown, unknown>,
    next: middy.NextFunction
  ) => {
    const acceptEncodingHeader = get(
      handler,
      'event.headers["accept-encoding"]',
      get(handler, 'event.headers["Accept-Encoding"]', "")
    );
    const contentTypeHeader = get(
      handler,
      'event.headers["content-type"]',
      get(handler, 'event.headers["Content-Type"]')
    );

    // Gzip if:
    // - client accepts encoding
    // - it is a JSON response
    // - the response is not already encoded
    if (
      contentTypeHeader === "application/json" &&
      get(handler, "response.isBase64Encoded", false) === false &&
      isString(acceptEncodingHeader) &&
      acceptEncodingHeader.includes("gzip")
    ) {
      const normalizedResponse = normalizeHttpResponse(
        get(handler, "response")
      );
      set(handler, "response", normalizedResponse);

      // Do not gzip empty response
      const responseBody = get(normalizedResponse, "body");
      if (responseBody !== undefined) {
        // Gzip response
        const bufferedResponse = Buffer.from(JSON.stringify(responseBody));
        const compressedResponse =
          gzipSync(bufferedResponse).toString("base64");

        set(handler, "response.body", compressedResponse);
        set(handler, "response.isBase64Encoded", true);
        set(handler, 'response.headers["Content-Encoding"]', "gzip");
      }
    }

    next();
  },
});


// gzipMiddleware.test.ts
import middy from "@middy/core";
import get from "lodash/get";
import { gzipMiddleware } from "./gzipMiddleware";

beforeEach(() => {
  jest.clearAllMocks();
});

const getHandlerMock = (
  headers: { [header: string]: string },
  response: unknown
) =>
  ({
    event: {
      headers: { "Content-Type": "application/json", ...headers },
    },
    response,
  } as middy.HandlerLambda);

describe("gzipMiddleware", () => {
  const middleware = gzipMiddleware({});
  const nextMock = jest.fn();

  it("should compress payload for standard response", async () => {
    if (middleware.after === undefined) {
      return;
    }
    const handler = getHandlerMock({ "Accept-Encoding": "gzip" }, "ahah");

    await middleware.after(handler, nextMock);

    const responseBody = get(handler.response, "body");
    const responseHeaders: Map<string, string> = get(
      handler.response,
      "headers"
    );
    expect(typeof responseBody).toBe("string");
    expect(responseBody).not.toBe("ahah");
    expect(responseHeaders["Content-Encoding"]).toBe("gzip");
    expect(nextMock).toHaveBeenCalled();
  });

  it("should compress payload for extended response", async () => {
    if (middleware.after === undefined) {
      return;
    }

    const handler = getHandlerMock(
      { "Accept-Encoding": "gzip" },
      { body: "ahah" }
    );

    await middleware.after(handler, nextMock);

    const responseBody = get(handler.response, "body");
    const responseHeaders: Map<string, string> = get(
      handler.response,
      "headers"
    );
    expect(typeof responseBody).toBe("string");
    expect(responseBody).not.toBe("ahah");
    expect(responseHeaders["Content-Encoding"]).toBe("gzip");
    expect(nextMock).toHaveBeenCalled();
  });

  it("should not compress payload if no accept encoding header", async () => {
    if (middleware.after === undefined) {
      return;
    }

    const handler = getHandlerMock({}, { body: "ahah" });

    await middleware.after(handler, nextMock);

    const responseBody = get(handler.response, "body");
    expect(typeof responseBody).toBe("string");
    expect(responseBody).toBe("ahah");
    expect(get(handler.response, "headers['Content-Encoding']")).toBe(
      undefined
    );
    expect(nextMock).toHaveBeenCalled();
  });

  it("should not compress payload if already encoded", async () => {
    if (middleware.after === undefined) {
      return;
    }

    const handler = getHandlerMock(
      { "Accept-Encoding": "gzip" },
      { body: "somebase64string", statusCode: 200, isBase64Encoded: true }
    );

    await middleware.after(handler, nextMock);

    const responseBody = get(handler.response, "body");
    expect(typeof responseBody).toBe("string");
    expect(responseBody).toBe("somebase64string");
    expect(get(handler.response, "headers['Content-Encoding']")).toBe(
      undefined
    );
    expect(nextMock).toHaveBeenCalled();
  });

  it("should not compress payload if is not a JSON response", async () => {
    if (middleware.after === undefined) {
      return;
    }

    const handler = getHandlerMock(
      { "Accept-Encoding": "gzip", "Content-Type": "text/plain" },
      "ahah"
    );

    await middleware.after(handler, nextMock);

    const responseBody = handler.response;
    expect(typeof responseBody).toBe("string");
    expect(responseBody).toBe("ahah");
    expect(get(handler.response, "headers['Content-Encoding']")).toBe(
      undefined
    );
    expect(nextMock).toHaveBeenCalled();
  });
});