Create EC2MetadataClientException for IMDS 4XX errors (#5947)

* Create EC2MetadataClientException for IMDS 400 errors

* Creating EC2MetadataClientException for IMDS 4XX errors

* Added Changelog

* Additional changes for EC2MetadataClientException handling

* Fixing checkstyle issues

* Additional changes for EC2MetadataClientException - added clear description and some code refactoring

* Additional EC2MetadataClientException improvements

Author: S-Saranya1

.changes/next-release/feature-AWSSDKforJavav2-8a464f3.json

@@ -0,0 +1,6 @@
+{
+    "type": "feature",
+    "category": "EC2 Metadata Client",
+    "contributor": "",
+    "description": "Added new Ec2MetadataClientException extending SdkClientException for IMDS unsuccessful responses that captures HTTP status codes, headers, and raw response content for improved error handling. See [#5786](https://github.com/aws/aws-sdk-java-v2/issues/5786)"
+}

core/imds/src/main/java/software/amazon/awssdk/imds/Ec2MetadataAsyncClient.java

@@ -21,6 +21,7 @@
 import software.amazon.awssdk.annotations.Immutable;
 import software.amazon.awssdk.annotations.SdkPublicApi;
 import software.amazon.awssdk.annotations.ThreadSafe;
+import software.amazon.awssdk.core.exception.RetryableException;
 import software.amazon.awssdk.core.exception.SdkClientException;
 import software.amazon.awssdk.http.async.SdkAsyncHttpClient;
 import software.amazon.awssdk.imds.internal.DefaultEc2MetadataAsyncClient;
@@ -68,6 +69,11 @@ public interface Ec2MetadataAsyncClient extends SdkAutoCloseable {
      *
      * @param path Input path
      * @return A CompletableFuture that completes when the MetadataResponse is made available.
+     * @throws Ec2MetadataClientException if the request returns a 4XX error response. The exception includes
+     *          the HTTP status code, headers, and error response body
+     * @throws RetryableException if the request returns a 5XX error response and should be retried
+     * @throws SdkClientException if the maximum number of retries is reached, if there's an IO error during
+     *          the request, or if the response is empty when success is expected
      */
     CompletableFuture<Ec2MetadataResponse> get(String path);
 

core/imds/src/main/java/software/amazon/awssdk/imds/Ec2MetadataClient.java

@@ -18,6 +18,7 @@
 import software.amazon.awssdk.annotations.Immutable;
 import software.amazon.awssdk.annotations.SdkPublicApi;
 import software.amazon.awssdk.annotations.ThreadSafe;
+import software.amazon.awssdk.core.exception.RetryableException;
 import software.amazon.awssdk.core.exception.SdkClientException;
 import software.amazon.awssdk.http.SdkHttpClient;
 import software.amazon.awssdk.imds.internal.DefaultEc2MetadataClient;
@@ -66,6 +67,11 @@ public interface Ec2MetadataClient extends SdkAutoCloseable {
      *
      * @param path  Input path
      * @return Instance metadata value as part of MetadataResponse Object
+     * @throws Ec2MetadataClientException if the request returns a 4XX error response. The exception includes
+     *         the HTTP status code, headers, and error response body
+     * @throws RetryableException if the request returns a 5XX error response and should be retried
+     * @throws SdkClientException if the maximum number of retries is reached, if there's an IO error during
+     *         the request, or if the response is empty when success is expected
      */
     Ec2MetadataResponse get(String path);
 

core/imds/src/main/java/software/amazon/awssdk/imds/Ec2MetadataClientException.java

@@ -0,0 +1,111 @@
+/*
+ * 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.imds;
+
+
+import software.amazon.awssdk.annotations.SdkPublicApi;
+import software.amazon.awssdk.core.SdkBytes;
+import software.amazon.awssdk.core.exception.SdkClientException;
+import software.amazon.awssdk.http.SdkHttpResponse;
+
+/**
+ * Extends {@link SdkClientException} for EC2 Instance Metadata Service (IMDS) non-successful
+ * responses (4XX codes). Provides detailed error information through:
+ * <p>
+ * - HTTP status code via {@link #statusCode()} for specific error handling
+ * - Raw response content via {@link #rawResponse()} containing the error response body
+ * - HTTP headers via {@link #sdkHttpResponse()} providing additional error context from the response
+ */
+
+@SdkPublicApi
+public final class Ec2MetadataClientException extends SdkClientException {
+
+    private final int statusCode;
+    private final SdkBytes rawResponse;
+    private final SdkHttpResponse sdkHttpResponse;
+
+    private Ec2MetadataClientException(BuilderImpl builder) {
+        super(builder);
+        this.statusCode = builder.statusCode;
+        this.rawResponse = builder.rawResponse;
+        this.sdkHttpResponse = builder.sdkHttpResponse;
+    }
+
+    /**
+     * @return The HTTP status code returned by the IMDS service.
+     */
+    public int statusCode() {
+        return statusCode;
+    }
+
+    /**
+     * Returns the response payload as bytes.
+     */
+    public SdkBytes rawResponse() {
+        return rawResponse;
+    }
+
+    /**
+     * Returns a map of HTTP headers associated with the error response.
+     */
+    public SdkHttpResponse sdkHttpResponse() {
+        return sdkHttpResponse;
+    }
+
+    public static Builder builder() {
+        return new BuilderImpl();
+    }
+
+    public interface Builder extends SdkClientException.Builder {
+        Builder statusCode(int statusCode);
+
+        Builder rawResponse(SdkBytes rawResponse);
+
+        Builder sdkHttpResponse(SdkHttpResponse sdkHttpResponse);
+
+        @Override
+        Ec2MetadataClientException build();
+    }
+
+    private static final class BuilderImpl extends SdkClientException.BuilderImpl implements Builder {
+        private int statusCode;
+        private SdkBytes rawResponse;
+        private SdkHttpResponse sdkHttpResponse;
+
+        @Override
+        public Builder statusCode(int statusCode) {
+            this.statusCode = statusCode;
+            return this;
+        }
+
+        @Override
+        public Builder rawResponse(SdkBytes rawResponse) {
+            this.rawResponse = rawResponse;
+            return this;
+        }
+
+        @Override
+        public Builder sdkHttpResponse(SdkHttpResponse sdkHttpResponse) {
+            this.sdkHttpResponse = sdkHttpResponse;
+            return this;
+        }
+
+        @Override
+        public Ec2MetadataClientException build() {
+            return new Ec2MetadataClientException(this);
+        }
+    }
+}

core/imds/src/main/java/software/amazon/awssdk/imds/internal/AsyncHttpRequestHelper.java

@@ -23,6 +23,7 @@
 import java.util.concurrent.CompletableFuture;
 import java.util.function.Function;
 import software.amazon.awssdk.annotations.SdkInternalApi;
+import software.amazon.awssdk.core.SdkBytes;
 import software.amazon.awssdk.core.exception.RetryableException;
 import software.amazon.awssdk.core.exception.SdkClientException;
 import software.amazon.awssdk.core.http.HttpResponseHandler;
@@ -37,6 +38,7 @@
 import software.amazon.awssdk.http.async.AsyncExecuteRequest;
 import software.amazon.awssdk.http.async.SdkAsyncHttpClient;
 import software.amazon.awssdk.http.async.SdkHttpContentPublisher;
+import software.amazon.awssdk.imds.Ec2MetadataClientException;
 import software.amazon.awssdk.utils.CompletableFutureUtils;
 
 @SdkInternalApi
@@ -84,15 +86,22 @@ private static String handleResponse(SdkHttpFullResponse response, ExecutionAttr
             response.content().orElseThrow(() -> SdkClientException.create("Unexpected error: empty response content"));
         String responseContent = uncheckedInputStreamToUtf8(inputStream);
 
-        // non-retryable error
-        if (statusCode.isOneOf(HttpStatusFamily.CLIENT_ERROR)) {
-            throw SdkClientException.builder().message(responseContent).build();
-        }
-
         // retryable error
         if (statusCode.isOneOf(HttpStatusFamily.SERVER_ERROR)) {
             throw RetryableException.create(responseContent);
         }
+
+        // non-retryable error
+        if (!statusCode.isOneOf(HttpStatusFamily.SUCCESSFUL)) {
+            throw Ec2MetadataClientException.builder()
+                                            .statusCode(response.statusCode())
+                                            .sdkHttpResponse(response)
+                                            .rawResponse(SdkBytes.fromUtf8String(responseContent))
+                                            .message(String.format("Failed to send request to IMDS. "
+                                                                   + "Service returned %d error",
+                                                                   response.statusCode()))
+                                            .build();
+        }
         return responseContent;
     }
 

core/imds/src/main/java/software/amazon/awssdk/imds/internal/DefaultEc2MetadataClient.java

@@ -30,6 +30,7 @@
 import software.amazon.awssdk.annotations.Immutable;
 import software.amazon.awssdk.annotations.SdkInternalApi;
 import software.amazon.awssdk.annotations.ThreadSafe;
+import software.amazon.awssdk.core.SdkBytes;
 import software.amazon.awssdk.core.exception.RetryableException;
 import software.amazon.awssdk.core.exception.SdkClientException;
 import software.amazon.awssdk.core.internal.http.loader.DefaultSdkHttpClientBuilder;
@@ -40,6 +41,7 @@
 import software.amazon.awssdk.http.HttpStatusFamily;
 import software.amazon.awssdk.http.SdkHttpClient;
 import software.amazon.awssdk.imds.Ec2MetadataClient;
+import software.amazon.awssdk.imds.Ec2MetadataClientException;
 import software.amazon.awssdk.imds.Ec2MetadataResponse;
 import software.amazon.awssdk.imds.Ec2MetadataRetryPolicy;
 import software.amazon.awssdk.imds.EndpointMode;
@@ -149,6 +151,19 @@ public Ec2MetadataResponse get(String path) {
         throw sdkClientExceptionBuilder.build();
     }
 
+    private void handleUnsuccessfulResponse(int statusCode, Optional<AbortableInputStream> responseBody,
+                                            HttpExecuteResponse response, Supplier<String> errorMessageSupplier) {
+        String responseContent = responseBody.map(BaseEc2MetadataClient::uncheckedInputStreamToUtf8)
+                                             .orElse("");
+
+        throw Ec2MetadataClientException.builder()
+                                        .statusCode(statusCode)
+                                        .sdkHttpResponse(response.httpResponse())
+                                        .rawResponse(SdkBytes.fromUtf8String(responseContent))
+                                        .message(errorMessageSupplier.get())
+                                        .build();
+    }
+
     private Ec2MetadataResponse sendRequest(String path, String token) throws IOException {
 
         HttpExecuteRequest httpExecuteRequest =
@@ -170,12 +185,9 @@ private Ec2MetadataResponse sendRequest(String path, String token) throws IOExce
         }
 
         if (!HttpStatusFamily.of(statusCode).isOneOf(HttpStatusFamily.SUCCESSFUL)) {
-            responseBody.map(BaseEc2MetadataClient::uncheckedInputStreamToUtf8)
-                        .ifPresent(str -> log.debug(() -> "Metadata request response body: " + str));
-            throw SdkClientException
-                .builder()
-                .message(String.format("The requested metadata at path '%s' returned Http code %s", path, statusCode))
-                .build();
+            handleUnsuccessfulResponse(statusCode, responseBody, response,
+                () -> String.format("The requested metadata at path '%s' returned Http code %s", path, statusCode)
+            );
         }
 
         AbortableInputStream abortableInputStream = responseBody.orElseThrow(
@@ -219,11 +231,9 @@ private Token getToken() {
         }
 
         if (!HttpStatusFamily.of(statusCode).isOneOf(HttpStatusFamily.SUCCESSFUL)) {
-            response.responseBody().map(BaseEc2MetadataClient::uncheckedInputStreamToUtf8)
-                    .ifPresent(body -> log.debug(() -> "Token request response body: " + body));
-            throw SdkClientException.builder()
-                                    .message("Could not retrieve token, " + statusCode + " error occurred.")
-                                    .build();
+            handleUnsuccessfulResponse(statusCode, response.responseBody(), response,
+                () -> String.format("Could not retrieve token, %d error occurred", statusCode)
+            );
         }
 
         String ttl = response.httpResponse()

core/imds/src/test/java/software/amazon/awssdk/imds/internal/BaseEc2MetadataClientTest.java

@@ -50,6 +50,7 @@
 import software.amazon.awssdk.core.exception.SdkClientException;
 import software.amazon.awssdk.core.retry.backoff.FixedDelayBackoffStrategy;
 import software.amazon.awssdk.imds.Ec2MetadataClientBuilder;
+import software.amazon.awssdk.imds.Ec2MetadataClientException;
 import software.amazon.awssdk.imds.Ec2MetadataResponse;
 import software.amazon.awssdk.imds.Ec2MetadataRetryPolicy;
 import software.amazon.awssdk.imds.EndpointMode;
@@ -356,4 +357,49 @@ void get_successOnFirstTry_shouldNotRetryAndSucceed_whenConnectionTakesMoreThanO
                 .withHeader(TOKEN_HEADER, equalTo("some-token")));
         });
     }
+
+    @Test
+    void clientError_throwsEc2MetadataClientExceptionWithFullResponse() {
+        stubFor(put(urlPathEqualTo(TOKEN_RESOURCE_PATH))
+                                    .willReturn(aResponse()
+                                    .withBody("token")
+                                    .withHeader(EC2_METADATA_TOKEN_TTL_HEADER, "21600")));
+
+        stubFor(get(urlPathEqualTo(AMI_ID_RESOURCE))
+                                    .willReturn(aResponse()
+                                    .withStatus(400)
+                                    .withHeader("x-custom-header", "custom-value")
+                                    .withBody("Bad Request")));
+
+        failureAssertions(
+            AMI_ID_RESOURCE,
+            Ec2MetadataClientException.class,
+            exception -> {
+                assertThat(exception.statusCode()).isEqualTo(400);
+                assertThat(exception.sdkHttpResponse().firstMatchingHeader("x-custom-header"))
+                    .hasValue("custom-value");
+                assertThat(exception.rawResponse().asUtf8String()).isEqualTo("Bad Request");
+            }
+        );
+    }
+
+    @Test
+    void tokenRequest_clientError_throwsEc2MetadataClientExceptionWithFullResponse() {
+        stubFor(put(urlPathEqualTo(TOKEN_RESOURCE_PATH))
+                    .willReturn(aResponse()
+                                    .withStatus(403)
+                                    .withHeader("x-error-type", "AccessDenied")
+                                    .withBody("Forbidden")));
+
+        failureAssertions(
+            AMI_ID_RESOURCE,
+            Ec2MetadataClientException.class,
+            exception -> {
+                assertThat(exception.statusCode()).isEqualTo(403);
+                assertThat(exception.sdkHttpResponse().firstMatchingHeader("x-error-type"))
+                    .hasValue("AccessDenied");
+                assertThat(exception.rawResponse().asUtf8String()).isEqualTo("Forbidden");
+            }
+        );
+    }
 }