Add Gzip request compression feature

.changelog/80ed28327bcd4301a264f318efaf8216.json

@@ -0,0 +1,8 @@
+{
+    "id": "80ed2832-7bcd-4301-a264-f318efaf8216",
+    "type": "feature",
+    "description": "Support modeled request compression.",
+    "modules": [
+        "."
+    ]
+}

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/SmithyGoTypes.java

@@ -105,4 +105,10 @@ public static final class Bearer {
             public static final Symbol NewSignHTTPSMessage = SmithyGoDependency.SMITHY_AUTH_BEARER.valueSymbol("NewSignHTTPSMessage");
         }
     }
+
+    public static final class Private {
+        public static final class RequestCompression {
+            public static final Symbol AddRequestCompression = SmithyGoDependency.SMITHY_REQUEST_COMPRESSION.valueSymbol("AddRequestCompression");
+        }
+    }
 }

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

@@ -0,0 +1,160 @@
+/*
+ * 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 static software.amazon.smithy.go.codegen.GoWriter.goTemplate;
+
+import java.util.ArrayList;
+import java.util.List;
+import software.amazon.smithy.codegen.core.SymbolProvider;
+import software.amazon.smithy.go.codegen.GoCodegenPlugin;
+import software.amazon.smithy.go.codegen.GoDelegator;
+import software.amazon.smithy.go.codegen.GoSettings;
+import software.amazon.smithy.go.codegen.GoUniverseTypes;
+import software.amazon.smithy.go.codegen.GoWriter;
+import software.amazon.smithy.go.codegen.SmithyGoTypes;
+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.ShapeId;
+import software.amazon.smithy.model.traits.RequestCompressionTrait;
+import software.amazon.smithy.utils.ListUtils;
+import software.amazon.smithy.utils.MapUtils;
+
+
+public final class RequestCompression implements GoIntegration {
+    private static final String DISABLE_REQUEST_COMPRESSION = "DisableRequestCompression";
+
+    private static final String REQUEST_MIN_COMPRESSION_SIZE_BYTES = "RequestMinCompressSizeBytes";
+
+    private final List<RuntimeClientPlugin> runtimeClientPlugins = new ArrayList<>();
+
+    private static String getAddRequestCompressionMiddlewareFuncName(String operationName) {
+        return String.format("addOperation%sRequestCompressionMiddleware", operationName);
+    }
+
+    // Write operation plugin for request compression middleware
+    @Override
+    public void processFinalizedModel(GoSettings settings, Model model) {
+        ServiceShape service = settings.getService(model);
+        TopDownIndex.of(model)
+                .getContainedOperations(service).forEach(operation -> {
+                    if (!operation.hasTrait(RequestCompressionTrait.class)) {
+                        return;
+                    }
+                    SymbolProvider symbolProvider = GoCodegenPlugin.createSymbolProvider(model, settings);
+                    String funcName = getAddRequestCompressionMiddlewareFuncName(
+                            symbolProvider.toSymbol(operation).getName()
+                    );
+                    runtimeClientPlugins.add(RuntimeClientPlugin.builder().operationPredicate((m, s, o) -> {
+                        if (!o.hasTrait(RequestCompressionTrait.class)) {
+                            return false;
+                        }
+                        return o.equals(operation);
+                    }).registerMiddleware(MiddlewareRegistrar.builder()
+                        .resolvedFunction(SymbolUtils.buildPackageSymbol(funcName))
+                        .useClientOptions().build())
+                        .build());
+                });
+    }
+
+    @Override
+    public void writeAdditionalFiles(
+            GoSettings settings,
+            Model model,
+            SymbolProvider symbolProvider,
+            GoDelegator goDelegator
+    ) {
+        ServiceShape service = settings.getService(model);
+        for (ShapeId operationID : service.getAllOperations()) {
+            OperationShape operation = model.expectShape(operationID, OperationShape.class);
+            if (!operation.hasTrait(RequestCompressionTrait.class)) {
+                continue;
+            }
+            goDelegator.useShapeWriter(operation, writeMiddlewareHelper(symbolProvider, operation));
+        }
+    }
+
+
+    public static boolean isRequestCompressionService(Model model, ServiceShape service) {
+        return TopDownIndex.of(model)
+                .getContainedOperations(service).stream()
+                .anyMatch(it -> it.hasTrait(RequestCompressionTrait.class));
+    }
+
+    private GoWriter.Writable writeMiddlewareHelper(SymbolProvider symbolProvider, OperationShape operation) {
+        String operationName = symbolProvider.toSymbol(operation).getName();
+        RequestCompressionTrait trait = operation.expectTrait(RequestCompressionTrait.class);
+
+        return goTemplate("""
+                func $add:L(stack $stack:P, options Options) error {
+                    return $addInternal:T(stack, options.DisableRequestCompression, options.RequestMinCompressSizeBytes,
+                    $algorithms:W)
+                }
+                """,
+                MapUtils.of(
+                "add", getAddRequestCompressionMiddlewareFuncName(operationName),
+                "stack", SmithyGoTypes.Middleware.Stack,
+                "addInternal", SmithyGoTypes.Private.RequestCompression.AddRequestCompression,
+                "algorithms", generateAlgorithmList(trait.getEncodings())
+                ));
+    }
+
+    @Override
+    public List<RuntimeClientPlugin> getClientPlugins() {
+       runtimeClientPlugins.add(
+                RuntimeClientPlugin.builder()
+                        .servicePredicate(RequestCompression::isRequestCompressionService)
+                        .configFields(ListUtils.of(
+                                ConfigField.builder()
+                                        .name(DISABLE_REQUEST_COMPRESSION)
+                                        .type(GoUniverseTypes.Bool)
+                                        .documentation(
+                                        "Whether to disable automatic request compression for supported operations.")
+                                        .build(),
+                                ConfigField.builder()
+                                        .name(REQUEST_MIN_COMPRESSION_SIZE_BYTES)
+                                        .type(GoUniverseTypes.Int64)
+                                        .documentation("The minimum request body size, in bytes, at which compression "
+                                        + "should occur. The default value is 10 KiB. Values must fall within "
+                                        + "[0, 1MiB].")
+                                        .build()
+                        ))
+                        .build()
+       );
+
+       return runtimeClientPlugins;
+    }
+
+    private GoWriter.Writable generateAlgorithmList(List<String> algorithms) {
+        return goTemplate("""
+        []string{
+            $W
+        }
+        """,
+                GoWriter.ChainWritable.of(
+                        algorithms.stream()
+                                .map(it -> goTemplate("$S,", it))
+                                .toList()
+                ).compose(false));
+    }
+}

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

@@ -8,6 +8,7 @@ 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
 
 # modeled auth schemes
 software.amazon.smithy.go.codegen.integration.auth.SigV4AuthScheme

private/requestcompression/gzip.go

@@ -0,0 +1,30 @@
+package requestcompression
+
+import (
+	"bytes"
+	"compress/gzip"
+	"fmt"
+	"io"
+)
+
+func gzipCompress(input io.Reader) ([]byte, error) {
+	var b bytes.Buffer
+	w, err := gzip.NewWriterLevel(&b, gzip.DefaultCompression)
+	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
+}

private/requestcompression/request_compression.go

@@ -0,0 +1,103 @@
+// Package requestcompression implements runtime support for smithy-modeled
+// request compression.
+//
+// This package is designated as private and is intended for use only by the
+// smithy client runtime. The exported API therein is not considered stable and
+// is subject to breaking changes without notice.
+package requestcompression
+
+import (
+	"bytes"
+	"context"
+	"fmt"
+	"github.com/aws/smithy-go/middleware"
+	"github.com/aws/smithy-go/transport/http"
+	"io"
+)
+
+const maxRequestMinCompressSizeBytes = 10485760
+
+// Enumeration values for supported compress Algorithms.
+const (
+	GZIP = "gzip"
+)
+
+type compressFunc func(io.Reader) ([]byte, error)
+
+var allowedAlgorithms = map[string]compressFunc{
+	GZIP: gzipCompress,
+}
+
+// AddRequestCompression add requestCompression middleware to op stack
+func AddRequestCompression(stack *middleware.Stack, disabled bool, minBytes int64, algorithms []string) error {
+	return stack.Serialize.Add(&requestCompression{
+		disableRequestCompression:   disabled,
+		requestMinCompressSizeBytes: minBytes,
+		compressAlgorithms:          algorithms,
+	}, middleware.After)
+}
+
+type requestCompression struct {
+	disableRequestCompression   bool
+	requestMinCompressSizeBytes int64
+	compressAlgorithms          []string
+}
+
+// ID returns the ID of the middleware
+func (m requestCompression) ID() string {
+	return "RequestCompression"
+}
+
+// HandleSerialize gzip compress the request's stream/body if enabled by config fields
+func (m requestCompression) HandleSerialize(
+	ctx context.Context, in middleware.SerializeInput, next middleware.SerializeHandler,
+) (
+	out middleware.SerializeOutput, metadata middleware.Metadata, err error,
+) {
+	if m.disableRequestCompression {
+		return next.HandleSerialize(ctx, in)
+	}
+	// still need to check requestMinCompressSizeBytes in case it is out of range after service client config
+	if m.requestMinCompressSizeBytes < 0 || m.requestMinCompressSizeBytes > maxRequestMinCompressSizeBytes {
+		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)
+	}
+
+	for _, algorithm := range m.compressAlgorithms {
+		compressFunc := allowedAlgorithms[algorithm]
+		if compressFunc != nil {
+			if stream := req.GetStream(); stream != nil {
+				size, found, err := req.StreamLength()
+				if err != nil {
+					return out, metadata, fmt.Errorf("error while finding request stream length, %v", err)
+				} else if !found || size < m.requestMinCompressSizeBytes {
+					return next.HandleSerialize(ctx, in)
+				}
+
+				compressedBytes, err := compressFunc(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
+
+				if val := req.Header.Get("Content-Encoding"); val != "" {
+					req.Header.Set("Content-Encoding", fmt.Sprintf("%s, %s", val, algorithm))
+				} else {
+					req.Header.Set("Content-Encoding", algorithm)
+				}
+			}
+			break
+		}
+	}
+
+	return next.HandleSerialize(ctx, in)
+}