Add an Option to Disable v2 to v1 Fallback (#218)

Co-authored-by: Michael Graeb <graebm@amazon.com>

include/aws/auth/aws_imds_client.h

@@ -50,6 +50,11 @@ struct aws_imds_client_options {
      */
     enum aws_imds_protocol_version imds_version;
 
+    /*
+     * If true, fallback from v2 to v1 will be disabled for all cases
+     */
+    bool ec2_metadata_v1_disabled;
+
     /*
      * Table holding all cross-system functional dependencies for an imds client.
      *

include/aws/auth/credentials.h

@@ -216,6 +216,11 @@ struct aws_credentials_provider_imds_options {
      */
     enum aws_imds_protocol_version imds_version;
 
+    /*
+     * If true, fallback from v2 to v1 will be disabled for all cases
+     */
+    bool ec2_metadata_v1_disabled;
+
     /* For mocking the http layer in tests, leave NULL otherwise */
     struct aws_auth_http_system_vtable *function_table;
 };

source/aws_imds_client.c

@@ -67,6 +67,7 @@ struct aws_imds_client {
     struct aws_linked_list pending_queries;
     struct aws_mutex token_lock;
     struct aws_condition_variable token_signal;
+    bool ec2_metadata_v1_disabled;
 
     struct aws_atomic_var ref_count;
 };
@@ -146,6 +147,7 @@ struct aws_imds_client *aws_imds_client_new(
     client->function_table =
         options->function_table ? options->function_table : g_aws_credentials_provider_http_function_table;
     client->token_required = options->imds_version == IMDS_PROTOCOL_V1 ? false : true;
+    client->ec2_metadata_v1_disabled = options->ec2_metadata_v1_disabled;
     client->shutdown_options = options->shutdown_options;
 
     struct aws_socket_options socket_options;
@@ -220,6 +222,7 @@ struct imds_user_data {
     /* Indicate the request is a fallback from a failure call. */
     bool is_fallback_request;
     bool is_imds_token_request;
+    bool ec2_metadata_v1_disabled;
     int status_code;
     int error_code;
 
@@ -281,6 +284,7 @@ static struct imds_user_data *s_user_data_new(
     }
 
     wrapped_user_data->imds_token_required = client->token_required;
+    wrapped_user_data->ec2_metadata_v1_disabled = client->ec2_metadata_v1_disabled;
     aws_atomic_store_int(&wrapped_user_data->ref_count, 1);
     return wrapped_user_data;
 
@@ -504,26 +508,15 @@ static void s_client_on_token_response(struct imds_user_data *user_data) {
         s_update_token_safely(user_data->client, NULL, true, 0 /*expire_timestamp*/);
         return;
     }
-    /*
-     * Other than that, if meets any error, then token is not required,
-     * we should fall back to insecure request. Otherwise, we should use
-     * token in following requests.
-     */
-    if (user_data->status_code != AWS_HTTP_STATUS_CODE_200_OK || user_data->current_result.len == 0) {
-        AWS_LOGF_DEBUG(
-            AWS_LS_IMDS_CLIENT,
-            "(id=%p) IMDS client failed to fetch token for requester %p, fall back to v1 for the same "
-            "requester. Received response status code: %d",
-            (void *)user_data->client,
-            (void *)user_data,
-            user_data->status_code);
-        s_update_token_safely(user_data->client, NULL, false, 0 /*expire_timestamp*/);
-    } else {
+
+    if (user_data->status_code == AWS_HTTP_STATUS_CODE_200_OK && user_data->current_result.len != 0) {
+        AWS_LOGF_DEBUG(AWS_LS_IMDS_CLIENT, "(id=%p) IMDS client has fetched the token", (void *)user_data->client);
+
         struct aws_byte_cursor cursor = aws_byte_cursor_from_buf(&(user_data->current_result));
         aws_byte_cursor_trim_pred(&cursor, aws_char_is_space);
         aws_byte_buf_reset(&user_data->imds_token, true /*zero contents*/);
         if (aws_byte_buf_append_and_update(&user_data->imds_token, &cursor)) {
-            s_update_token_safely(user_data->client, NULL, true, 0 /*expire_timestamp*/);
+            s_update_token_safely(user_data->client, NULL /*token*/, true /*token_required*/, 0 /*expire_timestamp*/);
             return;
         }
         /* The token was ALWAYS last for 6 hours, 21600 secs. Use current timestamp plus 21595 secs as the expiration
@@ -532,8 +525,31 @@ static void s_client_on_token_response(struct imds_user_data *user_data) {
         user_data->client->function_table->aws_high_res_clock_get_ticks(&current);
         uint64_t expire_timestamp = aws_add_u64_saturating(
             current, aws_timestamp_convert(s_imds_token_ttl_secs, AWS_TIMESTAMP_SECS, AWS_TIMESTAMP_NANOS, NULL));
-        s_update_token_safely(
-            user_data->client, cursor.len == 0 ? NULL : &user_data->imds_token, cursor.len != 0, expire_timestamp);
+
+        AWS_ASSERT(cursor.len != 0);
+        s_update_token_safely(user_data->client, &user_data->imds_token, true /*token_required*/, expire_timestamp);
+    } else if (user_data->ec2_metadata_v1_disabled) {
+        AWS_LOGF_DEBUG(
+            AWS_LS_IMDS_CLIENT,
+            "(id=%p) IMDS client failed to fetch token for requester %p, and fall back to v1 is disabled."
+            "Received response status code: %d",
+            (void *)user_data->client,
+            (void *)user_data,
+            user_data->status_code);
+        s_update_token_safely(user_data->client, NULL /*token*/, true /*token_required*/, 0 /*expire_timestamp*/);
+    } else {
+        /* Request failed; falling back to insecure request.
+         * TODO: The retryable error (503 throttle) will also fall back to v1. Instead, we should just resend the token
+         * request.
+         */
+        AWS_LOGF_DEBUG(
+            AWS_LS_IMDS_CLIENT,
+            "(id=%p) IMDS client failed to fetch token for requester %p, fall back to v1 for the same "
+            "requester. Received response status code: %d",
+            (void *)user_data->client,
+            (void *)user_data,
+            user_data->status_code);
+        s_update_token_safely(user_data->client, NULL /*token*/, false /* token_required*/, 0 /*expire_timestamp*/);
     }
 }
 

source/credentials_provider_imds.c

@@ -84,6 +84,7 @@ struct aws_credentials_provider *aws_credentials_provider_new_imds(
         .bootstrap = options->bootstrap,
         .function_table = options->function_table,
         .imds_version = options->imds_version,
+        .ec2_metadata_v1_disabled = options->ec2_metadata_v1_disabled,
         .shutdown_options =
             {
                 .shutdown_callback = s_on_imds_client_shutdown,

tests/CMakeLists.txt

@@ -147,6 +147,7 @@ add_test_case(imds_client_new_release)
 add_test_case(imds_client_connect_failure)
 add_test_case(imds_client_token_request_failure)
 add_test_case(imds_client_insecure_fallback_request_failure)
+add_test_case(imds_client_v1_fallback_disabled_failure)
 add_test_case(imds_client_resource_request_failure)
 add_test_case(imds_client_resource_request_success)
 add_test_case(imds_client_insecure_resource_request_success)

tests/aws_imds_client_test.c

@@ -609,6 +609,67 @@ static int s_imds_client_insecure_fallback_request_failure(struct aws_allocator
 
 AWS_TEST_CASE(imds_client_insecure_fallback_request_failure, s_imds_client_insecure_fallback_request_failure);
 
+AWS_STATIC_STRING_FROM_LITERAL(
+    s_good_response,
+    "{\"AccessKeyId\":\"SuccessfulAccessKey\", \n  \"SecretAccessKey\":\"SuccessfulSecret\", \n  "
+    "\"Token\":\"TokenSuccess\", \n \"Expiration\":\"2020-02-25T06:03:31Z\"}");
+AWS_STATIC_STRING_FROM_LITERAL(s_access_key, "SuccessfulAccessKey");
+AWS_STATIC_STRING_FROM_LITERAL(s_secret_key, "SuccessfulSecret");
+AWS_STATIC_STRING_FROM_LITERAL(s_token, "TokenSuccess");
+AWS_STATIC_STRING_FROM_LITERAL(s_expiration, "2020-02-25T06:03:31Z");
+
+static int s_imds_client_v1_fallback_disabled_failure(struct aws_allocator *allocator, void *ctx) {
+    (void)ctx;
+
+    s_aws_imds_tester_init(allocator);
+    /* secure data flow is not supported */
+    s_tester.response_code[0] = AWS_HTTP_STATUS_CODE_403_FORBIDDEN;
+    /* v1 would have worked if fallback was not disabled */
+    struct aws_byte_cursor good_response_cursor = aws_byte_cursor_from_string(s_good_response);
+    aws_array_list_push_back(&s_tester.response_data_callbacks[1], &good_response_cursor);
+
+    struct aws_imds_client_options options = {
+        .bootstrap = s_tester.bootstrap,
+        .function_table = &s_mock_function_table,
+        .ec2_metadata_v1_disabled = true,
+        .shutdown_options =
+            {
+                .shutdown_callback = s_on_shutdown_complete,
+                .shutdown_user_data = NULL,
+            },
+    };
+
+    struct aws_imds_client *client = aws_imds_client_new(allocator, &options);
+
+    aws_imds_client_get_resource_async(
+        client, aws_byte_cursor_from_string(s_expected_imds_resource_uri), s_get_resource_callback, NULL);
+
+    s_aws_wait_for_resource_result();
+
+    ASSERT_TRUE(s_tester.has_received_resource_callback == true);
+    ASSERT_TRUE(s_tester.resource.len == 0);
+    /* There was no fallback so we didn't get any resource */
+    ASSERT_TRUE(s_validate_uri_path_and_resource(1, false /*no resource*/) == 0);
+    ASSERT_TRUE(s_tester.token_ttl_header_exist[0]);
+    ASSERT_TRUE(s_tester.token_ttl_header_expected[0]);
+    ASSERT_FALSE(s_tester.token_header_exist[0]);
+
+    ASSERT_UINT_EQUALS(s_tester.error_code, AWS_AUTH_IMDS_CLIENT_SOURCE_FAILURE);
+
+    aws_imds_client_release(client);
+
+    s_aws_wait_for_imds_client_shutdown_callback();
+
+    /* Because we mock the http connection manager, we never get a callback back from it */
+    aws_mem_release(allocator, client);
+
+    ASSERT_SUCCESS(s_aws_imds_tester_cleanup());
+
+    return 0;
+}
+
+AWS_TEST_CASE(imds_client_v1_fallback_disabled_failure, s_imds_client_v1_fallback_disabled_failure);
+
 static int s_imds_client_resource_request_failure(struct aws_allocator *allocator, void *ctx) {
     (void)ctx;
 
@@ -655,15 +716,6 @@ static int s_imds_client_resource_request_failure(struct aws_allocator *allocato
 
 AWS_TEST_CASE(imds_client_resource_request_failure, s_imds_client_resource_request_failure);
 
-AWS_STATIC_STRING_FROM_LITERAL(
-    s_good_response,
-    "{\"AccessKeyId\":\"SuccessfulAccessKey\", \n  \"SecretAccessKey\":\"SuccessfulSecret\", \n  "
-    "\"Token\":\"TokenSuccess\", \n \"Expiration\":\"2020-02-25T06:03:31Z\"}");
-AWS_STATIC_STRING_FROM_LITERAL(s_access_key, "SuccessfulAccessKey");
-AWS_STATIC_STRING_FROM_LITERAL(s_secret_key, "SuccessfulSecret");
-AWS_STATIC_STRING_FROM_LITERAL(s_token, "TokenSuccess");
-AWS_STATIC_STRING_FROM_LITERAL(s_expiration, "2020-02-25T06:03:31Z");
-
 static int s_imds_client_resource_request_success(struct aws_allocator *allocator, void *ctx) {
     (void)ctx;