Refactoring interceptors for async payloads and adding S3 checksum validation

Author: spfink

codegen/src/main/java/software/amazon/awssdk/codegen/poet/client/specs/Ec2ProtocolSpec.java

@@ -51,7 +51,7 @@ private MethodSpec dryRunMethod() {
                         StaxResponseHandler.class,
                         dryRunResult,
                         VoidStaxUnmarshaller.class)
-                .addStatement("\nclientHandler.execute(new $T<$T, $T>().withMarshaller($L).withResponseHandler($N)" +
+                .addStatement("\nclientHandler.execute(new $T<$T, $T>().marshaller($L).withResponseHandler($N)" +
                                 ".withInput($L))",
                         ClientExecutionParams.class,
                         Request.class,

codegen/src/main/java/software/amazon/awssdk/codegen/poet/client/specs/JsonProtocolSpec.java

@@ -174,7 +174,8 @@ public CodeBlock executionHandler(OperationModel opModel) {
                  opModel.getInput().getVariableName());
 
         if (opModel.hasStreamingInput()) {
-            codeBlock.add(".withMarshaller(new $T(new $T(protocolFactory), requestBody))",
+            codeBlock.add(".withRequestBody(requestBody)")
+                     .add(".withMarshaller(new $T(new $T(protocolFactory), requestBody))",
                           ParameterizedTypeName.get(ClassName.get(StreamingRequestMarshaller.class), requestType),
                           marshaller);
         } else {

codegen/src/main/java/software/amazon/awssdk/codegen/poet/client/specs/QueryProtocolSpec.java

@@ -106,7 +106,8 @@ public CodeBlock executionHandler(OperationModel opModel) {
                  "errorResponseHandler",
                  opModel.getInput().getVariableName());
         if (opModel.hasStreamingInput()) {
-            return codeBlock.add(".withMarshaller(new $T(new $T(protocolFactory), requestBody)));",
+            return codeBlock.add(".withRequestBody(requestBody)")
+                            .add(".withMarshaller(new $T(new $T(protocolFactory), requestBody)));",
                                  ParameterizedTypeName.get(ClassName.get(StreamingRequestMarshaller.class), requestType),
                                  marshaller)
                             .build();

codegen/src/test/resources/software/amazon/awssdk/codegen/poet/client/test-json-client-class.java

@@ -39,6 +39,7 @@
 import software.amazon.awssdk.core.interceptor.ExecutionAttributes;
 import software.amazon.awssdk.core.interceptor.ExecutionInterceptorChain;
 import software.amazon.awssdk.core.interceptor.InterceptorContext;
+import software.amazon.awssdk.core.interceptor.SdkExecutionAttribute;
 import software.amazon.awssdk.core.interceptor.SdkInternalExecutionAttribute;
 import software.amazon.awssdk.core.signer.Signer;
 import software.amazon.awssdk.http.SdkHttpFullRequest;
@@ -55,7 +56,8 @@ private AwsClientHandlerUtils() {
     }
 
     public static <InputT extends SdkRequest, OutputT extends SdkResponse> ExecutionContext createExecutionContext(
-        ClientExecutionParams<InputT, OutputT> executionParams, SdkClientConfiguration clientConfig) {
+        ClientExecutionParams<InputT, OutputT> executionParams,
+        SdkClientConfiguration clientConfig) {
 
         SdkRequest originalRequest = executionParams.getInput();
         AwsCredentialsProvider clientCredentials = clientConfig.option(AwsClientOption.CREDENTIALS_PROVIDER);
@@ -76,14 +78,17 @@ public static <InputT extends SdkRequest, OutputT extends SdkResponse> Execution
                           clientConfig.option(AwsClientOption.SERVICE_SIGNING_NAME))
             .putAttribute(AwsExecutionAttribute.AWS_REGION, clientConfig.option(AwsClientOption.AWS_REGION))
             .putAttribute(AwsSignerExecutionAttribute.SIGNING_REGION, clientConfig.option(AwsClientOption.SIGNING_REGION))
-            .putAttribute(SdkInternalExecutionAttribute.IS_FULL_DUPLEX, executionParams.isFullDuplex());
+            .putAttribute(SdkInternalExecutionAttribute.IS_FULL_DUPLEX, executionParams.isFullDuplex())
+            .putAttribute(SdkExecutionAttribute.CLIENT_TYPE, clientConfig.option(SdkClientOption.CLIENT_TYPE));
 
         ExecutionInterceptorChain executionInterceptorChain =
                 new ExecutionInterceptorChain(clientConfig.option(SdkClientOption.EXECUTION_INTERCEPTORS));
         return ExecutionContext.builder()
                                .interceptorChain(executionInterceptorChain)
                                .interceptorContext(InterceptorContext.builder()
                                                                      .request(originalRequest)
+                                                                     .asyncRequestBody(executionParams.getAsyncRequestBody())
+                                                                     .requestBody(executionParams.getRequestBody())
                                                                      .build())
                                .executionAttributes(executionAttributes)
                                .signer(computeSigner(originalRequest, clientConfig))

core/aws-core/src/main/java/software/amazon/awssdk/awscore/client/handler/AwsClientHandlerUtils.java

@@ -29,9 +29,10 @@
 import software.amazon.awssdk.core.exception.AbortedException;
 import software.amazon.awssdk.core.exception.ApiCallTimeoutException;
 import software.amazon.awssdk.core.internal.http.AmazonSyncHttpClient;
-import software.amazon.awssdk.http.InvokeableHttpRequest;
+import software.amazon.awssdk.http.HttpExecuteResponse;
+import software.amazon.awssdk.http.ExecutableHttpRequest;
 import software.amazon.awssdk.http.SdkHttpClient;
-import software.amazon.awssdk.http.SdkHttpFullResponse;
+import software.amazon.awssdk.http.SdkHttpResponse;
 import utils.HttpTestUtils;
 
 @RunWith(MockitoJUnitRunner.class)
@@ -43,16 +44,17 @@ public class AbortedExceptionClientExecutionTimerIntegrationTest  {
     private SdkHttpClient sdkHttpClient;
 
     @Mock
-    private InvokeableHttpRequest abortableCallable;
+    private ExecutableHttpRequest abortableCallable;
 
     @Before
     public void setup() throws Exception {
         when(sdkHttpClient.prepareRequest(any())).thenReturn(abortableCallable);
         httpClient = HttpTestUtils.testClientBuilder().httpClient(sdkHttpClient)
                                   .apiCallTimeout(Duration.ofMillis(1000))
                                   .build();
-        when(abortableCallable.call()).thenReturn(SdkHttpFullResponse.builder()
-                                                                     .statusCode(200)
+        when(abortableCallable.call()).thenReturn(HttpExecuteResponse.builder().response(SdkHttpResponse.builder()
+                                                                                                        .statusCode(200)
+                                                                                                        .build())
                                                                      .build());
     }
 

core/sdk-core/src/it/java/software/amazon/awssdk/core/http/timers/client/AbortedExceptionClientExecutionTimerIntegrationTest.java

@@ -0,0 +1,42 @@
+/*
+ * Copyright 2010-2018 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.awssdk.core;
+
+import software.amazon.awssdk.annotations.SdkPublicApi;
+
+/**
+ * Enum that represents the type of client being used.
+ */
+@SdkPublicApi
+public enum ClientType {
+
+    ASYNC("Async"),
+    SYNC("Sync");
+
+    private final String clientType;
+
+    ClientType(String clientType) {
+        this.clientType = clientType;
+    }
+
+    /* (non-Javadoc)
+     * @see java.lang.Enum#toString()
+     */
+    @Override
+    public String toString() {
+        return clientType;
+    }
+}

core/sdk-core/src/main/java/software/amazon/awssdk/core/ClientType.java

@@ -0,0 +1,84 @@
+/*
+ * Copyright 2010-2018 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.awssdk.core.checksums;
+
+import java.security.MessageDigest;
+import software.amazon.awssdk.annotations.SdkInternalApi;
+
+/**
+ * Implementation of {@link SdkChecksum} to calculate an MD5 checksum.
+ */
+@SdkInternalApi
+public class Md5Checksum implements SdkChecksum {
+
+    private MessageDigest digest;
+
+    private MessageDigest digestLastMarked;
+
+    public Md5Checksum() {
+        this.digest = getDigest();
+    }
+
+    @Override
+    public void update(int b) {
+        digest.update((byte) b);
+    }
+
+    @Override
+    public void update(byte[] b, int off, int len) {
+        digest.update(b, off, len);
+    }
+
+    @Override
+    public long getValue() {
+        throw new UnsupportedOperationException("Use getChecksumBytes() instead.");
+    }
+
+    @Override
+    public void reset() {
+        digest = (digestLastMarked == null)
+                   // This is necessary so that should there be a reset without a
+                   // preceding mark, the MD5 would still be computed correctly.
+                   ? getDigest()
+                   : cloneFrom(digestLastMarked);
+    }
+
+    private MessageDigest getDigest() {
+        try {
+            return MessageDigest.getInstance("MD5");
+        } catch (Exception e) {
+            throw new IllegalStateException("Unexpected error creating MD5 checksum", e);
+        }
+    }
+
+    @Override
+    public byte[] getChecksumBytes() {
+        return digest.digest();
+    }
+
+    @Override
+    public void mark(int readLimit) {
+        digestLastMarked = cloneFrom(digest);
+    }
+
+    private MessageDigest cloneFrom(MessageDigest from) {
+        try {
+            return (MessageDigest) from.clone();
+        } catch (CloneNotSupportedException e) { // should never occur
+            throw new IllegalStateException("unexpected", e);
+        }
+    }
+}

core/sdk-core/src/main/java/software/amazon/awssdk/core/checksums/Md5Checksum.java

@@ -0,0 +1,42 @@
+/*
+ * Copyright 2010-2018 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.awssdk.core.checksums;
+
+import java.util.zip.Checksum;
+import software.amazon.awssdk.annotations.SdkPublicApi;
+
+/**
+ * Extension of {@link Checksum} to support checksums and checksum validations used by the SDK that
+ * are not provided by the JDK.
+ */
+@SdkPublicApi
+public interface SdkChecksum extends Checksum {
+
+    /**
+     * Returns the computed checksum in a byte array rather than the long provided by
+     * {@link #getValue()}.
+     *
+     * @return byte[] containing the checksum
+     */
+    byte[] getChecksumBytes();
+
+    /**
+     * Allows marking a checksum for checksums that support the ability to mark and reset.
+     *
+     * @param readLimit the maximum limit of bytes that can be read before the mark position becomes invalid.
+     */
+    void mark(int readLimit);
+}

core/sdk-core/src/main/java/software/amazon/awssdk/core/checksums/SdkChecksum.java

@@ -15,6 +15,8 @@
 
 package software.amazon.awssdk.core.client.builder;
 
+import static software.amazon.awssdk.core.ClientType.ASYNC;
+import static software.amazon.awssdk.core.ClientType.SYNC;
 import static software.amazon.awssdk.core.client.config.SdkAdvancedAsyncClientOption.FUTURE_COMPLETION_EXECUTOR;
 import static software.amazon.awssdk.core.client.config.SdkAdvancedClientOption.SIGNER;
 import static software.amazon.awssdk.core.client.config.SdkAdvancedClientOption.USER_AGENT_PREFIX;
@@ -55,8 +57,8 @@
 import software.amazon.awssdk.core.internal.http.loader.DefaultSdkHttpClientBuilder;
 import software.amazon.awssdk.core.internal.util.UserAgentUtils;
 import software.amazon.awssdk.core.retry.RetryPolicy;
-import software.amazon.awssdk.http.ExecuteRequest;
-import software.amazon.awssdk.http.InvokeableHttpRequest;
+import software.amazon.awssdk.http.ExecutableHttpRequest;
+import software.amazon.awssdk.http.HttpExecuteRequest;
 import software.amazon.awssdk.http.SdkHttpClient;
 import software.amazon.awssdk.http.async.AsyncExecuteRequest;
 import software.amazon.awssdk.http.async.SdkAsyncHttpClient;
@@ -84,6 +86,7 @@
  */
 @SdkProtectedApi
 public abstract class SdkDefaultClientBuilder<B extends SdkClientBuilder<B, C>, C> implements SdkClientBuilder<B, C> {
+
     private static final SdkHttpClient.Builder DEFAULT_HTTP_CLIENT_BUILDER = new DefaultSdkHttpClientBuilder();
     private static final SdkAsyncHttpClient.Builder DEFAULT_ASYNC_HTTP_CLIENT_BUILDER = new DefaultSdkAsyncHttpClientBuilder();
 
@@ -204,6 +207,7 @@ protected SdkClientConfiguration finalizeChildConfiguration(SdkClientConfigurati
     private SdkClientConfiguration finalizeSyncConfiguration(SdkClientConfiguration config) {
         return config.toBuilder()
                      .option(SdkClientOption.SYNC_HTTP_CLIENT, resolveSyncHttpClient(config))
+                     .option(SdkClientOption.CLIENT_TYPE, SYNC)
                      .build();
     }
 
@@ -214,6 +218,7 @@ private SdkClientConfiguration finalizeAsyncConfiguration(SdkClientConfiguration
         return config.toBuilder()
                      .option(FUTURE_COMPLETION_EXECUTOR, resolveAsyncFutureCompletionExecutor(config))
                      .option(ASYNC_HTTP_CLIENT, resolveAsyncHttpClient(config))
+                     .option(SdkClientOption.CLIENT_TYPE, ASYNC)
                      .build();
     }
 
@@ -365,7 +370,7 @@ private NonManagedSdkHttpClient(SdkHttpClient delegate) {
         }
 
         @Override
-        public InvokeableHttpRequest prepareRequest(ExecuteRequest request) {
+        public ExecutableHttpRequest prepareRequest(HttpExecuteRequest request) {
             return delegate.prepareRequest(request);
         }
 

core/sdk-core/src/main/java/software/amazon/awssdk/core/client/builder/SdkDefaultClientBuilder.java

@@ -21,6 +21,7 @@
 import java.util.Map;
 import java.util.concurrent.ScheduledExecutorService;
 import software.amazon.awssdk.annotations.SdkProtectedApi;
+import software.amazon.awssdk.core.ClientType;
 import software.amazon.awssdk.core.ServiceConfiguration;
 import software.amazon.awssdk.core.client.builder.SdkClientBuilder;
 import software.amazon.awssdk.core.interceptor.ExecutionInterceptor;
@@ -86,6 +87,11 @@ public final class SdkClientOption<T> extends ClientOption<T> {
     public static final SdkClientOption<SdkHttpClient> SYNC_HTTP_CLIENT =
             new SdkClientOption<>(SdkHttpClient.class);
 
+    /**
+     * The type of client used to make requests.
+     */
+    public static final SdkClientOption<ClientType> CLIENT_TYPE = new SdkClientOption<ClientType>(ClientType.class);
+
     /**
      * @see ClientOverrideConfiguration#apiCallAttemptTimeout()
      */