Add checksum support in request (#911)

* Add checksum support in request

* Address comments

Author: joewyz

aws/client/aws-client-restjson/src/it/java/software/amazon/smithy/java/client/aws/restjson/RestJson1ProtocolTests.java

@@ -35,8 +35,6 @@ public class RestJson1ProtocolTests {
     @HttpClientRequestTests
     @ProtocolTestFilter(
             skipTests = {
-                    // TODO: support checksums in requests
-                    "RestJsonHttpChecksumRequired",
                     // TODO: These tests require a payload even when the httpPayload member is null. Should it?
                     "RestJsonHttpWithHeadersButNoPayload",
                     "RestJsonHttpWithEmptyStructurePayload",

client/client-core/src/test/java/software/amazon/smithy/java/client/core/ClientTest.java

@@ -32,6 +32,7 @@
 import software.amazon.smithy.java.client.http.mock.MockPlugin;
 import software.amazon.smithy.java.client.http.mock.MockQueue;
 import software.amazon.smithy.java.client.http.plugins.ApplyHttpRetryInfoPlugin;
+import software.amazon.smithy.java.client.http.plugins.HttpChecksumPlugin;
 import software.amazon.smithy.java.client.http.plugins.UserAgentPlugin;
 import software.amazon.smithy.java.core.serde.document.Document;
 import software.amazon.smithy.java.dynamicclient.DynamicClient;
@@ -95,6 +96,7 @@ public void tracksPlugins() throws URISyntaxException {
                         // And HttpMessageExchange applies the UserAgentPlugin and ApplyHttpRetryInfoPlugin.
                         UserAgentPlugin.class,
                         ApplyHttpRetryInfoPlugin.class,
+                        HttpChecksumPlugin.class,
                         // User plugins are applied last.
                         FooPlugin.class));
     }

client/client-http/src/main/java/software/amazon/smithy/java/client/http/HttpMessageExchange.java

@@ -8,6 +8,7 @@
 import software.amazon.smithy.java.client.core.ClientConfig;
 import software.amazon.smithy.java.client.core.MessageExchange;
 import software.amazon.smithy.java.client.http.plugins.ApplyHttpRetryInfoPlugin;
+import software.amazon.smithy.java.client.http.plugins.HttpChecksumPlugin;
 import software.amazon.smithy.java.client.http.plugins.UserAgentPlugin;
 import software.amazon.smithy.java.http.api.HttpRequest;
 import software.amazon.smithy.java.http.api.HttpResponse;
@@ -19,6 +20,7 @@
  * <ul>
  *     <li>{@link UserAgentPlugin}</li>
  *     <li>{@link ApplyHttpRetryInfoPlugin}</li>
+ *     <li>{@link HttpChecksumPlugin}</li>
  * </ul>
  */
 public final class HttpMessageExchange implements MessageExchange<HttpRequest, HttpResponse> {
@@ -33,5 +35,6 @@ private HttpMessageExchange() {}
     public void configureClient(ClientConfig.Builder config) {
         config.applyPlugin(new UserAgentPlugin());
         config.applyPlugin(new ApplyHttpRetryInfoPlugin());
+        config.applyPlugin(new HttpChecksumPlugin());
     }
 }

client/client-http/src/main/java/software/amazon/smithy/java/client/http/plugins/HttpChecksumPlugin.java

@@ -0,0 +1,68 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package software.amazon.smithy.java.client.http.plugins;
+
+import java.security.MessageDigest;
+import java.security.NoSuchAlgorithmException;
+import java.util.Base64;
+import software.amazon.smithy.java.client.core.ClientConfig;
+import software.amazon.smithy.java.client.core.ClientPlugin;
+import software.amazon.smithy.java.client.core.interceptors.ClientInterceptor;
+import software.amazon.smithy.java.client.core.interceptors.RequestHook;
+import software.amazon.smithy.java.core.schema.TraitKey;
+import software.amazon.smithy.java.http.api.HttpRequest;
+import software.amazon.smithy.java.io.ByteBufferUtils;
+import software.amazon.smithy.model.traits.HttpChecksumRequiredTrait;
+import software.amazon.smithy.utils.ListUtils;
+import software.amazon.smithy.utils.SmithyInternalApi;
+
+/**
+ * Plugin that adds Content-MD5 header for operations with @httpChecksumRequired trait.
+ */
+@SmithyInternalApi
+public final class HttpChecksumPlugin implements ClientPlugin {
+
+    @Override
+    public void configureClient(ClientConfig.Builder config) {
+        config.addInterceptor(HttpChecksumInterceptor.INSTANCE);
+    }
+
+    static final class HttpChecksumInterceptor implements ClientInterceptor {
+        private static final ClientInterceptor INSTANCE = new HttpChecksumInterceptor();
+        private static final TraitKey<HttpChecksumRequiredTrait> CHECKSUM_REQUIRED_TRAIT_KEY =
+                TraitKey.get(HttpChecksumRequiredTrait.class);
+
+        @Override
+        public <RequestT> RequestT modifyBeforeTransmit(RequestHook<?, ?, RequestT> hook) {
+            return hook.mapRequest(HttpRequest.class, HttpChecksumInterceptor::processRequest);
+        }
+
+        private static HttpRequest processRequest(RequestHook<?, ?, HttpRequest> hook) {
+            if (hook.operation().schema().hasTrait(CHECKSUM_REQUIRED_TRAIT_KEY)) {
+                return addContentMd5Header(hook.request());
+            }
+            return hook.request();
+        }
+
+        static HttpRequest addContentMd5Header(HttpRequest request) {
+            var body = request.body();
+            if (body != null) {
+                var buffer = body.waitForByteBuffer();
+                var bytes = ByteBufferUtils.getBytes(buffer);
+                try {
+                    byte[] hash = MessageDigest.getInstance("MD5").digest(bytes);
+                    String base64Hash = Base64.getEncoder().encodeToString(hash);
+                    return request.toBuilder()
+                            .withReplacedHeader("Content-MD5", ListUtils.of(base64Hash))
+                            .build();
+                } catch (NoSuchAlgorithmException e) {
+                    throw new IllegalStateException("Unable to fetch message digest instance for MD5", e);
+                }
+            }
+            return request;
+        }
+    }
+}

client/client-http/src/test/java/software/amazon/smithy/java/client/http/plugins/HttpChecksumPluginTest.java

@@ -0,0 +1,53 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package software.amazon.smithy.java.client.http.plugins;
+
+import static org.hamcrest.MatcherAssert.assertThat;
+import static org.hamcrest.Matchers.equalTo;
+import static org.hamcrest.Matchers.hasSize;
+
+import java.net.URI;
+import java.nio.charset.StandardCharsets;
+import org.junit.jupiter.api.Test;
+import software.amazon.smithy.java.http.api.HttpRequest;
+import software.amazon.smithy.java.io.datastream.DataStream;
+
+public class HttpChecksumPluginTest {
+
+    @Test
+    public void interceptorAddsContentMd5HeaderForKnownBody() throws Exception {
+        var interceptor = new HttpChecksumPlugin.HttpChecksumInterceptor();
+        var req = HttpRequest.builder()
+                .uri(new URI("/"))
+                .method("POST")
+                .body(DataStream.ofBytes("test body".getBytes(StandardCharsets.UTF_8)))
+                .build();
+
+        var result = interceptor.addContentMd5Header(req);
+
+        var headers = result.headers().allValues("Content-MD5");
+        assertThat(headers, hasSize(1));
+        assertThat(headers.get(0), equalTo("u/mv50Mcr1+Jpgi8MejYIg=="));
+    }
+
+    @Test
+    public void interceptorReplacesExistingContentMd5Header() throws Exception {
+        var interceptor = new HttpChecksumPlugin.HttpChecksumInterceptor();
+        var req = HttpRequest.builder()
+                .uri(new URI("/"))
+                .method("POST")
+                .body(DataStream.ofBytes("test body".getBytes(StandardCharsets.UTF_8)))
+                .withAddedHeader("Content-MD5", "wrong-hash")
+                .build();
+
+        var result = interceptor.addContentMd5Header(req);
+
+        var headers = result.headers().allValues("Content-MD5");
+        assertThat(headers, hasSize(1));
+        assertThat(headers.get(0), equalTo("u/mv50Mcr1+Jpgi8MejYIg=="));
+    }
+
+}