From 54732ce126104f0f6b322a8f633c82112e0dd497 Mon Sep 17 00:00:00 2001
From: skotambkar <kotambka@amazon.com>
Date: Fri, 11 Dec 2020 02:53:41 -0800
Subject: [PATCH] remove content type header from streaming operations if
 payload is empty or nil

---
 aws/middleware/content_type.go                | 42 +++++++++++++++++++
 .../AwsHttpPresignURLClientGenerator.java     | 18 ++++++++
 .../integrationtest/s3/presign_test.go        |  6 +--
 service/s3/api_op_PutObject.go                |  3 ++
 4 files changed, 65 insertions(+), 4 deletions(-)
 create mode 100644 aws/middleware/content_type.go

diff --git a/aws/middleware/content_type.go b/aws/middleware/content_type.go
new file mode 100644
index 00000000000..9264d69add5
--- /dev/null
+++ b/aws/middleware/content_type.go
@@ -0,0 +1,42 @@
+package middleware
+
+import (
+	"context"
+	"fmt"
+	"github.com/awslabs/smithy-go/middleware"
+	smithyhttp "github.com/awslabs/smithy-go/transport/http"
+)
+
+// removeContentTypeHeader is a build middleware that removes
+// content type header if content-length header is unset or
+// is set to zero,
+type removeContentTypeHeader struct {
+}
+
+// ID the name of the middleware.
+func (m *removeContentTypeHeader) ID() string {
+	return "RemoveContentTypeHeader"
+}
+
+// HandleBuild adds or appends the constructed user agent to the request.
+func (m *removeContentTypeHeader) HandleBuild(ctx context.Context, in middleware.BuildInput, next middleware.BuildHandler) (
+	out middleware.BuildOutput, metadata middleware.Metadata, err error,
+) {
+	req, ok := in.Request.(*smithyhttp.Request)
+	if !ok {
+		return out, metadata, fmt.Errorf("unknown transport type %T", in)
+	}
+
+	// remove contentTypeHeader when content-length is zero
+	if req.ContentLength == 0 {
+		req.Header.Del("content-type")
+	}
+
+	return next.HandleBuild(ctx, in)
+}
+
+// RemoveContentTypeHeader removes content-type header if
+// content length is unset or equal to zero.
+func RemoveContentTypeHeader(stack *middleware.Stack) error {
+	return stack.Build.Add(&removeContentTypeHeader{}, middleware.After)
+}
diff --git a/codegen/smithy-aws-go-codegen/src/main/java/software/amazon/smithy/aws/go/codegen/AwsHttpPresignURLClientGenerator.java b/codegen/smithy-aws-go-codegen/src/main/java/software/amazon/smithy/aws/go/codegen/AwsHttpPresignURLClientGenerator.java
index f7d625845a5..1ada67074a1 100644
--- a/codegen/smithy-aws-go-codegen/src/main/java/software/amazon/smithy/aws/go/codegen/AwsHttpPresignURLClientGenerator.java
+++ b/codegen/smithy-aws-go-codegen/src/main/java/software/amazon/smithy/aws/go/codegen/AwsHttpPresignURLClientGenerator.java
@@ -93,6 +93,10 @@ private static final String addAsUnsignedPayloadName(String operationName) {
         return String.format("add%sPayloadAsUnsigned", operationName);
     }
 
+    private static final Symbol removeContentTypeHeader = SymbolUtils.createValueSymbolBuilder(
+            "RemoveContentTypeHeader", AwsGoDependency.AWS_MIDDLEWARE
+    ).build();
+
     // build pointable symbols
     private static Symbol buildSymbol(String name, boolean exported) {
         if (!exported) {
@@ -255,6 +259,20 @@ private void writePresignOperationFunction(
                                         .getAddOperationMiddlewareFuncName(operationSymbol));
                                 writer.write("c.$L,", CONVERT_TO_PRESIGN_MIDDLEWARE_NAME);
 
+                                // we should remove Content-Type header if input is a stream and
+                                // payload is empty/nil stream.
+                                if (operationInputShape.members().stream().anyMatch(memberShape -> {
+                                    return memberShape.getMemberTrait(model, StreamingTrait.class).isPresent();
+                                })) {
+                                    writer.addUseImports(SmithyGoDependency.SMITHY_MIDDLEWARE);
+                                    writer.addUseImports(AwsGoDependency.AWS_MIDDLEWARE);
+
+                                    writer.openBlock("func(stack *middleware.Stack, options Options) error {", "}," ,
+                                            () -> {
+                                        writer.write("return $T(stack)", removeContentTypeHeader);
+                                    });
+                                }
+
                                 // s3 needs to add a middleware to switch to using unsigned payload .
                                 if (isS3ServiceShape(model, serviceShape)) {
                                     writer.write("$L,", addAsUnsignedPayloadName(operationSymbol.getName()));
diff --git a/service/internal/integrationtest/s3/presign_test.go b/service/internal/integrationtest/s3/presign_test.go
index 027329cd383..cc7b8f310af 100644
--- a/service/internal/integrationtest/s3/presign_test.go
+++ b/service/internal/integrationtest/s3/presign_test.go
@@ -37,10 +37,8 @@ func TestInteg_PresignURL(t *testing.T) {
 			expectedSignedHeader: http.Header{},
 		},
 		"empty-body": {
-			body: bytes.NewReader([]byte("")),
-			expectedSignedHeader: http.Header{
-				"content-type": {"application/octet-stream"},
-			},
+			body:                 bytes.NewReader([]byte("")),
+			expectedSignedHeader: http.Header{},
 		},
 	}
 
diff --git a/service/s3/api_op_PutObject.go b/service/s3/api_op_PutObject.go
index 8e5f4c92723..80284a913b6 100644
--- a/service/s3/api_op_PutObject.go
+++ b/service/s3/api_op_PutObject.go
@@ -449,6 +449,9 @@ func (c *PresignClient) PresignPutObject(ctx context.Context, params *PutObjectI
 	result, _, err := c.client.invokeOperation(ctx, "PutObject", params, clientOptFns,
 		addOperationPutObjectMiddlewares,
 		c.convertToPresignMiddleware,
+		func(stack *middleware.Stack, options Options) error {
+			return awsmiddleware.RemoveContentTypeHeader(stack)
+		},
 		addPutObjectPayloadAsUnsigned,
 	)
 	if err != nil {