Account ID support for Crendentials Providers

include/aws/auth/credentials.h

@@ -715,6 +715,24 @@ struct aws_credentials_provider_cognito_options {
     struct aws_auth_http_system_vtable *function_table;
 };
 
+/**
+ * Configuration options for `aws_credentials_new_with_options`
+ */
+struct aws_credentials_options {
+
+    /* Value for the aws access key id field */
+    struct aws_byte_cursor access_key_id_cursor;
+    /* Value for the secret access key field */
+    struct aws_byte_cursor secret_access_key_cursor;
+    /* (Optional) security token associated with the credentials */
+    struct aws_byte_cursor session_token_cursor;
+    /* (Optional) account id associated with the credentials */
+    struct aws_byte_cursor account_id_cursor;
+    /* Time point, in seconds since epoch, that the credentials will no longer be valid.
+     * For credentials that do not expire, use UINT64_MAX. */
+    uint64_t expiration_timepoint_seconds;
+};
+
 AWS_EXTERN_C_BEGIN
 
 /*
@@ -758,13 +776,9 @@ struct aws_credentials *aws_credentials_new(
  * @return a valid credentials object, or NULL
  */
 AWS_AUTH_API
-struct aws_credentials *aws_credentials_new_with_account_id(
+struct aws_credentials *aws_credentials_new_with_options(
     struct aws_allocator *allocator,
-    struct aws_byte_cursor access_key_id_cursor,
-    struct aws_byte_cursor secret_access_key_cursor,
-    struct aws_byte_cursor session_token_cursor,
-    struct aws_byte_cursor account_id_cursor,
-    uint64_t expiration_timepoint_seconds);
+    const struct aws_credentials_options *options);
 
 /**
  * Creates a new set of aws anonymous credentials.

include/aws/auth/private/credentials_utils.h

@@ -84,6 +84,7 @@ struct aws_parse_credentials_from_json_doc_options {
     const char *secret_access_key_name;
     const char *token_name;
     const char *expiration_name;
+    const char *account_id_name;
     const char *top_level_object_name;
     enum aws_parse_credentials_expiration_format expiration_format;
     bool token_required;
@@ -192,6 +193,13 @@ struct aws_profile_collection *aws_load_profile_collection_from_config_file(
 AWS_AUTH_API
 struct aws_string *aws_credentials_provider_resolve_region_from_env(struct aws_allocator *allocator);
 
+/*
+ * Parses account_id from arn with the format arn:partition:service:region:account-id:resource-id
+ * Returns an empty cursor if the account_id is not present or unable to parse the account_id
+ */
+AWS_AUTH_API
+struct aws_byte_cursor aws_parse_account_id_from_arn(struct aws_byte_cursor arn);
+
 AWS_EXTERN_C_END
 
 #endif /* AWS_AUTH_CREDENTIALS_PRIVATE_H */

source/aws_profile.c

@@ -24,6 +24,7 @@ static const struct aws_string *s_profile_get_property_value(
 AWS_STATIC_STRING_FROM_LITERAL(s_access_key_id_profile_var, "aws_access_key_id");
 AWS_STATIC_STRING_FROM_LITERAL(s_secret_access_key_profile_var, "aws_secret_access_key");
 AWS_STATIC_STRING_FROM_LITERAL(s_session_token_profile_var, "aws_session_token");
+AWS_STATIC_STRING_FROM_LITERAL(s_account_id_profile_var, "aws_account_id");
 
 struct aws_credentials *aws_credentials_new_from_profile(
     struct aws_allocator *allocator,
@@ -35,6 +36,14 @@ struct aws_credentials *aws_credentials_new_from_profile(
     }
 
     const struct aws_string *session_token = s_profile_get_property_value(profile, s_session_token_profile_var);
-
-    return aws_credentials_new_from_string(allocator, access_key, secret_key, session_token, UINT64_MAX);
+    const struct aws_string *account_id = s_profile_get_property_value(profile, s_account_id_profile_var);
+    struct aws_credentials_options creds_option = {
+        .access_key_id_cursor = aws_byte_cursor_from_string(access_key),
+        .secret_access_key_cursor = aws_byte_cursor_from_string(secret_key),
+        .session_token_cursor = aws_byte_cursor_from_string(session_token),
+        .account_id_cursor = aws_byte_cursor_from_string(account_id),
+        .expiration_timepoint_seconds = UINT64_MAX,
+    };
+
+    return aws_credentials_new_with_options(allocator, &creds_option);
 }

source/credentials.c

@@ -96,30 +96,25 @@ struct aws_credentials *aws_credentials_new(
 
     struct aws_byte_cursor account_id;
     AWS_ZERO_STRUCT(account_id);
-
-    return aws_credentials_new_with_account_id(
-        allocator,
-        access_key_id_cursor,
-        secret_access_key_cursor,
-        session_token_cursor,
-        account_id,
-        expiration_timepoint_seconds);
+    struct aws_credentials_options creds_option = {
+        .access_key_id_cursor = access_key_id_cursor,
+        .secret_access_key_cursor = secret_access_key_cursor,
+        .session_token_cursor = session_token_cursor,
+        .account_id_cursor = account_id,
+        .expiration_timepoint_seconds = expiration_timepoint_seconds,
+    };
+    return aws_credentials_new_with_options(allocator, &creds_option);
 }
-
-struct aws_credentials *aws_credentials_new_with_account_id(
+struct aws_credentials *aws_credentials_new_with_options(
     struct aws_allocator *allocator,
-    struct aws_byte_cursor access_key_id_cursor,
-    struct aws_byte_cursor secret_access_key_cursor,
-    struct aws_byte_cursor session_token_cursor,
-    struct aws_byte_cursor account_id_cursor,
-    uint64_t expiration_timepoint_seconds) {
+    const struct aws_credentials_options *options) {
 
-    if (access_key_id_cursor.ptr == NULL || access_key_id_cursor.len == 0) {
+    if (options->access_key_id_cursor.ptr == NULL || options->access_key_id_cursor.len == 0) {
         aws_raise_error(AWS_ERROR_INVALID_ARGUMENT);
         return NULL;
     }
 
-    if (secret_access_key_cursor.ptr == NULL || secret_access_key_cursor.len == 0) {
+    if (options->secret_access_key_cursor.ptr == NULL || options->secret_access_key_cursor.len == 0) {
         aws_raise_error(AWS_ERROR_INVALID_ARGUMENT);
         return NULL;
     }
@@ -136,34 +131,34 @@ struct aws_credentials *aws_credentials_new_with_account_id(
     credentials->identity_type = AWS_CREDENTIALS_IDENTITY;
     struct aws_credentials_identity *credentials_identity = &credentials->identity.credentials_identity;
     credentials_identity->access_key_id =
-        aws_string_new_from_array(allocator, access_key_id_cursor.ptr, access_key_id_cursor.len);
+        aws_string_new_from_array(allocator, options->access_key_id_cursor.ptr, options->access_key_id_cursor.len);
     if (credentials_identity->access_key_id == NULL) {
         goto error;
     }
 
-    credentials_identity->secret_access_key =
-        aws_string_new_from_array(allocator, secret_access_key_cursor.ptr, secret_access_key_cursor.len);
+    credentials_identity->secret_access_key = aws_string_new_from_array(
+        allocator, options->secret_access_key_cursor.ptr, options->secret_access_key_cursor.len);
     if (credentials_identity->secret_access_key == NULL) {
         goto error;
     }
 
-    if (session_token_cursor.ptr != NULL && session_token_cursor.len > 0) {
+    if (options->session_token_cursor.ptr != NULL && options->session_token_cursor.len > 0) {
         credentials_identity->session_token =
-            aws_string_new_from_array(allocator, session_token_cursor.ptr, session_token_cursor.len);
+            aws_string_new_from_array(allocator, options->session_token_cursor.ptr, options->session_token_cursor.len);
         if (credentials_identity->session_token == NULL) {
             goto error;
         }
     }
 
-    if (account_id_cursor.ptr != NULL && account_id_cursor.len > 0) {
+    if (options->account_id_cursor.ptr != NULL && options->account_id_cursor.len > 0) {
         credentials_identity->account_id =
-            aws_string_new_from_array(allocator, account_id_cursor.ptr, account_id_cursor.len);
+            aws_string_new_from_array(allocator, options->account_id_cursor.ptr, options->account_id_cursor.len);
         if (credentials_identity->account_id == NULL) {
             goto error;
         }
     }
 
-    credentials->expiration_timepoint_seconds = expiration_timepoint_seconds;
+    credentials->expiration_timepoint_seconds = options->expiration_timepoint_seconds;
 
     return credentials;
 
@@ -334,6 +329,7 @@ struct aws_credentials *aws_credentials_new_from_string(
     const struct aws_string *secret_access_key,
     const struct aws_string *session_token,
     uint64_t expiration_timepoint_seconds) {
+
     struct aws_byte_cursor access_key_cursor = aws_byte_cursor_from_string(access_key_id);
     struct aws_byte_cursor secret_access_key_cursor = aws_byte_cursor_from_string(secret_access_key);
     struct aws_byte_cursor session_token_cursor;

source/credentials_provider_ecs.c

@@ -188,6 +188,7 @@ static void s_ecs_finalize_get_credentials_query(struct aws_credentials_provider
         .access_key_id_name = "AccessKeyId",
         .secret_access_key_name = "SecretAccessKey",
         .token_name = "Token",
+        .account_id_name = "AccountId",
         .expiration_name = "Expiration",
         .token_required = true,
         .expiration_required = true,

source/credentials_provider_environment.c

@@ -12,6 +12,7 @@
 AWS_STATIC_STRING_FROM_LITERAL(s_access_key_id_env_var, "AWS_ACCESS_KEY_ID");
 AWS_STATIC_STRING_FROM_LITERAL(s_secret_access_key_env_var, "AWS_SECRET_ACCESS_KEY");
 AWS_STATIC_STRING_FROM_LITERAL(s_session_token_env_var, "AWS_SESSION_TOKEN");
+AWS_STATIC_STRING_FROM_LITERAL(s_account_id_env_var, "AWS_ACCOUNT_ID");
 
 static int s_credentials_provider_environment_get_credentials_async(
     struct aws_credentials_provider *provider,
@@ -23,16 +24,24 @@ static int s_credentials_provider_environment_get_credentials_async(
     struct aws_string *access_key_id = NULL;
     struct aws_string *secret_access_key = NULL;
     struct aws_string *session_token = NULL;
+    struct aws_string *account_id = NULL;
     struct aws_credentials *credentials = NULL;
     int error_code = AWS_ERROR_SUCCESS;
 
     aws_get_environment_value(allocator, s_access_key_id_env_var, &access_key_id);
     aws_get_environment_value(allocator, s_secret_access_key_env_var, &secret_access_key);
     aws_get_environment_value(allocator, s_session_token_env_var, &session_token);
+    aws_get_environment_value(allocator, s_account_id_env_var, &account_id);
 
     if (access_key_id != NULL && access_key_id->len > 0 && secret_access_key != NULL && secret_access_key->len > 0) {
-        credentials =
-            aws_credentials_new_from_string(allocator, access_key_id, secret_access_key, session_token, UINT64_MAX);
+        struct aws_credentials_options creds_option = {
+            .access_key_id_cursor = aws_byte_cursor_from_string(access_key_id),
+            .secret_access_key_cursor = aws_byte_cursor_from_string(secret_access_key),
+            .session_token_cursor = aws_byte_cursor_from_string(session_token),
+            .account_id_cursor = aws_byte_cursor_from_string(account_id),
+            .expiration_timepoint_seconds = UINT64_MAX,
+        };
+        credentials = aws_credentials_new_with_options(allocator, &creds_option);
         if (credentials == NULL) {
             error_code = aws_last_error();
         }
@@ -54,6 +63,7 @@ static int s_credentials_provider_environment_get_credentials_async(
     callback(credentials, error_code, user_data);
 
     aws_credentials_release(credentials);
+    aws_string_destroy(account_id);
     aws_string_destroy(session_token);
     aws_string_destroy(secret_access_key);
     aws_string_destroy(access_key_id);

source/credentials_provider_process.c

@@ -50,6 +50,7 @@ static int s_get_credentials_from_process(
         .access_key_id_name = "AccessKeyId",
         .secret_access_key_name = "SecretAccessKey",
         .token_name = "SessionToken",
+        .account_id_name = "AccountId",
         .expiration_name = "Expiration",
         .token_required = false,
         .expiration_required = false,

source/credentials_provider_sso.c

@@ -56,6 +56,7 @@ struct aws_sso_query_context {
     struct aws_retry_token *retry_token;
     struct aws_byte_buf path_and_query;
     struct aws_string *token;
+    struct aws_string *sso_account_id;
 
     int status_code;
     int error_code;
@@ -91,6 +92,7 @@ static void s_sso_query_context_destroy(struct aws_sso_query_context *sso_query_
     s_sso_query_context_reset_request_specific_data(sso_query_context);
     aws_byte_buf_clean_up(&sso_query_context->payload);
     aws_byte_buf_clean_up(&sso_query_context->path_and_query);
+    aws_string_destroy(sso_query_context->sso_account_id);
     aws_credentials_provider_release(sso_query_context->provider);
     aws_retry_token_release(sso_query_context->retry_token);
     aws_mem_release(sso_query_context->allocator, sso_query_context);
@@ -108,6 +110,7 @@ static struct aws_sso_query_context *s_sso_query_context_new(
     sso_query_context->provider = aws_credentials_provider_acquire(provider);
     sso_query_context->original_user_data = user_data;
     sso_query_context->original_callback = callback;
+    sso_query_context->sso_account_id = aws_string_new_from_string(provider->allocator, impl->sso_account_id);
 
     /* construct path and query */
     struct aws_byte_cursor account_id_cursor = aws_byte_cursor_from_string(impl->sso_account_id);
@@ -139,6 +142,7 @@ static struct aws_sso_query_context *s_sso_query_context_new(
  */
 static void s_finalize_get_credentials_query(struct aws_sso_query_context *sso_query_context) {
     struct aws_credentials *credentials = NULL;
+    struct aws_credentials *credentials_with_account_id = NULL;
 
     if (sso_query_context->error_code == AWS_ERROR_SUCCESS) {
         /* parse credentials */
@@ -162,6 +166,14 @@ static void s_finalize_get_credentials_query(struct aws_sso_query_context *sso_q
             AWS_LS_AUTH_CREDENTIALS_PROVIDER,
             "(id=%p) successfully queried credentials",
             (void *)sso_query_context->provider);
+        struct aws_credentials_options creds_option = {
+            .access_key_id_cursor = aws_credentials_get_access_key_id(credentials),
+            .secret_access_key_cursor = aws_credentials_get_secret_access_key(credentials),
+            .session_token_cursor = aws_credentials_get_session_token(credentials),
+            .account_id_cursor = aws_byte_cursor_from_string(sso_query_context->sso_account_id),
+            .expiration_timepoint_seconds = aws_credentials_get_expiration_timepoint_seconds(credentials),
+        };
+        credentials_with_account_id = aws_credentials_new_with_options(sso_query_context->allocator, &creds_option);
     } else {
         AWS_LOGF_ERROR(
             AWS_LS_AUTH_CREDENTIALS_PROVIDER,
@@ -175,11 +187,12 @@ static void s_finalize_get_credentials_query(struct aws_sso_query_context *sso_q
 
     /* pass the credentials back */
     sso_query_context->original_callback(
-        credentials, sso_query_context->error_code, sso_query_context->original_user_data);
+        credentials_with_account_id, sso_query_context->error_code, sso_query_context->original_user_data);
 
     /* clean up */
     s_sso_query_context_destroy(sso_query_context);
     aws_credentials_release(credentials);
+    aws_credentials_release(credentials_with_account_id);
 }
 static void s_on_retry_ready(struct aws_retry_token *token, int error_code, void *user_data);
 

source/credentials_provider_static.c

@@ -50,14 +50,14 @@ struct aws_credentials_provider *aws_credentials_provider_new_static(
     }
 
     AWS_ZERO_STRUCT(*provider);
-
-    struct aws_credentials *credentials = aws_credentials_new_with_account_id(
-        allocator,
-        options->access_key_id,
-        options->secret_access_key,
-        options->session_token,
-        options->account_id,
-        UINT64_MAX);
+    struct aws_credentials_options creds_option = {
+        .access_key_id_cursor = options->access_key_id,
+        .secret_access_key_cursor = options->secret_access_key,
+        .session_token_cursor = options->session_token,
+        .account_id_cursor = options->account_id,
+        .expiration_timepoint_seconds = UINT64_MAX,
+    };
+    struct aws_credentials *credentials = aws_credentials_new_with_options(allocator, &creds_option);
     if (credentials == NULL) {
         goto on_new_credentials_failure;
     }