Add an Option to Disable v2 to v1 Fallback

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;
 
@@ -509,7 +513,8 @@ static void s_client_on_token_response(struct imds_user_data *user_data) {
      * 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) {
+    if (!user_data->ec2_metadata_v1_disabled &&
+        (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 "
@@ -533,7 +538,10 @@ static void s_client_on_token_response(struct imds_user_data *user_data) {
         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);
+            user_data->client,
+            cursor.len == 0 ? NULL : &user_data->imds_token,
+            cursor.len != 0 || user_data->ec2_metadata_v1_disabled,
+            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;