chore(middleware-flexible-checksums): delay checksum computation till stream is read

Author: trivikr

packages/middleware-flexible-checksums/package.json

@@ -40,6 +40,7 @@
     "@smithy/types": "^3.6.0",
     "@smithy/util-middleware": "^3.0.8",
     "@smithy/util-utf8": "^3.0.0",
+    "@smithy/util-stream": "^3.2.1",
     "tslib": "^2.6.2"
   },
   "devDependencies": {

packages/middleware-flexible-checksums/src/getChecksum.spec.ts

@@ -1,15 +1,12 @@
 import { afterEach, beforeEach, describe, expect, test as it, vi } from "vitest";
 
 import { getChecksum } from "./getChecksum";
-import { isStreaming } from "./isStreaming";
 import { stringHasher } from "./stringHasher";
 
-vi.mock("./isStreaming");
 vi.mock("./stringHasher");
 
 describe(getChecksum.name, () => {
   const mockOptions = {
-    streamHasher: vi.fn(),
     checksumAlgorithmFn: vi.fn(),
     base64Encoder: vi.fn(),
   };
@@ -26,21 +23,10 @@ describe(getChecksum.name, () => {
     vi.clearAllMocks();
   });
 
-  it("gets checksum from streamHasher if body is streaming", async () => {
-    vi.mocked(isStreaming).mockReturnValue(true);
-    mockOptions.streamHasher.mockResolvedValue(mockRawOutput);
-    const checksum = await getChecksum(mockBody, mockOptions);
-    expect(checksum).toEqual(mockOutput);
-    expect(stringHasher).not.toHaveBeenCalled();
-    expect(mockOptions.streamHasher).toHaveBeenCalledWith(mockOptions.checksumAlgorithmFn, mockBody);
-  });
-
-  it("gets checksum from stringHasher if body is not streaming", async () => {
-    vi.mocked(isStreaming).mockReturnValue(false);
+  it("gets checksum from stringHasher", async () => {
     vi.mocked(stringHasher).mockResolvedValue(mockRawOutput);
     const checksum = await getChecksum(mockBody, mockOptions);
     expect(checksum).toEqual(mockOutput);
-    expect(mockOptions.streamHasher).not.toHaveBeenCalled();
     expect(stringHasher).toHaveBeenCalledWith(mockOptions.checksumAlgorithmFn, mockBody);
   });
 });

packages/middleware-flexible-checksums/src/getChecksum.ts

@@ -1,18 +1,11 @@
-import { ChecksumConstructor, Encoder, HashConstructor, StreamHasher } from "@smithy/types";
+import { ChecksumConstructor, Encoder, HashConstructor } from "@smithy/types";
 
-import { isStreaming } from "./isStreaming";
 import { stringHasher } from "./stringHasher";
 
 export interface GetChecksumDigestOptions {
-  streamHasher: StreamHasher<any>;
   checksumAlgorithmFn: ChecksumConstructor | HashConstructor;
   base64Encoder: Encoder;
 }
 
-export const getChecksum = async (
-  body: unknown,
-  { streamHasher, checksumAlgorithmFn, base64Encoder }: GetChecksumDigestOptions
-) => {
-  const digest = isStreaming(body) ? streamHasher(checksumAlgorithmFn, body) : stringHasher(checksumAlgorithmFn, body);
-  return base64Encoder(await digest);
-};
+export const getChecksum = async (body: unknown, { checksumAlgorithmFn, base64Encoder }: GetChecksumDigestOptions) =>
+  base64Encoder(await stringHasher(checksumAlgorithmFn, body));

packages/middleware-flexible-checksums/src/validateChecksumFromResponse.ts

@@ -1,10 +1,13 @@
 import { HttpResponse } from "@smithy/protocol-http";
+import { Checksum } from "@smithy/types";
+import { createChecksumStream } from "@smithy/util-stream";
 
 import { PreviouslyResolved } from "./configuration";
 import { ChecksumAlgorithm } from "./constants";
 import { getChecksum } from "./getChecksum";
 import { getChecksumAlgorithmListForResponse } from "./getChecksumAlgorithmListForResponse";
 import { getChecksumLocationName } from "./getChecksumLocationName";
+import { isStreaming } from "./isStreaming";
 import { selectChecksumAlgorithmFunction } from "./selectChecksumAlgorithmFunction";
 
 export interface ValidateChecksumFromResponseOptions {
@@ -29,9 +32,20 @@ export const validateChecksumFromResponse = async (
     const checksumFromResponse = responseHeaders[responseHeader];
     if (checksumFromResponse) {
       const checksumAlgorithmFn = selectChecksumAlgorithmFunction(algorithm as ChecksumAlgorithm, config);
-      const { streamHasher, base64Encoder } = config;
-      const checksum = await getChecksum(responseBody, { streamHasher, checksumAlgorithmFn, base64Encoder });
+      const { base64Encoder } = config;
 
+      if (isStreaming(responseBody)) {
+        createChecksumStream({
+          expectedChecksum: checksumFromResponse,
+          checksumSourceLocation: responseHeader,
+          checksum: new checksumAlgorithmFn() as Checksum,
+          source: responseBody,
+          base64Encoder,
+        });
+        return;
+      }
+
+      const checksum = await getChecksum(responseBody, { checksumAlgorithmFn, base64Encoder });
       if (checksum === checksumFromResponse) {
         // The checksum for response payload is valid.
         break;