fix: move eventstream handler after signing step

Author: AllanZhengYP

packages/eventstream-handler-node/src/EventStreamPayloadHandler.ts

@@ -9,7 +9,10 @@ import {
   Encoder,
   Decoder,
   HttpRequest,
-  HandlerExecutionContext
+  HandlerExecutionContext,
+  FinalizeHandlerArguments,
+  FinalizeHandler,
+  FinalizeHandlerOutput
 } from "@aws-sdk/types";
 import { EventStreamMarshaller as EventMarshaller } from "@aws-sdk/eventstream-marshaller";
 import { Readable, PassThrough, pipeline } from "stream";
@@ -33,10 +36,10 @@ export class EventStreamPayloadHandler implements IEventStreamPayloadHandler {
   }
 
   async handle<T extends MetadataBearer>(
-    next: BuildHandler<any, T>,
-    args: BuildHandlerArguments<any>,
+    next: FinalizeHandler<any, T>,
+    args: FinalizeHandlerArguments<any>,
     context: HandlerExecutionContext = {} as any
-  ): Promise<BuildHandlerOutput<T>> {
+  ): Promise<FinalizeHandlerOutput<T>> {
     const request = args.request as HttpRequest;
     const { body: payload } = request;
     if (!(payload instanceof Readable)) {
@@ -46,14 +49,19 @@ export class EventStreamPayloadHandler implements IEventStreamPayloadHandler {
     request.body = new PassThrough({
       objectMode: true
     });
-    const result = await next(args);
-    if (!context.signature) {
-      throw new Error(
-        "Initial request signature is not available in eventstream handler."
-      );
+    let result: FinalizeHandlerOutput<any>;
+    try {
+      result = await next(args);
+    } catch (e) {
+      request.body.end();
+      throw e;
     }
+    const match = (request.headers["authorization"] || "").match(
+      /Signature=([\w]+)$/
+    );
+    const priorSignature = (match || [])[1];
     const signingStream = new EventSigningStream({
-      priorSignature: context.signature,
+      priorSignature,
       eventMarshaller: this.eventMarshaller,
       eventSigner: await this.eventSigner()
     });

packages/middleware-eventstream/src/handling-middleware.ts

@@ -1,17 +1,26 @@
-import { BuildMiddleware, BuildHandlerOptions } from "@aws-sdk/types";
+import {
+  BuildMiddleware,
+  BuildHandlerOptions,
+  FinalizeRequestMiddleware,
+  FinalizeRequestHandlerOptions,
+  RelativeLocation
+} from "@aws-sdk/types";
 import { EventStreamResolvedConfig } from "./configuration";
 import { HttpRequest } from "@aws-sdk/protocol-http";
 
 export const eventStreamHandlingMiddleware = (
   options: EventStreamResolvedConfig
-): BuildMiddleware<any, any> => (next, context) => async args => {
+): FinalizeRequestMiddleware<any, any> => (next, context) => async args => {
   const { request } = args;
   if (!HttpRequest.isInstance(request)) return next(args);
   return options.eventStreamPayloadHandler.handle(next, args, context);
 };
 
-export const eventStreamHandlingMiddlewareOptions: BuildHandlerOptions = {
-  step: "build",
+export const eventStreamHandlingMiddlewareOptions: FinalizeRequestHandlerOptions &
+  RelativeLocation<any, any> = {
+  step: "finalizeRequest",
   tags: ["EVENT_STREAM", "SIGNATURE", "HANDLE"],
-  name: "eventStreamHandlingMiddleware"
+  name: "eventStreamHandlingMiddleware",
+  relation: "after",
+  toMiddleware: "awsAuthMiddleware"
 };

packages/middleware-eventstream/src/pluggin.ts

@@ -8,10 +8,6 @@ import {
   eventStreamHeaderMiddleware,
   eventStreamHeaderMiddlewareOptions
 } from "./headers-middleware";
-import {
-  recordSignatureMiddleware,
-  recordSignatureMiddlewareOptions
-} from "./record-signature-middleware";
 
 export const getEventStreamPlugin = (
   options: EventStreamResolvedConfig
@@ -25,9 +21,5 @@ export const getEventStreamPlugin = (
       eventStreamHeaderMiddleware,
       eventStreamHeaderMiddlewareOptions
     );
-    clientStack.addRelativeTo(
-      recordSignatureMiddleware,
-      recordSignatureMiddlewareOptions
-    );
   }
 });

packages/middleware-eventstream/src/record-signature-middleware.ts

@@ -1,10 +1,10 @@
 import { HttpRequest } from "./http";
 import { MetadataBearer } from "./response";
 import {
-  BuildHandler,
-  BuildHandlerArguments,
-  BuildHandlerOutput,
-  HandlerExecutionContext
+  HandlerExecutionContext,
+  FinalizeHandler,
+  FinalizeHandlerArguments,
+  FinalizeHandlerOutput
 } from "./middleware";
 /**
  * An event stream message. The headers and body properties will always be
@@ -106,10 +106,10 @@ export interface EventStreamRequestSigner {
 
 export interface EventStreamPayloadHandler {
   handle: <Input extends object, Output extends MetadataBearer>(
-    next: BuildHandler<Input, Output>,
-    args: BuildHandlerArguments<Input>,
+    next: FinalizeHandler<Input, Output>,
+    args: FinalizeHandlerArguments<Input>,
     context?: HandlerExecutionContext
-  ) => Promise<BuildHandlerOutput<Output>>;
+  ) => Promise<FinalizeHandlerOutput<Output>>;
 }
 
 export interface EventStreamPayloadHandlerProvider {