From f2ba37b8859c4bfdc66c71e84d32d2daedc096ce Mon Sep 17 00:00:00 2001
From: "Kamat, Trivikram" <16024985+trivikr@users.noreply.github.com>
Date: Fri, 29 Dec 2023 18:18:54 +0000
Subject: [PATCH] chore: use bodyLengthChecker for calculating body length

---
 .../src/compressionMiddleware.spec.ts               |  3 ++-
 .../src/compressionMiddleware.ts                    |  9 ++++++---
 .../middleware-compression/src/configurations.ts    |  7 +++++++
 .../src/getCompressionPlugin.spec.ts                |  1 +
 .../src/resolveCompressionConfig.spec.ts            | 13 +++++++++----
 5 files changed, 25 insertions(+), 8 deletions(-)

diff --git a/packages/middleware-compression/src/compressionMiddleware.spec.ts b/packages/middleware-compression/src/compressionMiddleware.spec.ts
index e7a6db9c0550..1691d977c206 100644
--- a/packages/middleware-compression/src/compressionMiddleware.spec.ts
+++ b/packages/middleware-compression/src/compressionMiddleware.spec.ts
@@ -12,7 +12,9 @@ jest.mock("./compressStream");
 jest.mock("./isStreaming");
 
 describe(compressionMiddleware.name, () => {
+  const mockBody = "body";
   const mockConfig = {
+    bodyLengthChecker: jest.fn().mockReturnValue(mockBody.length),
     disableRequestCompression: false,
     requestMinCompressionSizeBytes: 0,
   };
@@ -22,7 +24,6 @@ describe(compressionMiddleware.name, () => {
 
   const mockNext = jest.fn();
   const mockContext = {};
-  const mockBody = "body";
   const mockArgs = { request: { headers: {}, body: mockBody } };
 
   afterEach(() => {
diff --git a/packages/middleware-compression/src/compressionMiddleware.ts b/packages/middleware-compression/src/compressionMiddleware.ts
index 03d6e2b27aff..2b41a7d5052d 100644
--- a/packages/middleware-compression/src/compressionMiddleware.ts
+++ b/packages/middleware-compression/src/compressionMiddleware.ts
@@ -60,9 +60,12 @@ export const compressionMiddleware =
             // Invalid case. We should never get here.
             throw new Error("Compression is not supported for streaming blobs that require a length.");
           }
-        } else if (body.length >= config.requestMinCompressionSizeBytes) {
-          updatedBody = await compressString(body);
-          isRequestCompressed = true;
+        } else {
+          const bodyLength = config.bodyLengthChecker(body);
+          if (bodyLength && bodyLength >= config.requestMinCompressionSizeBytes) {
+            updatedBody = await compressString(body);
+            isRequestCompressed = true;
+          }
         }
 
         if (isRequestCompressed) {
diff --git a/packages/middleware-compression/src/configurations.ts b/packages/middleware-compression/src/configurations.ts
index 44c0f83a1779..e25c1e23bc33 100644
--- a/packages/middleware-compression/src/configurations.ts
+++ b/packages/middleware-compression/src/configurations.ts
@@ -1,7 +1,14 @@
+import { BodyLengthCalculator } from "@smithy/types";
+
 /**
  * @public
  */
 export interface CompressionInputConfig {
+  /**
+   * A function that can calculate the length of a body.
+   */
+  bodyLengthChecker: BodyLengthCalculator;
+
   /**
    * Whether to disable request compression.
    */
diff --git a/packages/middleware-compression/src/getCompressionPlugin.spec.ts b/packages/middleware-compression/src/getCompressionPlugin.spec.ts
index ec5b01e32eea..7752d980bdea 100644
--- a/packages/middleware-compression/src/getCompressionPlugin.spec.ts
+++ b/packages/middleware-compression/src/getCompressionPlugin.spec.ts
@@ -5,6 +5,7 @@ jest.mock("./compressionMiddleware");
 
 describe(getCompressionPlugin.name, () => {
   const config = {
+    bodyLengthChecker: jest.fn(),
     disableRequestCompression: false,
     requestMinCompressionSizeBytes: 0,
   };
diff --git a/packages/middleware-compression/src/resolveCompressionConfig.spec.ts b/packages/middleware-compression/src/resolveCompressionConfig.spec.ts
index b0ffb18a3552..e94a789a2629 100644
--- a/packages/middleware-compression/src/resolveCompressionConfig.spec.ts
+++ b/packages/middleware-compression/src/resolveCompressionConfig.spec.ts
@@ -1,10 +1,15 @@
 import { resolveCompressionConfig } from "./resolveCompressionConfig";
 
 describe(resolveCompressionConfig.name, () => {
+  const mockConfig = {
+    bodyLengthChecker: jest.fn(),
+    disableRequestCompression: false,
+    requestMinCompressionSizeBytes: 0,
+  };
   it("should throw an error if requestMinCompressionSizeBytes is less than 0", () => {
     const requestMinCompressionSizeBytes = -1;
     expect(() => {
-      resolveCompressionConfig({ disableRequestCompression: false, requestMinCompressionSizeBytes });
+      resolveCompressionConfig({ ...mockConfig, requestMinCompressionSizeBytes });
     }).toThrow(
       new RangeError(
         "The value for requestMinCompressionSizeBytes must be between 0 and 10485760 inclusive. " +
@@ -16,7 +21,7 @@ describe(resolveCompressionConfig.name, () => {
   it("should throw an error if requestMinCompressionSizeBytes is greater than 10485760", () => {
     const requestMinCompressionSizeBytes = 10485761;
     expect(() => {
-      resolveCompressionConfig({ disableRequestCompression: false, requestMinCompressionSizeBytes });
+      resolveCompressionConfig({ ...mockConfig, requestMinCompressionSizeBytes });
     }).toThrow(
       new RangeError(
         "The value for requestMinCompressionSizeBytes must be between 0 and 10485760 inclusive. " +
@@ -26,13 +31,13 @@ describe(resolveCompressionConfig.name, () => {
   });
 
   it.each([0, 10240, 10485760])("returns requestMinCompressionSizeBytes value %s", (requestMinCompressionSizeBytes) => {
-    const inputConfig = { disableRequestCompression: false, requestMinCompressionSizeBytes };
+    const inputConfig = { ...mockConfig, requestMinCompressionSizeBytes };
     const resolvedConfig = resolveCompressionConfig(inputConfig);
     expect(inputConfig).toEqual(resolvedConfig);
   });
 
   it.each([false, true])("returns disableRequestCompression value %s", (disableRequestCompression) => {
-    const inputConfig = { disableRequestCompression, requestMinCompressionSizeBytes: 0 };
+    const inputConfig = { ...mockConfig, disableRequestCompression };
     const resolvedConfig = resolveCompressionConfig(inputConfig);
     expect(inputConfig).toEqual(resolvedConfig);
   });