Avoid throwing NPE in the CombinedResponseAsyncHttpResponseHandler#onStream

core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/http/async/CombinedResponseAsyncHttpResponseHandler.java

@@ -19,14 +19,14 @@
 
 import java.nio.ByteBuffer;
 import java.util.concurrent.CompletableFuture;
-import java.util.concurrent.atomic.AtomicReference;
 import org.reactivestreams.Publisher;
 import software.amazon.awssdk.annotations.SdkInternalApi;
 import software.amazon.awssdk.core.Response;
 import software.amazon.awssdk.core.exception.SdkException;
 import software.amazon.awssdk.core.internal.http.TransformingAsyncResponseHandler;
 import software.amazon.awssdk.http.SdkHttpFullResponse;
 import software.amazon.awssdk.http.SdkHttpResponse;
+import software.amazon.awssdk.utils.Validate;
 
 /**
  * Detects whether the response succeeded or failed by just checking the HTTP status and delegates to appropriate
@@ -39,7 +39,6 @@ public final class CombinedResponseAsyncHttpResponseHandler<OutputT>
     private final TransformingAsyncResponseHandler<OutputT> successResponseHandler;
     private final TransformingAsyncResponseHandler<? extends SdkException> errorResponseHandler;
     private CompletableFuture<SdkHttpResponse> headersFuture;
-    private final AtomicReference<SdkHttpFullResponse> response = new AtomicReference<>();
 
     public CombinedResponseAsyncHttpResponseHandler(
         TransformingAsyncResponseHandler<OutputT> successResponseHandler,
@@ -51,6 +50,7 @@ public CombinedResponseAsyncHttpResponseHandler(
 
     @Override
     public void onHeaders(SdkHttpResponse response) {
+        Validate.isTrue(headersFuture != null, "onHeaders() invoked without prepare().");
         headersFuture.complete(response);
         logRequestId(response);
 
@@ -59,7 +59,6 @@ public void onHeaders(SdkHttpResponse response) {
         } else {
             errorResponseHandler.onHeaders(response);
         }
-        this.response.set(toFullResponse(response));
     }
 
     @Override
@@ -73,7 +72,11 @@ public void onError(Throwable error) {
 
     @Override
     public void onStream(Publisher<ByteBuffer> publisher) {
-        if (this.response.get().isSuccessful()) {
+        Validate.isTrue(headersFuture != null, "onStream() invoked without prepare().");
+        Validate.isTrue(headersFuture.isDone(), "headersFuture is still not completed when onStream() is "
+                                                + "invoked.");
+        SdkHttpResponse sdkHttpResponse = headersFuture.join();
+        if (sdkHttpResponse.isSuccessful()) {
             successResponseHandler.onStream(publisher);
         } else {
             errorResponseHandler.onStream(publisher);
@@ -82,32 +85,31 @@ public void onStream(Publisher<ByteBuffer> publisher) {
 
     @Override
     public CompletableFuture<Response<OutputT>> prepare() {
-        this.response.set(null);
+        headersFuture = new CompletableFuture<>();
         CompletableFuture<OutputT> preparedTransformFuture = successResponseHandler.prepare();
 
         CompletableFuture<? extends SdkException> preparedErrorTransformFuture = errorResponseHandler == null ? null :
             errorResponseHandler.prepare();
 
-        headersFuture = new CompletableFuture<>();
-
         return headersFuture.thenCompose(headers -> {
+            SdkHttpFullResponse sdkHttpFullResponse = toFullResponse(headers);
             if (headers.isSuccessful()) {
                 return preparedTransformFuture.thenApply(
                     r -> Response.<OutputT>builder().response(r)
-                                                    .httpResponse(response.get())
+                                                    .httpResponse(sdkHttpFullResponse)
                                                     .isSuccess(true)
                                                     .build());
             }
 
             if (preparedErrorTransformFuture != null) {
                 return preparedErrorTransformFuture.thenApply(
                     e -> Response.<OutputT>builder().exception(e)
-                                                    .httpResponse(response.get())
+                                                    .httpResponse(sdkHttpFullResponse)
                                                     .isSuccess(false)
                                                     .build());
             }
             return CompletableFuture.completedFuture(
-                Response.<OutputT>builder().httpResponse(response.get())
+                Response.<OutputT>builder().httpResponse(sdkHttpFullResponse)
                                            .isSuccess(false)
                                            .build());
         });

core/sdk-core/src/test/java/software/amazon/awssdk/core/internal/http/async/CombinedResponseAsyncHttpResponseHandlerTest.java

@@ -0,0 +1,113 @@
+/*
+ * Copyright 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.internal.http.async;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+import io.reactivex.Flowable;
+import java.nio.ByteBuffer;
+import java.nio.charset.StandardCharsets;
+import java.util.concurrent.CompletableFuture;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import org.mockito.Mockito;
+import software.amazon.awssdk.core.Response;
+import software.amazon.awssdk.core.exception.SdkClientException;
+import software.amazon.awssdk.core.internal.http.TransformingAsyncResponseHandler;
+import software.amazon.awssdk.http.SdkHttpFullResponse;
+
+class CombinedResponseAsyncHttpResponseHandlerTest {
+
+    private CombinedResponseAsyncHttpResponseHandler<Void> responseHandler;
+    private TransformingAsyncResponseHandler<Void> successResponseHandler;
+    private TransformingAsyncResponseHandler<SdkClientException> errorResponseHandler;
+
+    @BeforeEach
+    public void setup() {
+        successResponseHandler = Mockito.mock(TransformingAsyncResponseHandler.class);
+        errorResponseHandler = Mockito.mock(TransformingAsyncResponseHandler.class);
+        responseHandler = new CombinedResponseAsyncHttpResponseHandler<>(successResponseHandler, errorResponseHandler);
+    }
+
+    @Test
+    void onStream_invokedWithoutPrepare_shouldThrowException() {
+        assertThatThrownBy(() ->
+                               responseHandler.onStream(publisher()))
+            .isInstanceOf(IllegalArgumentException.class).hasMessageContaining("onStream() invoked");
+
+    }
+
+    @Test
+    void onHeaders_invokedWithoutPrepare_shouldThrowException() {
+        assertThatThrownBy(() ->
+                               responseHandler.onHeaders(SdkHttpFullResponse.builder().build()))
+            .isInstanceOf(IllegalArgumentException.class).hasMessageContaining("onHeaders() invoked");
+
+    }
+
+    @Test
+    void onStream_invokedWithoutOnHeaders_shouldThrowException() {
+        when(successResponseHandler.prepare()).thenReturn(CompletableFuture.completedFuture(null));
+
+        responseHandler.prepare();
+        assertThatThrownBy(() ->
+                               responseHandler.onStream(publisher()))
+            .isInstanceOf(IllegalArgumentException.class).hasMessageContaining("headersFuture is still not completed when onStream()");
+
+    }
+
+    @Test
+    void successResponse_shouldCompleteHeaderFuture() {
+        when(successResponseHandler.prepare()).thenReturn(CompletableFuture.completedFuture(null));
+
+        CompletableFuture<Response<Void>> future = responseHandler.prepare();
+        SdkHttpFullResponse sdkHttpFullResponse = SdkHttpFullResponse.builder()
+                                                       .statusCode(200)
+                                                       .build();
+        Flowable<ByteBuffer> publisher = publisher();
+        responseHandler.onHeaders(sdkHttpFullResponse);
+        responseHandler.onStream(publisher);
+        verify(successResponseHandler).prepare();
+        verify(successResponseHandler).onStream(publisher);
+        assertThat(future).isDone();
+        assertThat(future.join().httpResponse()).isEqualTo(sdkHttpFullResponse);
+    }
+
+    @Test
+    void errorResponse_shouldCompleteHeaderFuture() {
+        when(errorResponseHandler.prepare()).thenReturn(CompletableFuture.completedFuture(null));
+
+        CompletableFuture<Response<Void>> future = responseHandler.prepare();
+        SdkHttpFullResponse sdkHttpFullResponse = SdkHttpFullResponse.builder()
+                                                                     .statusCode(400)
+                                                                     .build();
+        Flowable<ByteBuffer> publisher = publisher();
+        responseHandler.onHeaders(sdkHttpFullResponse);
+        responseHandler.onStream(publisher);
+        verify(errorResponseHandler).prepare();
+        verify(errorResponseHandler).onStream(publisher);
+        assertThat(future).isDone();
+        assertThat(future.join().httpResponse()).isEqualTo(sdkHttpFullResponse);
+    }
+
+    private static Flowable<ByteBuffer> publisher() {
+        return Flowable.just(ByteBuffer.wrap("string".getBytes(StandardCharsets.UTF_8)));
+    }
+
+}