chore: remove input[requestAlgorithmMember] if set by SDK

aws · Nov 26, 2024 · 9ebb747 · 9ebb747
1 parent 5ab602c
commit 9ebb747
Show file tree

Hide file tree

Showing 2 changed files with 77 additions and 18 deletions.
diff --git a/packages/middleware-flexible-checksums/src/flexibleChecksumsMiddleware.spec.ts b/packages/middleware-flexible-checksums/src/flexibleChecksumsMiddleware.spec.ts
@@ -3,7 +3,7 @@ import { BuildHandlerArguments } from "@smithy/types";
 import { afterEach, beforeEach, describe, expect, test as it, vi } from "vitest";
 
 import { PreviouslyResolved } from "./configuration";
-import { ChecksumAlgorithm, RequestChecksumCalculation } from "./constants";
+import { ChecksumAlgorithm, DEFAULT_CHECKSUM_ALGORITHM, RequestChecksumCalculation } from "./constants";
 import { flexibleChecksumsMiddleware } from "./flexibleChecksumsMiddleware";
 import { getChecksumAlgorithmForRequest } from "./getChecksumAlgorithmForRequest";
 import { getChecksumLocationName } from "./getChecksumLocationName";
@@ -34,7 +34,7 @@ describe(flexibleChecksumsMiddleware.name, () => {
   const mockConfig = {
     requestChecksumCalculation: () => Promise.resolve(RequestChecksumCalculation.WHEN_REQUIRED),
   } as PreviouslyResolved;
-  const mockMiddlewareConfig = { input: mockInput, requestChecksumRequired: false };
+  const mockMiddlewareConfig = { requestChecksumRequired: false };
 
   const mockBody = { body: "mockRequestBody" };
   const mockHeaders = { "content-length": 100, "content-encoding": "gzip" };
@@ -78,17 +78,66 @@ describe(flexibleChecksumsMiddleware.name, () => {
         expect(getChecksumAlgorithmForRequest).toHaveBeenCalledTimes(1);
       });
 
-      it("skip if header is already present", async () => {
-        const handler = flexibleChecksumsMiddleware(mockConfig, mockMiddlewareConfig)(mockNext, {});
-        vi.mocked(hasHeaderWithPrefix).mockReturnValue(true);
-
-        await handler(mockArgs);
-
-        expect(hasHeaderWithPrefix).toHaveBeenCalledTimes(1);
-        expect(getChecksumLocationName).not.toHaveBeenCalled();
-        expect(selectChecksumAlgorithmFunction).not.toHaveBeenCalled();
-        expect(hasHeader).not.toHaveBeenCalled();
-        expect(mockNext).toHaveBeenCalledWith(mockArgs);
+      describe("skip if header is already present", async () => {
+        beforeEach(() => {
+          vi.mocked(hasHeaderWithPrefix).mockReturnValue(true);
+        });
+
+        afterEach(() => {
+          expect(hasHeaderWithPrefix).toHaveBeenCalledTimes(1);
+          expect(getChecksumLocationName).not.toHaveBeenCalled();
+          expect(selectChecksumAlgorithmFunction).not.toHaveBeenCalled();
+          expect(hasHeader).not.toHaveBeenCalled();
+        });
+
+        it("with no changes input", async () => {
+          const handler = flexibleChecksumsMiddleware(mockConfig, mockMiddlewareConfig)(mockNext, {});
+
+          await handler(mockArgs);
+          expect(mockNext).toHaveBeenCalledWith(mockArgs);
+        });
+
+        describe("handles input[requestAlgorithmMember]", () => {
+          it("removes if respective checksum header is not present", async () => {
+            const mockRequestAlgorithmMember = "mockRequestAlgorithmMember";
+            const handler = flexibleChecksumsMiddleware(mockConfig, {
+              ...mockMiddlewareConfig,
+              requestAlgorithmMember: mockRequestAlgorithmMember,
+            })(mockNext, {});
+
+            const mockArgsWithInput = {
+              ...mockArgs,
+              input: { [mockRequestAlgorithmMember]: DEFAULT_CHECKSUM_ALGORITHM },
+            };
+            await handler(mockArgsWithInput);
+            expect(mockNext).toHaveBeenCalledWith(mockArgs);
+            expect(mockNext.mock.calls[0][0].input[mockRequestAlgorithmMember]).toBeUndefined();
+          });
+
+          // This means user set the checksum algorithm, as well as the actual checksum.
+          it("retains if respective checksum header is present", async () => {
+            const mockRequestAlgorithmMember = "mockRequestAlgorithmMember";
+            const handler = flexibleChecksumsMiddleware(mockConfig, {
+              ...mockMiddlewareConfig,
+              requestAlgorithmMember: mockRequestAlgorithmMember,
+            })(mockNext, {});
+
+            const mockArgsWithInputAndRequest = {
+              ...mockArgs,
+              input: { [mockRequestAlgorithmMember]: DEFAULT_CHECKSUM_ALGORITHM },
+              request: {
+                ...mockRequest,
+                headers: {
+                  ...mockHeaders,
+                  [`x-amz-checksum-${DEFAULT_CHECKSUM_ALGORITHM.toLowerCase()}`]: mockChecksum,
+                },
+              },
+            };
+            await handler(mockArgsWithInputAndRequest);
+            expect(mockNext).toHaveBeenCalledWith(mockArgsWithInputAndRequest);
+            expect(mockNext.mock.calls[0][0].input[mockRequestAlgorithmMember]).toBe(DEFAULT_CHECKSUM_ALGORITHM);
+          });
+        });
       });
     });
   });

diff --git a/packages/middleware-flexible-checksums/src/flexibleChecksumsMiddleware.ts b/packages/middleware-flexible-checksums/src/flexibleChecksumsMiddleware.ts
@@ -11,7 +11,7 @@ import {
 } from "@smithy/types";
 
 import { PreviouslyResolved } from "./configuration";
-import { ChecksumAlgorithm } from "./constants";
+import { ChecksumAlgorithm, DEFAULT_CHECKSUM_ALGORITHM } from "./constants";
 import { getChecksumAlgorithmForRequest } from "./getChecksumAlgorithmForRequest";
 import { getChecksumLocationName } from "./getChecksumLocationName";
 import { hasHeader } from "./hasHeader";
@@ -60,15 +60,25 @@ export const flexibleChecksumsMiddleware =
       return next(args);
     }
 
-    if (hasHeaderWithPrefix("x-amz-checksum-", args.request.headers)) {
+    const { request, input } = args;
+    const { body: requestBody, headers } = request;
+    const { requestChecksumRequired, requestAlgorithmMember } = middlewareConfig;
+
+    if (hasHeaderWithPrefix("x-amz-checksum-", headers)) {
+      // Remove input[requestAlgorithmMember] and header, if it was added by flexibleChecksumsInputMiddleware
+      if (
+        requestAlgorithmMember &&
+        input[requestAlgorithmMember] === DEFAULT_CHECKSUM_ALGORITHM &&
+        !headers[`x-amz-checksum-${DEFAULT_CHECKSUM_ALGORITHM.toLowerCase()}`]
+      ) {
+        delete input[requestAlgorithmMember];
+        delete headers["x-amz-sdk-checksum-algorithm"];
+      }
       return next(args);
     }
 
-    const { request, input } = args;
-    const { body: requestBody, headers } = request;
     const { base64Encoder, streamHasher } = config;
     const requestChecksumCalculation = await config.requestChecksumCalculation();
-    const { requestChecksumRequired, requestAlgorithmMember } = middlewareConfig;
 
     const checksumAlgorithm = getChecksumAlgorithmForRequest(input, {
       requestChecksumRequired,