chore: use bodyLengthChecker for calculating body length

aws · Dec 29, 2023 · f2ba37b · f2ba37b
1 parent 07de246
commit f2ba37b
Show file tree

Hide file tree

Showing 5 changed files with 25 additions and 8 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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
@@ -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);
   });