Add Gzip request compression feature

codegen/smithy-go-codegen/src/main/java/software/amazon/smithy/go/codegen/SmithyGoDependency.java

@@ -53,6 +53,8 @@ public final class SmithyGoDependency {
     public static final GoDependency SMITHY_HTTP_TRANSPORT = smithy("transport/http", "smithyhttp");
     public static final GoDependency SMITHY_MIDDLEWARE = smithy("middleware");
     public static final GoDependency SMITHY_PRIVATE_PROTOCOL = smithy("private/protocol", "smithyprivateprotocol");
+    public static final GoDependency SMITHY_REQUEST_COMPRESSION =
+    smithy("private/requestcompression", "smithyrequestcompression");
     public static final GoDependency SMITHY_TIME = smithy("time", "smithytime");
     public static final GoDependency SMITHY_HTTP_BINDING = smithy("encoding/httpbinding");
     public static final GoDependency SMITHY_JSON = smithy("encoding/json", "smithyjson");

codegen/smithy-go-codegen/src/main/java/software/amazon/smithy/go/codegen/requestcompression/RequestCompression.java

@@ -0,0 +1,118 @@
+/*
+ * Copyright 2020 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License").
+ * You may not use this file except in compliance with the License.
+ * A copy of the License is located at
+ *
+ *  http://aws.amazon.com/apache2.0
+ *
+ * or in the "license" file accompanying this file. This file is distributed
+ * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+ * express or implied. See the License for the specific language governing
+ * permissions and limitations under the License.
+ */
+
+package software.amazon.smithy.go.codegen.requestcompression;
+
+import java.util.List;
+import software.amazon.smithy.codegen.core.SymbolProvider;
+import software.amazon.smithy.go.codegen.GoDelegator;
+import software.amazon.smithy.go.codegen.GoSettings;
+import software.amazon.smithy.go.codegen.GoWriter;
+import software.amazon.smithy.go.codegen.SmithyGoDependency;
+import software.amazon.smithy.go.codegen.SymbolUtils;
+import software.amazon.smithy.go.codegen.integration.ConfigField;
+import software.amazon.smithy.go.codegen.integration.GoIntegration;
+import software.amazon.smithy.go.codegen.integration.MiddlewareRegistrar;
+import software.amazon.smithy.go.codegen.integration.RuntimeClientPlugin;
+import software.amazon.smithy.model.Model;
+import software.amazon.smithy.model.knowledge.TopDownIndex;
+import software.amazon.smithy.model.shapes.OperationShape;
+import software.amazon.smithy.model.shapes.ServiceShape;
+import software.amazon.smithy.model.shapes.ToShapeId;
+import software.amazon.smithy.model.traits.RequestCompressionTrait;
+import software.amazon.smithy.utils.ListUtils;
+
+public final class RequestCompression implements GoIntegration {
+    private static final String ADD_REQUEST_COMPRESSION = "addRequestCompression";
+
+    private static final String ADD_REQUEST_COMPRESSION_INTERNAL = "AddRequestCompression";
+
+    private static final String DISABLE_REQUEST_COMPRESSION = "DisableRequestCompression";
+
+    private static final String REQUEST_MIN_COMPRESSION_SIZE_BYTES = "RequestMinCompressSizeBytes";
+
+    @Override
+    public void writeAdditionalFiles(
+            GoSettings settings,
+            Model model,
+            SymbolProvider symbolProvider,
+            GoDelegator goDelegator
+    ) {
+        ServiceShape service = settings.getService(model);
+        if (!isRequestCompressionService(model, service)) {
+            return;
+        }
+
+        goDelegator.useShapeWriter(service, this::writeMiddlewareHelper);
+    }
+
+
+    public static boolean isRequestCompressionService(Model model, ServiceShape service) {
+        TopDownIndex topDownIndex = TopDownIndex.of(model);
+        for (ToShapeId operation : topDownIndex.getContainedOperations(service)) {
+            OperationShape operationShape = model.expectShape(operation.toShapeId(), OperationShape.class);
+            if (operationShape.hasTrait(RequestCompressionTrait.class)) {
+                return true;
+            }
+        }
+        return false;
+    }
+
+    private void writeMiddlewareHelper(GoWriter writer) {
+        writer.openBlock("func $L(stack *middleware.Stack, options Options) error {", "}",
+                ADD_REQUEST_COMPRESSION, () -> {
+                    writer.write(
+                    "return $T(stack, options.DisableRequestCompression, options.RequestMinCompressSizeBytes)",
+                            SymbolUtils.createValueSymbolBuilder(ADD_REQUEST_COMPRESSION_INTERNAL,
+                                    SmithyGoDependency.SMITHY_REQUEST_COMPRESSION).build()
+                    );
+                });
+        writer.insertTrailingNewline();
+    }
+
+    @Override
+    public List<RuntimeClientPlugin> getClientPlugins() {
+        return ListUtils.of(
+                RuntimeClientPlugin.builder()
+                        .operationPredicate((model, service, operation) ->
+                        operation.hasTrait(RequestCompressionTrait.class))
+                        .registerMiddleware(MiddlewareRegistrar.builder()
+                                .resolvedFunction(SymbolUtils.createValueSymbolBuilder(ADD_REQUEST_COMPRESSION).build())
+                                .useClientOptions()
+                                .build())
+                        .build(),
+                RuntimeClientPlugin.builder()
+                        .servicePredicate(RequestCompression::isRequestCompressionService)
+                        .configFields(ListUtils.of(
+                                ConfigField.builder()
+                                        .name(DISABLE_REQUEST_COMPRESSION)
+                                        .type(SymbolUtils.createValueSymbolBuilder("bool")
+                                                .putProperty(SymbolUtils.GO_UNIVERSE_TYPE, true)
+                                                .build())
+                                        .documentation("Determine if request compression is allowed, default to false")
+                                        .build(),
+                                ConfigField.builder()
+                                        .name(REQUEST_MIN_COMPRESSION_SIZE_BYTES)
+                                        .type(SymbolUtils.createValueSymbolBuilder("int64")
+                                                .putProperty(SymbolUtils.GO_UNIVERSE_TYPE, true)
+                                                .build())
+                                        .documentation("Inclusive threshold request body size to trigger compression, "
+                                         + "default to 10240 and must be within 0 and 10485760 bytes inclusively")
+                                        .build()
+                        ))
+                        .build()
+        );
+    }
+}

codegen/smithy-go-codegen/src/main/resources/META-INF/services/software.amazon.smithy.go.codegen.integration.GoIntegration

@@ -9,3 +9,4 @@ software.amazon.smithy.go.codegen.integration.Paginators
 software.amazon.smithy.go.codegen.integration.Waiters
 software.amazon.smithy.go.codegen.integration.ClientLogger
 software.amazon.smithy.go.codegen.endpoints.EndpointClientPluginsGenerator
+software.amazon.smithy.go.codegen.requestcompression.RequestCompression

private/requestcompression/request_compression.go

@@ -0,0 +1,105 @@
+package requestcompression
+
+import (
+	"bytes"
+	"compress/gzip"
+	"context"
+	"fmt"
+	"github.com/aws/smithy-go/middleware"
+	"github.com/aws/smithy-go/transport/http"
+	"io"
+)
+
+// AddRequestCompression add requestCompression middleware to op stack
+func AddRequestCompression(stack *middleware.Stack, DisableRequestCompression bool, RequestMinCompressSizeBytes int64) error {
+	return stack.Build.Add(&requestCompression{
+		disableRequestCompression:   DisableRequestCompression,
+		requestMinCompressSizeBytes: RequestMinCompressSizeBytes,
+	}, middleware.After)
+}
+
+type requestCompression struct {
+	disableRequestCompression   bool
+	requestMinCompressSizeBytes int64
+}
+
+// ID returns the ID of the middleware
+func (m requestCompression) ID() string {
+	return "RequestCompression"
+}
+
+// HandleBuild gzip compress the request's stream/body if enabled by config fields
+func (m requestCompression) HandleBuild(
+	ctx context.Context, in middleware.BuildInput, next middleware.BuildHandler,
+) (
+	out middleware.BuildOutput, metadata middleware.Metadata, err error,
+) {
+	if m.disableRequestCompression {
+		return next.HandleBuild(ctx, in)
+	}
+	// still need to check requestMinCompressSizeBytes in case it is out of range after service client config
+	if m.requestMinCompressSizeBytes < 0 || m.requestMinCompressSizeBytes > 10485760 {
+		return out, metadata, fmt.Errorf("invalid range for min request compression size bytes %d, must be within 0 and 10485760 inclusively", m.requestMinCompressSizeBytes)
+	}
+
+	req, ok := in.Request.(*http.Request)
+	if !ok {
+		return out, metadata, fmt.Errorf("unknown request type %T", req)
+	}
+
+	var isCompressed bool
+	if stream := req.GetStream(); stream != nil {
+		compressedBytes, err := compress(stream)
+		if err != nil {
+			return out, metadata, fmt.Errorf("failed to compress request stream, %v", err)
+		}
+
+		var newReq *http.Request
+		if newReq, err = req.SetStream(bytes.NewReader(compressedBytes)); err != nil {
+			return out, metadata, fmt.Errorf("failed to set request stream, %v", err)
+		}
+		*req = *newReq
+		isCompressed = true
+	} else if req.ContentLength >= m.requestMinCompressSizeBytes {
+		compressedBytes, err := compress(req.Body)
+		if err != nil {
+			return out, metadata, fmt.Errorf("failed to compress request body, %v", err)
+		}
+
+		isCompressed = true
+		req.Body = io.NopCloser(bytes.NewReader(compressedBytes))
+	}
+
+	if isCompressed {
+		// Either append to the header if it already exists, else set it
+		if len(req.Header["Content-Encoding"]) != 0 {
+			req.Header["Content-Encoding"][0] += ", gzip"
+		} else {
+			req.Header.Set("Content-Encoding", "gzip")
+		}
+	}
+
+	return next.HandleBuild(ctx, in)
+}
+
+func compress(input io.Reader) ([]byte, error) {
+	var b bytes.Buffer
+	w, err := gzip.NewWriterLevel(&b, gzip.BestCompression)
+	if err != nil {
+		return nil, fmt.Errorf("failed to create gzip writer, %v", err)
+	}
+
+	inBytes, err := io.ReadAll(input)
+	if err != nil {
+		return nil, fmt.Errorf("failed read payload to compress, %v", err)
+	}
+
+	if _, err = w.Write(inBytes); err != nil {
+		return nil, fmt.Errorf("failed to write payload to be compressed, %v", err)
+	}
+	if err = w.Close(); err != nil {
+		return nil, fmt.Errorf("failed to flush payload being compressed, %v", err)
+	}
+
+	return b.Bytes(), nil
+}