feat(middleware-flexible-checksums): allow custom checksums to be used in responses

Author: kuhe

packages-internal/middleware-flexible-checksums/src/flexibleChecksumsResponseMiddleware.ts

@@ -73,21 +73,32 @@ export const flexibleChecksumsResponseMiddleware =
     // @ts-ignore Element implicitly has an 'any' type for input[requestValidationModeMember]
     if (requestValidationModeMember && input[requestValidationModeMember] === "ENABLED") {
       const { clientName, commandName } = context;
+
+      const customChecksumAlgorithms = Object.keys(config.checksumAlgorithms ?? {}).filter((algorithm: string) => {
+        const responseHeader = getChecksumLocationName(algorithm);
+        return response.headers[responseHeader] !== undefined;
+      });
+      const algoList = getChecksumAlgorithmListForResponse([
+        ...(responseAlgorithms ?? []),
+        ...customChecksumAlgorithms,
+      ]);
+
       const isS3WholeObjectMultipartGetResponseChecksum =
         clientName === "S3Client" &&
         commandName === "GetObjectCommand" &&
-        getChecksumAlgorithmListForResponse(responseAlgorithms).every((algorithm: ChecksumAlgorithm) => {
+        algoList.every((algorithm: ChecksumAlgorithm) => {
           const responseHeader = getChecksumLocationName(algorithm);
           const checksumFromResponse = response.headers[responseHeader];
           return !checksumFromResponse || isChecksumWithPartNumber(checksumFromResponse);
         });
+
       if (isS3WholeObjectMultipartGetResponseChecksum) {
         return result;
       }
 
       await validateChecksumFromResponse(response as HttpResponse, {
         config,
-        responseAlgorithms,
+        responseAlgorithms: algoList,
         logger: context.logger,
       });
     }

packages-internal/middleware-flexible-checksums/src/getChecksumAlgorithmListForResponse.spec.ts

@@ -4,14 +4,16 @@ import { getChecksumAlgorithmListForResponse } from "./getChecksumAlgorithmListF
 import { PRIORITY_ORDER_ALGORITHMS } from "./types";
 
 describe(getChecksumAlgorithmListForResponse.name, () => {
-  const unknownAlgorithm = "UNKNOWNALGO";
+  const u1 = "UNKNOWNALGO1";
+  const u2 = "UNKNOWNALGO2";
+  const u3 = "UNKNOWNALGO3";
 
   it("returns empty if responseAlgorithms is empty", () => {
     expect(getChecksumAlgorithmListForResponse([])).toEqual([]);
   });
 
-  it("returns empty if contents of responseAlgorithms is not in priority order", () => {
-    expect(getChecksumAlgorithmListForResponse([unknownAlgorithm])).toEqual([]);
+  it("returns unknown algorithms in their existing order if no priority information is available", () => {
+    expect(getChecksumAlgorithmListForResponse([u1, u3, u2])).toEqual([u1, u3, u2]);
   });
 
   describe("returns list as per priority order", () => {
@@ -38,12 +40,14 @@ describe(getChecksumAlgorithmListForResponse.name, () => {
     );
   });
 
-  it("ignores algorithms not present in priority list", () => {
-    expect(getChecksumAlgorithmListForResponse([unknownAlgorithm, ...PRIORITY_ORDER_ALGORITHMS].reverse())).toEqual(
-      PRIORITY_ORDER_ALGORITHMS
-    );
-    expect(getChecksumAlgorithmListForResponse([...PRIORITY_ORDER_ALGORITHMS, unknownAlgorithm].reverse())).toEqual(
-      PRIORITY_ORDER_ALGORITHMS
-    );
+  it("does not ignore algorithms not present in the priority list. However, they receive lowest priority.", () => {
+    expect(getChecksumAlgorithmListForResponse([u1, ...PRIORITY_ORDER_ALGORITHMS].reverse())).toEqual([
+      ...PRIORITY_ORDER_ALGORITHMS,
+      u1,
+    ]);
+    expect(getChecksumAlgorithmListForResponse([...PRIORITY_ORDER_ALGORITHMS, u1].reverse())).toEqual([
+      ...PRIORITY_ORDER_ALGORITHMS,
+      u1,
+    ]);
   });
 });

packages-internal/middleware-flexible-checksums/src/getChecksumAlgorithmListForResponse.ts

@@ -1,5 +1,5 @@
 import type { ChecksumAlgorithm } from "./constants";
-import { CLIENT_SUPPORTED_ALGORITHMS, PRIORITY_ORDER_ALGORITHMS } from "./types";
+import { PRIORITY_ORDER_ALGORITHMS } from "./types";
 
 /**
  * Returns the priority array of algorithm to use to verify checksum and names
@@ -8,12 +8,16 @@ import { CLIENT_SUPPORTED_ALGORITHMS, PRIORITY_ORDER_ALGORITHMS } from "./types"
 export const getChecksumAlgorithmListForResponse = (responseAlgorithms: string[] = []): ChecksumAlgorithm[] => {
   const validChecksumAlgorithms: ChecksumAlgorithm[] = [];
 
-  for (const algorithm of PRIORITY_ORDER_ALGORITHMS) {
-    if (!responseAlgorithms.includes(algorithm) || !CLIENT_SUPPORTED_ALGORITHMS.includes(algorithm)) {
-      continue;
+  let i = PRIORITY_ORDER_ALGORITHMS.length;
+
+  for (const algorithm of responseAlgorithms) {
+    const priority = PRIORITY_ORDER_ALGORITHMS.indexOf(algorithm as ChecksumAlgorithm);
+    if (priority !== -1) {
+      validChecksumAlgorithms[priority] = algorithm as ChecksumAlgorithm;
+    } else {
+      validChecksumAlgorithms[i++] = algorithm as ChecksumAlgorithm;
     }
-    validChecksumAlgorithms.push(algorithm as ChecksumAlgorithm);
   }
 
-  return validChecksumAlgorithms;
+  return validChecksumAlgorithms.filter(Boolean);
 };

packages-internal/middleware-flexible-checksums/src/middleware-flexible-checksums.integ.spec.ts

@@ -1,7 +1,12 @@
 import { requireRequestsFrom } from "@aws-sdk/aws-util-test/src";
+import type { S3ExtensionConfiguration } from "@aws-sdk/client-s3";
 import { ChecksumAlgorithm, S3 } from "@aws-sdk/client-s3";
 import type { HttpHandler, HttpRequest } from "@smithy/protocol-http";
 import { HttpResponse } from "@smithy/protocol-http";
+import type { Checksum } from "@smithy/types";
+import { toBase64 } from "@smithy/util-base64";
+import { ChecksumStream } from "@smithy/util-stream";
+import { fromUtf8 } from "@smithy/util-utf8";
 import { Readable, Transform } from "node:stream";
 import { describe, expect, test as it } from "vitest";
 
@@ -236,4 +241,111 @@ describe("middleware-flexible-checksums", () => {
       });
     });
   });
+
+  describe("Novel checksums", () => {
+    /**
+     * This highly performant checksum algorithm always returns the bytes of the string "Hello, world."
+     * and the number of bytes seen.
+     */
+    class HelloWorldChecksum implements Checksum {
+      private hash = "Hello, world.";
+      private bytesSeen = 0;
+
+      public constructor() {}
+
+      public reset(): void {
+        this.bytesSeen = 0;
+      }
+
+      public update(data: Uint8Array): void {
+        this.bytesSeen += data.byteLength;
+      }
+
+      async digest(): Promise<Uint8Array> {
+        return fromUtf8(this.hash + this.bytesSeen);
+      }
+    }
+
+    class HelloWorldChecksumExtension {
+      configure(extensions: S3ExtensionConfiguration) {
+        extensions.addChecksumAlgorithm({
+          algorithmId() {
+            return "HELLOWORLD";
+          },
+          checksumConstructor() {
+            return HelloWorldChecksum;
+          },
+        });
+      }
+    }
+
+    describe("novel request checksum", () => {
+      it("should send a request with the novel checksum in the header", async () => {
+        const s3 = new S3({
+          credentials: {
+            accessKeyId: "INTEG",
+            secretAccessKey: "INTEG",
+          },
+          extensions: [new HelloWorldChecksumExtension()],
+        });
+        requireRequestsFrom(s3).toMatch({
+          headers: {
+            "x-amz-checksum-helloworld": toBase64(fromUtf8("Hello, world.2")),
+          },
+        });
+
+        await s3.putObject({
+          Bucket: "bucket",
+          Key: "key",
+          Body: "hi",
+          ChecksumAlgorithm: "HELLOWORLD" as any,
+        });
+
+        expect.assertions(1);
+      });
+
+      it("should receive a request and verify the novel checksum", async () => {
+        const s3 = new S3({
+          credentials: {
+            accessKeyId: "INTEG",
+            secretAccessKey: "INTEG",
+          },
+          extensions: [new HelloWorldChecksumExtension()],
+        });
+
+        requireRequestsFrom(s3)
+          .toMatch({
+            hostname: "bucket.s3.us-west-2.amazonaws.com",
+          })
+          .respondWith(
+            new HttpResponse({
+              statusCode: 200,
+              headers: {
+                "x-amz-checksum-helloworld": toBase64(fromUtf8("Hello, world.2")),
+              },
+              body: Readable.from(Buffer.from("hi_extra_bytes")),
+            })
+          );
+
+        const get = await s3.getObject({
+          Bucket: "bucket",
+          Key: "key",
+        });
+
+        expect.assertions(3);
+
+        expect(get.Body).toBeInstanceOf(ChecksumStream);
+        try {
+          await get.Body?.transformToByteArray();
+        } catch (e) {
+          const [ex, ac] = [toBase64(fromUtf8("Hello, world.2")), toBase64(fromUtf8("Hello, world.14"))];
+          expect(e.message).toEqual(
+            `
+Checksum mismatch: expected "${ex}" but received "${ac}" in response header "x-amz-checksum-helloworld".
+`.trim()
+          );
+        }
+      });
+    });
+  });
 });