From 4efdb7efb9ba9c1c5ab6be6eda5a4bcf637b7e96 Mon Sep 17 00:00:00 2001
From: Ran Vaknin <RanVaknin@users.noreply.github.com>
Date: Thu, 11 Jan 2024 21:17:53 +0000
Subject: [PATCH] fix(middleware-ssec): add logic to handle string input as
 specified by api model

---
 packages/middleware-ssec/src/index.spec.ts | 51 ++++++++++++++++++++--
 packages/middleware-ssec/src/index.ts      | 32 ++++++++------
 2 files changed, 66 insertions(+), 17 deletions(-)

diff --git a/packages/middleware-ssec/src/index.spec.ts b/packages/middleware-ssec/src/index.spec.ts
index a333412e126c..d8818024b129 100644
--- a/packages/middleware-ssec/src/index.spec.ts
+++ b/packages/middleware-ssec/src/index.spec.ts
@@ -1,11 +1,13 @@
 import { ChecksumConstructor } from "@smithy/types";
+import * as crypto from "crypto";
 
 import { ssecMiddleware } from "./";
 
 describe("ssecMiddleware", () => {
   const next = jest.fn();
   const decoder = jest.fn().mockResolvedValue(new Uint8Array(0));
-  const encoder = jest.fn().mockReturnValue("base64");
+  const encoder1 = jest.fn();
+  const encoder2 = jest.fn();
   const mockHashUpdate = jest.fn();
   const mockHashReset = jest.fn();
   const mockHashDigest = jest.fn().mockReturnValue(new Uint8Array(0));
@@ -17,13 +19,53 @@ describe("ssecMiddleware", () => {
   beforeEach(() => {
     next.mockClear();
     decoder.mockClear();
-    encoder.mockClear();
+    encoder1.mockClear();
+    encoder2.mockClear();
     mockHashUpdate.mockClear();
     mockHashDigest.mockClear();
     mockHashReset.mockClear();
   });
 
   it("should base64 encode input keys and set respective MD5 inputs", async () => {
+    encoder1.mockReturnValue("/+JF8FMG8UVMWSaNz0s6Wg==");
+    const key = "TestKey123";
+    const binaryRepresentationOfKey = Buffer.from(key);
+    const base64Key = binaryRepresentationOfKey.toString("base64");
+    const md5Hash = crypto.createHash("md5").update(binaryRepresentationOfKey).digest();
+    const base64Md5Hash = Buffer.from(md5Hash).toString("base64");
+
+    const args = {
+      input: {
+        SSECustomerKey: base64Key,
+        CopySourceSSECustomerKey: base64Key,
+      },
+    };
+
+    const handler = ssecMiddleware({
+      base64Encoder: encoder1,
+      utf8Decoder: decoder,
+      md5: MockHash,
+    })(next, {} as any);
+
+    await handler(args);
+
+    expect(next.mock.calls.length).toBe(1);
+    expect(next).toHaveBeenCalledWith({
+      input: {
+        SSECustomerKey: base64Key,
+        SSECustomerKeyMD5: base64Md5Hash,
+        CopySourceSSECustomerKey: base64Key,
+        CopySourceSSECustomerKeyMD5: base64Md5Hash,
+      },
+    });
+    expect(decoder.mock.calls.length).toBe(0);
+    expect(encoder1.mock.calls.length).toBe(2);
+    expect(mockHashUpdate.mock.calls.length).toBe(2);
+    expect(mockHashDigest.mock.calls.length).toBe(2);
+    encoder1.mockClear();
+  });
+  it("should base64 encode input keys and set respective MD5 inputs", async () => {
+    encoder2.mockReturnValue("base64");
     const args = {
       input: {
         SSECustomerKey: "foo",
@@ -32,7 +74,7 @@ describe("ssecMiddleware", () => {
     };
 
     const handler = ssecMiddleware({
-      base64Encoder: encoder,
+      base64Encoder: encoder2,
       utf8Decoder: decoder,
       md5: MockHash,
     })(next, {} as any);
@@ -49,8 +91,9 @@ describe("ssecMiddleware", () => {
       },
     });
     expect(decoder.mock.calls.length).toBe(2);
-    expect(encoder.mock.calls.length).toBe(4);
+    expect(encoder2.mock.calls.length).toBe(4);
     expect(mockHashUpdate.mock.calls.length).toBe(2);
     expect(mockHashDigest.mock.calls.length).toBe(2);
+    encoder2.mockClear();
   });
 });
diff --git a/packages/middleware-ssec/src/index.ts b/packages/middleware-ssec/src/index.ts
index 63c245eb988a..ec3add3d7ae7 100644
--- a/packages/middleware-ssec/src/index.ts
+++ b/packages/middleware-ssec/src/index.ts
@@ -21,7 +21,7 @@ interface PreviouslyResolved {
 export function ssecMiddleware(options: PreviouslyResolved): InitializeMiddleware<any, any> {
   return <Output extends MetadataBearer>(next: InitializeHandler<any, Output>): InitializeHandler<any, Output> =>
     async (args: InitializeHandlerArguments<any>): Promise<InitializeHandlerOutput<Output>> => {
-      let input = { ...args.input };
+      const input = { ...args.input };
       const properties = [
         {
           target: "SSECustomerKey",
@@ -36,19 +36,25 @@ export function ssecMiddleware(options: PreviouslyResolved): InitializeMiddlewar
       for (const prop of properties) {
         const value: SourceData | undefined = (input as any)[prop.target];
         if (value) {
-          const valueView: Uint8Array = ArrayBuffer.isView(value)
-            ? new Uint8Array(value.buffer, value.byteOffset, value.byteLength)
-            : typeof value === "string"
-            ? options.utf8Decoder(value)
-            : new Uint8Array(value);
-          const encoded = options.base64Encoder(valueView);
+          let valueForHash: Uint8Array;
+          if (typeof value === "string") {
+            const isBase64Encoded = /^(?:[A-Za-z0-9+/]{4})*([A-Za-z0-9+/]{2}==|[A-Za-z0-9+/]{3}=)?$/.test(value);
+            if (isBase64Encoded) {
+              valueForHash = new Uint8Array(Buffer.from(value, "base64"));
+            } else {
+              valueForHash = options.utf8Decoder(value);
+              input[prop.target] = options.base64Encoder(valueForHash);
+            }
+          } else {
+            valueForHash = ArrayBuffer.isView(value)
+              ? new Uint8Array(value.buffer, value.byteOffset, value.byteLength)
+              : new Uint8Array(value);
+            input[prop.target] = options.base64Encoder(valueForHash);
+          }
+
           const hash = new options.md5();
-          hash.update(valueView);
-          input = {
-            ...(input as any),
-            [prop.target]: encoded,
-            [prop.hash]: options.base64Encoder(await hash.digest()),
-          };
+          hash.update(valueForHash);
+          input[prop.hash] = options.base64Encoder(await hash.digest());
         }
       }