fix response payload and incorrectly parsing error response

smithy-typescript-codegen/src/main/java/software/amazon/smithy/typescript/codegen/integration/HttpBindingProtocolGenerator.java

@@ -119,6 +119,8 @@ public void generateSharedComponents(GenerationContext context) {
         generateDocumentBodyShapeSerializers(context, serializingDocumentShapes);
         generateDocumentBodyShapeDeserializers(context, deserializingDocumentShapes);
         HttpProtocolGeneratorUtils.generateMetadataDeserializer(context, getApplicationProtocol().getResponseType());
+        HttpProtocolGeneratorUtils.generateCollectBody(context);
+        HttpProtocolGeneratorUtils.generateCollectBodyString(context);
     }
 
     /**
@@ -593,7 +595,7 @@ private void generateOperationDeserializer(
 
         // Write out the error deserialization dispatcher.
         Set<StructureShape> errorShapes = HttpProtocolGeneratorUtils.generateErrorDispatcher(
-                context, operation, responseType, this::writeErrorCodeParser);
+                context, operation, responseType, this::writeErrorCodeParser, this.isErrorCodeInBody());
         deserializingErrorShapes.addAll(errorShapes);
     }
 
@@ -605,11 +607,13 @@ private void generateErrorDeserializer(GenerationContext context, StructureShape
         Symbol errorSymbol = symbolProvider.toSymbol(error);
         String errorDeserMethodName = ProtocolGenerator.getDeserFunctionName(errorSymbol,
                 context.getProtocolName()) + "Response";
+        boolean isBodyParsed = this.isErrorCodeInBody();
 
         writer.openBlock("const $L = async (\n"
-                       + "  output: any,\n"
+                       + "  $L: any,\n"
                        + "  context: __SerdeContext\n"
-                       + "): Promise<$T> => {", "};", errorDeserMethodName, errorSymbol, () -> {
+                       + "): Promise<$T> => {", "};",
+                errorDeserMethodName, isBodyParsed ? "parsedOutput" : "output", errorSymbol, () -> {
             writer.openBlock("const contents: $T = {", "};", errorSymbol, () -> {
                 writer.write("__type: $S,", error.getId().getName());
                 writer.write("$$fault: $S,", error.getTrait(ErrorTrait.class).get().getValue());
@@ -620,7 +624,7 @@ private void generateErrorDeserializer(GenerationContext context, StructureShape
             });
 
             readHeaders(context, error, bindingIndex);
-            List<HttpBinding> documentBindings = readResponseBody(context, error, bindingIndex);
+            List<HttpBinding> documentBindings = readErrorResponseBody(context, error, bindingIndex, isBodyParsed);
             // Track all shapes bound to the document so their deserializers may be generated.
             documentBindings.forEach(binding -> {
                 Shape target = model.expectShape(binding.getMember().getTarget());
@@ -632,6 +636,23 @@ private void generateErrorDeserializer(GenerationContext context, StructureShape
         writer.write("");
     }
 
+    private List<HttpBinding> readErrorResponseBody(
+            GenerationContext context,
+            Shape error,
+            HttpBindingIndex bindingIndex,
+            boolean isBodyParsed
+    ) {
+        TypeScriptWriter writer = context.getWriter();
+        if (isBodyParsed) {
+            // Body is already parsed in error dispatcher, simply assign body to data.
+            writer.write("const data: any = output.body;");
+            return ListUtils.of();
+        } else {
+            // Deserialize response body just like in normal response.
+            return readResponseBody(context, error, bindingIndex);
+        }
+    }
+
     private void readHeaders(
             GenerationContext context,
             Shape operationOrError,
@@ -690,42 +711,44 @@ private List<HttpBinding> readResponseBody(
         documentBindings.sort(Comparator.comparing(HttpBinding::getMemberName));
         List<HttpBinding> payloadBindings = bindingIndex.getResponseBindings(operationOrError, Location.PAYLOAD);
 
+        OperationIndex operationIndex = context.getModel().getKnowledge(OperationIndex.class);
+        StructureShape operationOutputOrError = operationOrError.asStructureShape()
+                .orElseGet(() -> operationIndex.getOutput(operationOrError).orElse(null));
+        boolean hasStreamingComponent = Optional.ofNullable(operationOutputOrError)
+                .map(structure -> structure.getAllMembers().values().stream()
+                        .anyMatch(memberShape -> memberShape.hasTrait(StreamingTrait.class)))
+                .orElse(false);
+
         if (!documentBindings.isEmpty()) {
-            readReponseBodyData(context, operationOrError);
+            // If response has document binding, the body can be parsed to JavaScript object.
+            writer.write("const data: any = await parseBody(output.body, context);");
             deserializeOutputDocument(context, operationOrError, documentBindings);
             return documentBindings;
         }
         if (!payloadBindings.isEmpty()) {
-            readReponseBodyData(context, operationOrError);
             // There can only be one payload binding.
             HttpBinding binding = payloadBindings.get(0);
             Shape target = context.getModel().expectShape(binding.getMember().getTarget());
+            if (hasStreamingComponent) {
+                // If payload is streaming, return raw low-level stream directly.
+                writer.write("const data: any = output.body;");
+            } else if (target instanceof BlobShape) {
+                // If payload is blob, only need to collect stream to binary data(Uint8Array).
+                writer.write("const data: any = await collectBody(output.body, context);");
+            } else if (target instanceof StructureShape || target instanceof UnionShape) {
+                // If body is Structure or Union, they we need to parse the string into JavaScript object.
+                writer.write("const data: any = await parseBody(output.body, context);");
+            } else {
+                // If payload is string, we need to collect body and encode binary to string.
+                writer.write("const data: any = await collectBodyString(output.body, context);");
+            }
             writer.write("contents.$L = $L;", binding.getMemberName(), getOutputValue(context,
                     Location.PAYLOAD, "data", binding.getMember(), target));
             return payloadBindings;
         }
         return ListUtils.of();
     }
 
-    private void readReponseBodyData(GenerationContext context, Shape operationOrError) {
-        TypeScriptWriter writer = context.getWriter();
-        // Prepare response body for deserializing.
-        OperationIndex operationIndex = context.getModel().getKnowledge(OperationIndex.class);
-        StructureShape operationOutputOrError = operationOrError.asStructureShape()
-                .orElseGet(() -> operationIndex.getOutput(operationOrError).orElse(null));
-        boolean hasStreamingComponent = Optional.ofNullable(operationOutputOrError)
-                .map(structure -> structure.getAllMembers().values().stream()
-                        .anyMatch(memberShape -> memberShape.hasTrait(StreamingTrait.class)))
-                .orElse(false);
-        if (hasStreamingComponent) {
-            // For operations with streaming output or errors with streaming body we keep the body intact.
-            writer.write("const data: any = output.body;");
-        } else {
-            // Otherwise, we collect the response body to structured object with parseBody().
-            writer.write("const data: any = await parseBody(output.body, context);");
-        }
-    }
-
     /**
      * Given context and a source of data, generate an output value provider for the
      * shape. This may use native types (like generating a Date for timestamps,)
@@ -906,6 +929,14 @@ private String getNumberOutputParam(Location bindingType, String dataSource, Sha
      */
     protected abstract void writeErrorCodeParser(GenerationContext context);
 
+    /**
+     * A boolean indicates whether body is collected and parsed in error code parser.
+     * If so, each error shape deserializer should not parse body again.
+     *
+     * @return returns whether the error code exists in response body
+     */
+    protected abstract boolean isErrorCodeInBody();
+
     /**
      * Writes the code needed to deserialize the output document of a response.
      *

smithy-typescript-codegen/src/main/java/software/amazon/smithy/typescript/codegen/integration/HttpProtocolGeneratorUtils.java

@@ -108,6 +108,44 @@ static void generateMetadataDeserializer(GenerationContext context, SymbolRefere
         writer.write("");
     }
 
+    /**
+     * Writes a response body stream collector. This function converts the low-level response body stream to
+     * Uint8Array binary data.
+     *
+     * @param context The generation context.
+     */
+    static void generateCollectBody(GenerationContext context) {
+        TypeScriptWriter writer = context.getWriter();
+
+        writer.addImport("SerdeContext", "__SerdeContext", "@aws-sdk/types");
+        writer.write("// Collect low-level response body stream to Uint8Array.");
+        writer.openBlock("const collectBody = (streamBody: any, context: __SerdeContext): Promise<Uint8Array> => {",
+                "};", () -> {
+            writer.write("return context.streamCollector(streamBody) || new Uint8Array();");
+        });
+
+        writer.write("");
+    }
+
+    /**
+     * Writes a function converting the low-level response body stream to utf-8 encoded string. It depends on
+     * response body stream collector{@link #generateCollectBody(GenerationContext)}.
+     *
+     * @param context The generation context
+     */
+    static void generateCollectBodyString(GenerationContext context) {
+        TypeScriptWriter writer = context.getWriter();
+
+        writer.addImport("SerdeContext", "__SerdeContext", "@aws-sdk/types");
+        writer.write("// Encode Uint8Array data into string with utf-8.");
+        writer.openBlock("const collectBodyString = (streamBody: any, context: __SerdeContext): Promise<string> => {",
+                "};", () -> {
+            writer.write("return collectBody(streamBody, context).then(body => context.utf8Encoder(body));");
+        });
+
+        writer.write("");
+    }
+
     /**
      * Writes a function used to dispatch to the proper error deserializer
      * for each error that the operation can return. The generated function
@@ -118,13 +156,15 @@ static void generateMetadataDeserializer(GenerationContext context, SymbolRefere
      * @param operation The operation to generate for.
      * @param responseType The response type for the HTTP protocol.
      * @param errorCodeGenerator A consumer
+     * @param shouldParseErrorBody Flag indicating whether need to parse response body
      * @return A set of all error structure shapes for the operation that were dispatched to.
      */
     static Set<StructureShape> generateErrorDispatcher(
             GenerationContext context,
             OperationShape operation,
             SymbolReference responseType,
-            Consumer<GenerationContext> errorCodeGenerator
+            Consumer<GenerationContext> errorCodeGenerator,
+            boolean shouldParseErrorBody
     ) {
         TypeScriptWriter writer = context.getWriter();
         SymbolProvider symbolProvider = context.getSymbolProvider();
@@ -138,15 +178,15 @@ static Set<StructureShape> generateErrorDispatcher(
                        + "  output: $T,\n"
                        + "  context: __SerdeContext,\n"
                        + "): Promise<$T> {", "}", errorMethodName, responseType, outputType, () -> {
-            writer.write("const data: any = await parseBody(output.body, context);");
-            // We only consume the parsedOutput if we're dispatching, so only generate if we will.
-            if (!operation.getErrors().isEmpty()) {
-                // Create a holding object since we have already parsed the body, but retain the rest of the output.
-                writer.openBlock("const parsedOutput: any = {", "};", () -> {
-                    writer.write("...output,");
-                    writer.write("body: data,");
-                });
-            }
+            // Prepare error response for parsing error code. If error code needs to be parsed from response body
+            // then we collect body and parse it to JS object, otherwise leave the response body as is.
+            writer.openBlock(
+                    "const $L: any = {", "};", shouldParseErrorBody ? "parsedOutput" : "errorOutput",
+                    () -> {
+                        writer.write("...output,");
+                        writer.write("body: $L,",
+                                shouldParseErrorBody ? "await parseBody(output.body, context)" : "output.body");
+                    });
 
             // Error responses must be at least SmithyException and MetadataBearer implementations.
             writer.addImport("SmithyException", "__SmithyException",
@@ -167,17 +207,19 @@ static Set<StructureShape> generateErrorDispatcher(
                             context.getProtocolName()) + "Response";
                     writer.openBlock("case $S:\ncase $S:", "  break;", errorId.getName(), errorId.toString(), () -> {
                         // Dispatch to the error deserialization function.
-                        writer.write("response = await $L(parsedOutput, context);", errorDeserMethodName);
+                        writer.write("response = await $L($L, context);",
+                                errorDeserMethodName, shouldParseErrorBody ? "parsedOutput" : "errorOutput");
                     });
                 });
 
                 // Build a generic error the best we can for ones we don't know about.
                 writer.write("default:").indent()
                         .write("errorCode = errorCode || \"UnknownError\";")
-                        .openBlock("response = {", "};", () -> {
+                        .openBlock("response = {", "} as any;", () -> {
                             writer.write("__type: `$L#$${errorCode}`,", operation.getId().getNamespace());
                             writer.write("$$fault: \"client\",");
                             writer.write("$$metadata: deserializeMetadata(output),");
+                            writer.write("message: await collectBodyString(output.body, context)");
                         }).dedent();
             });
             writer.write("return Promise.reject(Object.assign(new Error(response.__type), response));");

smithy-typescript-codegen/src/main/java/software/amazon/smithy/typescript/codegen/integration/HttpRpcProtocolGenerator.java

@@ -254,7 +254,7 @@ private void generateOperationDeserializer(GenerationContext context, OperationS
 
         // Write out the error deserialization dispatcher.
         Set<StructureShape> errorShapes = HttpProtocolGeneratorUtils.generateErrorDispatcher(
-                context, operation, responseType, this::writeErrorCodeParser);
+                context, operation, responseType, this::writeErrorCodeParser, this.isErrorCodeInBody());
         deserializingErrorShapes.addAll(errorShapes);
     }
 
@@ -311,10 +311,13 @@ private void readResponseBody(GenerationContext context, OperationShape operatio
      * Writes the code that loads an {@code errorCode} String with the content used
      * to dispatch errors to specific serializers.
      *
-     * <p>Three variables will be in scope:
+     * <p>Two variables will be in scope:
      *   <ul>
-     *       <li>{@code output}: a value of the HttpResponse type.</li>
-     *       <li>{@code data}: the contents of the response body.</li>
+     *       <li>{@code errorOutput} or {@code parsedOutput}: a value of the HttpResponse type.
+     *          {@code errorOutput} is a raw HttpResponse whereas {@code parsedOutput} is a HttpResponse type with
+     *          body parsed to JavaScript object.
+     *          The actual value available is determined by {@link #isErrorCodeInBody}
+     *       </li>
      *       <li>{@code context}: the SerdeContext.</li>
      *   </ul>
      *
@@ -328,6 +331,17 @@ private void readResponseBody(GenerationContext context, OperationShape operatio
      */
     protected abstract void writeErrorCodeParser(GenerationContext context);
 
+    /**
+     * Indicates whether body is collected and parsed in error dispatcher.
+     *
+     * <p>If returns true, {@link #writeErrorCodeParser} will have {@code parsedOutput} in scope
+     *
+     * <P>If returns false, {@link #writeErrorCodeParser} will have {@code errorOutput} in scope
+     *
+     * @return returns whether the error code exists in response body
+     */
+    protected abstract boolean isErrorCodeInBody();
+
     /**
      * Writes the code needed to deserialize the output document of a response.
      *