fix(middleware-flexible-checksums): split stream when validating response checksum (#5126)

Author: kuhe

packages/middleware-flexible-checksums/package.json

@@ -14,6 +14,9 @@
   },
   "main": "./dist-cjs/index.js",
   "module": "./dist-es/index.js",
+  "browser": {
+    "./dist-es/streams/create-read-stream-on-buffer": "./dist-es/streams/create-read-stream-on-buffer.browser"
+  },
   "types": "./dist-types/index.d.ts",
   "author": {
     "name": "AWS SDK for JavaScript Team",
@@ -31,6 +34,7 @@
     "tslib": "^2.5.0"
   },
   "devDependencies": {
+    "@smithy/node-http-handler": "^2.0.3",
     "concurrently": "7.0.0",
     "downlevel-dts": "0.10.1",
     "rimraf": "3.0.2",

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

@@ -4,6 +4,7 @@ import {
   Encoder,
   GetAwsChunkedEncodingStream,
   HashConstructor,
+  StreamCollector,
   StreamHasher,
 } from "@smithy/types";
 
@@ -47,4 +48,9 @@ export interface PreviouslyResolved {
    * @internal
    */
   streamHasher: StreamHasher<any>;
+
+  /**
+   * Collects streams into buffers.
+   */
+  streamCollector: StreamCollector;
 }

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

@@ -1,9 +1,10 @@
 import { HttpRequest } from "@smithy/protocol-http";
-import { BuildHandlerArguments } from "@smithy/types";
+import { BuildHandlerArguments, DeserializeHandlerArguments } from "@smithy/types";
 
 import { PreviouslyResolved } from "./configuration";
 import { ChecksumAlgorithm } from "./constants";
 import { flexibleChecksumsMiddleware } from "./flexibleChecksumsMiddleware";
+import { flexibleChecksumsResponseMiddleware } from "./flexibleChecksumsResponseMiddleware";
 import { getChecksumAlgorithmForRequest } from "./getChecksumAlgorithmForRequest";
 import { getChecksumLocationName } from "./getChecksumLocationName";
 import { hasHeader } from "./hasHeader";
@@ -194,14 +195,14 @@ describe(flexibleChecksumsMiddleware.name, () => {
     const mockInput = { [mockRequestValidationModeMember]: "ENABLED" };
     const mockResponseAlgorithms = ["ALGO1", "ALGO2"];
 
-    const handler = flexibleChecksumsMiddleware(mockConfig, {
+    const responseHandler = flexibleChecksumsResponseMiddleware(mockConfig, {
       ...mockMiddlewareConfig,
       input: mockInput,
       requestValidationModeMember: mockRequestValidationModeMember,
       responseAlgorithms: mockResponseAlgorithms,
     })(mockNext, {});
 
-    await handler(mockArgs);
+    await responseHandler({ ...mockArgs, input: mockInput } as DeserializeHandlerArguments<any>);
     expect(validateChecksumFromResponse).toHaveBeenCalledWith(mockResult.response, {
       config: mockConfig,
       responseAlgorithms: mockResponseAlgorithms,

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

@@ -1,4 +1,4 @@
-import { HttpRequest, HttpResponse } from "@smithy/protocol-http";
+import { HttpRequest } from "@smithy/protocol-http";
 import {
   BuildHandler,
   BuildHandlerArguments,
@@ -15,8 +15,10 @@ import { hasHeader } from "./hasHeader";
 import { isStreaming } from "./isStreaming";
 import { selectChecksumAlgorithmFunction } from "./selectChecksumAlgorithmFunction";
 import { stringHasher } from "./stringHasher";
-import { validateChecksumFromResponse } from "./validateChecksumFromResponse";
 
+/**
+ * @internal
+ */
 export const flexibleChecksumsMiddleware =
   (config: PreviouslyResolved, middlewareConfig: FlexibleChecksumsMiddlewareConfig): BuildMiddleware<any, any> =>
   <Output extends MetadataBearer>(next: BuildHandler<any, Output>): BuildHandler<any, Output> =>
@@ -78,14 +80,5 @@ export const flexibleChecksumsMiddleware =
       },
     });
 
-    const { requestValidationModeMember, responseAlgorithms } = middlewareConfig;
-    // @ts-ignore Element implicitly has an 'any' type for input[requestValidationModeMember]
-    if (requestValidationModeMember && input[requestValidationModeMember] === "ENABLED") {
-      await validateChecksumFromResponse(result.response as HttpResponse, {
-        config,
-        responseAlgorithms,
-      });
-    }
-
     return result;
   };

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

@@ -0,0 +1,68 @@
+import { HttpRequest, HttpResponse } from "@smithy/protocol-http";
+import {
+  DeserializeHandler,
+  DeserializeHandlerArguments,
+  DeserializeHandlerOutput,
+  DeserializeMiddleware,
+  MetadataBearer,
+  RelativeMiddlewareOptions,
+} from "@smithy/types";
+
+import { PreviouslyResolved } from "./configuration";
+import { FlexibleChecksumsMiddlewareConfig } from "./getFlexibleChecksumsPlugin";
+import { isStreaming } from "./isStreaming";
+import { createReadStreamOnBuffer } from "./streams/create-read-stream-on-buffer";
+import { validateChecksumFromResponse } from "./validateChecksumFromResponse";
+
+/**
+ * @internal
+ */
+export const flexibleChecksumsResponseMiddlewareOptions: RelativeMiddlewareOptions = {
+  name: "flexibleChecksumsResponseMiddleware",
+  toMiddleware: "deserializerMiddleware",
+  relation: "after",
+  tags: ["BODY_CHECKSUM"],
+  override: true,
+};
+
+/**
+ * @internal
+ *
+ * The validation counterpart to the flexibleChecksumsMiddleware.
+ */
+export const flexibleChecksumsResponseMiddleware =
+  (config: PreviouslyResolved, middlewareConfig: FlexibleChecksumsMiddlewareConfig): DeserializeMiddleware<any, any> =>
+  <Output extends MetadataBearer>(next: DeserializeHandler<any, Output>): DeserializeHandler<any, Output> =>
+  async (args: DeserializeHandlerArguments<any>): Promise<DeserializeHandlerOutput<Output>> => {
+    if (!HttpRequest.isInstance(args.request)) {
+      return next(args);
+    }
+
+    const input = args.input;
+    const result = await next(args);
+
+    const response = result.response as HttpResponse;
+    let collectedStream: Uint8Array | undefined = undefined;
+
+    const { requestValidationModeMember, responseAlgorithms } = middlewareConfig;
+    // @ts-ignore Element implicitly has an 'any' type for input[requestValidationModeMember]
+    if (requestValidationModeMember && input[requestValidationModeMember] === "ENABLED") {
+      const isStreamingBody = isStreaming(response.body);
+
+      if (isStreamingBody) {
+        collectedStream = await config.streamCollector(response.body);
+        response.body = createReadStreamOnBuffer(collectedStream);
+      }
+
+      await validateChecksumFromResponse(result.response as HttpResponse, {
+        config,
+        responseAlgorithms,
+      });
+
+      if (isStreamingBody && collectedStream) {
+        response.body = createReadStreamOnBuffer(collectedStream);
+      }
+    }
+
+    return result;
+  };

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

@@ -2,6 +2,10 @@ import { BuildHandlerOptions, Pluggable } from "@smithy/types";
 
 import { PreviouslyResolved } from "./configuration";
 import { flexibleChecksumsMiddleware } from "./flexibleChecksumsMiddleware";
+import {
+  flexibleChecksumsResponseMiddleware,
+  flexibleChecksumsResponseMiddlewareOptions,
+} from "./flexibleChecksumsResponseMiddleware";
 
 export const flexibleChecksumsMiddlewareOptions: BuildHandlerOptions = {
   name: "flexibleChecksumsMiddleware",
@@ -45,5 +49,9 @@ export const getFlexibleChecksumsPlugin = (
 ): Pluggable<any, any> => ({
   applyToStack: (clientStack) => {
     clientStack.add(flexibleChecksumsMiddleware(config, middlewareConfig), flexibleChecksumsMiddlewareOptions);
+    clientStack.addRelativeTo(
+      flexibleChecksumsResponseMiddleware(config, middlewareConfig),
+      flexibleChecksumsResponseMiddlewareOptions
+    );
   },
 });

packages/middleware-flexible-checksums/src/streams/create-read-stream-on-buffer.browser.ts

@@ -0,0 +1,6 @@
+/**
+ * Convert a buffer to a readable stream.
+ */
+export function createReadStreamOnBuffer(buffer: Uint8Array): ReadableStream {
+  return new Blob([buffer]).stream();
+}

packages/middleware-flexible-checksums/src/streams/create-read-stream-on-buffer.spec.ts

@@ -0,0 +1,23 @@
+import { streamCollector } from "@smithy/node-http-handler";
+import { toUtf8 } from "@smithy/util-utf8";
+
+import { createReadStreamOnBuffer } from "./create-read-stream-on-buffer";
+
+describe(createReadStreamOnBuffer.name, () => {
+  it("converts a buffer into stream of the same contents", async () => {
+    const buffer = Buffer.from("abcd");
+
+    const stream = createReadStreamOnBuffer(buffer);
+
+    // changing the buffer here also changes the stream since it has not been collected.
+    buffer[0] = "z".charCodeAt(0);
+
+    const bytes = await streamCollector(stream);
+
+    // changing the buffer here does not change the stream since it has been collected.
+    buffer[1] = "z".charCodeAt(0);
+
+    expect(toUtf8(bytes)).not.toEqual("zzcd");
+    expect(toUtf8(bytes)).toEqual("zbcd");
+  });
+});

packages/middleware-flexible-checksums/src/streams/create-read-stream-on-buffer.ts

@@ -0,0 +1,11 @@
+import { Readable, Transform } from "stream";
+
+/**
+ * Convert a buffer to a readable stream.
+ */
+export function createReadStreamOnBuffer(buffer: Uint8Array): Readable {
+  const stream = new Transform();
+  stream.push(buffer);
+  stream.push(null);
+  return stream;
+}