From 85dda37881044e7919c76bbc632d754d1d5db149 Mon Sep 17 00:00:00 2001
From: Jason Del Ponte <delpontej@gmail.com>
Date: Mon, 13 Jul 2020 15:46:26 -0700
Subject: [PATCH] Add setting content-length middleware into operation stack

Adds the ContentLengthMiddleware to the operation's stack, and updates
generated clients and protocol tests to use the content-length check.

Depends on:
* https://github.com/awslabs/smithy-go/pull/108
* https://github.com/awslabs/smithy/pull/491
---
 .../aws/go/codegen/AssembleMiddlewareStack.java       | 11 +++++++++++
 .../aws/go/codegen/RestJsonProtocolGenerator.java     |  1 -
 2 files changed, 11 insertions(+), 1 deletion(-)

diff --git a/codegen/smithy-aws-go-codegen/src/main/java/software/amazon/smithy/aws/go/codegen/AssembleMiddlewareStack.java b/codegen/smithy-aws-go-codegen/src/main/java/software/amazon/smithy/aws/go/codegen/AssembleMiddlewareStack.java
index 27a97ae3091..34a19c66f16 100644
--- a/codegen/smithy-aws-go-codegen/src/main/java/software/amazon/smithy/aws/go/codegen/AssembleMiddlewareStack.java
+++ b/codegen/smithy-aws-go-codegen/src/main/java/software/amazon/smithy/aws/go/codegen/AssembleMiddlewareStack.java
@@ -4,6 +4,7 @@
 import java.util.Map;
 import software.amazon.smithy.aws.traits.auth.SigV4Trait;
 import software.amazon.smithy.aws.traits.auth.UnsignedPayloadTrait;
+import software.amazon.smithy.go.codegen.SmithyGoDependency;
 import software.amazon.smithy.go.codegen.SymbolUtils;
 import software.amazon.smithy.go.codegen.integration.GoIntegration;
 import software.amazon.smithy.go.codegen.integration.MiddlewareRegistrar;
@@ -43,6 +44,16 @@ public List<RuntimeClientPlugin> getClientPlugins() {
                         )
                         .build(),
 
+                // Add ContentLengthMiddleware to operation stack
+                RuntimeClientPlugin.builder()
+                        .registerMiddleware(MiddlewareRegistrar.builder()
+                                .resolvedFunction(SymbolUtils.createValueSymbolBuilder(
+                                        "AddContentLengthMiddleware", SmithyGoDependency.SMITHY_HTTP_TRANSPORT)
+                                        .build())
+                                .build()
+                        )
+                        .build(),
+
                 // Add endpoint serialize middleware to operation stack
                 RuntimeClientPlugin.builder()
                         .registerMiddleware(MiddlewareRegistrar.builder()
diff --git a/codegen/smithy-aws-go-codegen/src/main/java/software/amazon/smithy/aws/go/codegen/RestJsonProtocolGenerator.java b/codegen/smithy-aws-go-codegen/src/main/java/software/amazon/smithy/aws/go/codegen/RestJsonProtocolGenerator.java
index 526a59d219b..7f94d5b35e6 100644
--- a/codegen/smithy-aws-go-codegen/src/main/java/software/amazon/smithy/aws/go/codegen/RestJsonProtocolGenerator.java
+++ b/codegen/smithy-aws-go-codegen/src/main/java/software/amazon/smithy/aws/go/codegen/RestJsonProtocolGenerator.java
@@ -218,7 +218,6 @@ public void generateProtocolTests(GenerationContext context) {
                             writer.addUseImports(SmithyGoDependency.SMITHY_MIDDLEWARE);
                             writer.openBlock("[]APIOptionFunc{", "},", () -> {
                                 writer.openBlock("func(s *middleware.Stack) error {", "},", () -> {
-                                    writer.write("s.Build.Clear()");
                                     writer.write("s.Finalize.Clear()");
                                     writer.write("return nil");
                                 });