From 969be814e10cf17d8c95d66ee769e9f0d4c3e0ca Mon Sep 17 00:00:00 2001
From: "Kamat, Trivikram" <16024985+trivikr@users.noreply.github.com>
Date: Fri, 22 Nov 2024 20:35:34 +0000
Subject: [PATCH] chore: leaves input.requestAlgorithmMember if checksum header
 is set

---
 .../flexibleChecksumsInputMiddleware.spec.ts  | 47 ++++++++++++++-----
 .../src/flexibleChecksumsInputMiddleware.ts   | 11 +++--
 2 files changed, 44 insertions(+), 14 deletions(-)

diff --git a/packages/middleware-flexible-checksums/src/flexibleChecksumsInputMiddleware.spec.ts b/packages/middleware-flexible-checksums/src/flexibleChecksumsInputMiddleware.spec.ts
index 46aa6c395cb4..94d050a11fb9 100644
--- a/packages/middleware-flexible-checksums/src/flexibleChecksumsInputMiddleware.spec.ts
+++ b/packages/middleware-flexible-checksums/src/flexibleChecksumsInputMiddleware.spec.ts
@@ -1,11 +1,13 @@
 import { setFeature } from "@aws-sdk/core";
-import { afterEach, describe, expect, test as it, vi } from "vitest";
+import { afterEach, beforeEach, describe, expect, test as it, vi } from "vitest";
 
 import { PreviouslyResolved } from "./configuration";
 import { DEFAULT_CHECKSUM_ALGORITHM, RequestChecksumCalculation, ResponseChecksumValidation } from "./constants";
 import { flexibleChecksumsInputMiddleware } from "./flexibleChecksumsInputMiddleware";
+import { hasHeaderWithPrefix } from "./hasHeaderWithPrefix";
 
 vi.mock("@aws-sdk/core");
+vi.mock("./hasHeaderWithPrefix");
 
 describe(flexibleChecksumsInputMiddleware.name, () => {
   const mockNext = vi.fn();
@@ -19,6 +21,10 @@ describe(flexibleChecksumsInputMiddleware.name, () => {
     responseChecksumValidation: () => Promise.resolve(ResponseChecksumValidation.WHEN_SUPPORTED),
   } as PreviouslyResolved;
 
+  beforeEach(() => {
+    vi.mocked(hasHeaderWithPrefix).mockReturnValue(false);
+  });
+
   afterEach(() => {
     expect(mockNext).toHaveBeenCalledTimes(1);
     vi.clearAllMocks();
@@ -36,8 +42,11 @@ describe(flexibleChecksumsInputMiddleware.name, () => {
           mockNext,
           {}
         );
-        await handler({ input: {} });
-        expect(mockNext).toHaveBeenCalledWith({ input: { [mockRequestAlgorithmMember]: DEFAULT_CHECKSUM_ALGORITHM } });
+        await handler({ input: {}, request: {} });
+        expect(mockNext).toHaveBeenCalledWith({
+          input: { [mockRequestAlgorithmMember]: DEFAULT_CHECKSUM_ALGORITHM },
+          request: {},
+        });
       });
 
       it("requestChecksumRequired is set to true", async () => {
@@ -51,8 +60,11 @@ describe(flexibleChecksumsInputMiddleware.name, () => {
           requestChecksumRequired: true,
         })(mockNext, {});
 
-        await handler({ input: {} });
-        expect(mockNext).toHaveBeenCalledWith({ input: { [mockRequestAlgorithmMember]: DEFAULT_CHECKSUM_ALGORITHM } });
+        await handler({ input: {}, request: {} });
+        expect(mockNext).toHaveBeenCalledWith({
+          input: { [mockRequestAlgorithmMember]: DEFAULT_CHECKSUM_ALGORITHM },
+          request: {},
+        });
       });
     });
   });
@@ -82,8 +94,21 @@ describe(flexibleChecksumsInputMiddleware.name, () => {
         mockConfigReqChecksumCalculationWhenRequired,
         mockMiddlewareConfigWithRequestAlgorithmMember
       )(mockNext, {});
-      await handler({ input: {} });
-      expect(mockNext).toHaveBeenCalledWith({ input: {} });
+      await handler({ input: {}, request: {} });
+      expect(mockNext).toHaveBeenCalledWith({ input: {}, request: {} });
+    });
+
+    it("if checksum header is set", async () => {
+      const mockArgs = { input: {}, request: { key: "value" } };
+      vi.mocked(hasHeaderWithPrefix).mockReturnValue(true);
+
+      const handler = flexibleChecksumsInputMiddleware(mockConfig, {
+        ...mockMiddlewareConfigWithRequestAlgorithmMember,
+        requestChecksumRequired: true,
+      })(mockNext, {});
+
+      await handler(mockArgs);
+      expect(mockNext).toHaveBeenCalledWith(mockArgs);
     });
   });
 
@@ -97,13 +122,13 @@ describe(flexibleChecksumsInputMiddleware.name, () => {
         mockConfig,
         mockMiddlewareConfigWithMockRequestValidationModeMember
       )(mockNext, {});
-      await handler({ input: {} });
-      expect(mockNext).toHaveBeenCalledWith({ input: { [mockRequestValidationModeMember]: "ENABLED" } });
+      await handler({ input: {}, request: {} });
+      expect(mockNext).toHaveBeenCalledWith({ input: { [mockRequestValidationModeMember]: "ENABLED" }, request: {} });
     });
   });
 
   describe("leaves input.requestValidationModeMember", () => {
-    const mockArgs = { input: {} };
+    const mockArgs = { input: {}, request: {} };
 
     it("when requestValidationModeMember is not defined", async () => {
       const handler = flexibleChecksumsInputMiddleware(mockConfig, mockMiddlewareConfig)(mockNext, {});
@@ -157,7 +182,7 @@ describe(flexibleChecksumsInputMiddleware.name, () => {
       } as PreviouslyResolved;
 
       const handler = flexibleChecksumsInputMiddleware(mockConfigOverride, mockMiddlewareConfig)(mockNext, {});
-      await handler({ input: {} });
+      await handler({ input: {}, request: {} });
 
       expect(setFeature).toHaveBeenCalledTimes(2);
       if (configKey === "requestChecksumCalculation") {
diff --git a/packages/middleware-flexible-checksums/src/flexibleChecksumsInputMiddleware.ts b/packages/middleware-flexible-checksums/src/flexibleChecksumsInputMiddleware.ts
index a3095f9fcad7..81cf1782d8e4 100644
--- a/packages/middleware-flexible-checksums/src/flexibleChecksumsInputMiddleware.ts
+++ b/packages/middleware-flexible-checksums/src/flexibleChecksumsInputMiddleware.ts
@@ -1,4 +1,5 @@
 import { setFeature } from "@aws-sdk/core";
+import { HttpRequest } from "@smithy/protocol-http";
 import {
   HandlerExecutionContext,
   MetadataBearer,
@@ -11,6 +12,7 @@ import {
 
 import { PreviouslyResolved } from "./configuration";
 import { DEFAULT_CHECKSUM_ALGORITHM, RequestChecksumCalculation, ResponseChecksumValidation } from "./constants";
+import { hasHeaderWithPrefix } from "./hasHeaderWithPrefix";
 
 export interface FlexibleChecksumsInputMiddlewareConfig {
   /**
@@ -82,10 +84,13 @@ export const flexibleChecksumsInputMiddleware =
 
     // The value for input member to configure flexible checksum is not set.
     if (requestAlgorithmMember && !input[requestAlgorithmMember]) {
-      // Set requestAlgorithmMember as default checksum algorithm only if request checksum calculation is supported
-      // or request checksum is required.
+      // Set requestAlgorithmMember as default checksum algorithm only if request checksum algorithm is not supported,
+      // and either request checksumcalculation is supported or request checksum is required.
       if (requestChecksumCalculation === RequestChecksumCalculation.WHEN_SUPPORTED || requestChecksumRequired) {
-        input[requestAlgorithmMember] = DEFAULT_CHECKSUM_ALGORITHM;
+        const request = args.request as HttpRequest;
+        if (!request.headers || hasHeaderWithPrefix("x-amz-checksum-", request.headers)) {
+          input[requestAlgorithmMember] = DEFAULT_CHECKSUM_ALGORITHM;
+        }
       }
     }