Add support for AWS_SESSION_TOKEN

* Add comprehensive tests for the `rabbitmq_aws` cuttlefish schema. * Move `aws_prefer_imdsv2` setting to the `rabbitmq_aws` application. * Use AWS session token when present in env or config file. It was only used with IMDSv2 previously. * Add rabbitmq_aws:api_post_request/4, README cleanup (cherry picked from commit 251405c4e8)
2025-09-18 15:47:53 +00:00 · 2025-09-18 15:47:53 +00:00 · 717a8730b2
parent 056efd2e30
commit 717a8730b2
10 changed files with 150 additions and 86 deletions
--- a/deps/rabbitmq_aws/.editorconfig
+++ b/deps/rabbitmq_aws/.editorconfig
@ -1,15 +0,0 @@
-# top-most EditorConfig file
-root = true
-
-# Unix-style newlines with a newline ending every file
-[*]
-end_of_line = lf
-insert_final_newline = true
-
-[Makefile]
-indent_style = tab
-
-# 2 space indentation
-[{*.erl, *.hrl, *.md}]
-indent_style = space
-indent_size = 2
--- a/deps/rabbitmq_aws/Makefile
+++ b/deps/rabbitmq_aws/Makefile
@ -7,9 +7,11 @@ define PROJECT_ENV
 []
 endef

+BUILD_DEPS = rabbit
+TEST_DEPS = meck rabbitmq_ct_helpers rabbitmq_ct_client_helpers
 LOCAL_DEPS = crypto inets ssl xmerl public_key
-BUILD_DEPS = rabbit_common
-TEST_DEPS = meck rabbit rabbitmq_ct_helpers rabbitmq_ct_client_helpers
+
+PLT_APPS = rabbit

 DEP_EARLY_PLUGINS = rabbit_common/mk/rabbitmq-early-plugin.mk
 DEP_PLUGINS = rabbit_common/mk/rabbitmq-plugin.mk
--- a/deps/rabbitmq_aws/README.md
+++ b/deps/rabbitmq_aws/README.md
@ -11,7 +11,7 @@ A fork of [gmr/httpc-aws](https://github.com/gmr/httpc-aws) for use in building
 Configuration for *rabbitmq-aws* is can be provided in multiple ways. It is designed
 to behave similarly to the [AWS Command Line Interface](http://docs.aws.amazon.com/cli/latest/userguide/cli-chap-getting-started.html)
 with respect to providing region and configuration information. Additionally it
-has two methods, ``rabbitmq_aws:set_region/1`` and ``rabbitmq_aws:set_credentials/2``
+has two methods, `rabbitmq_aws:set_region/1` and `rabbitmq_aws:set_credentials/2`
 to allow for application specific configuration, bypassing the automatic configuration
 behavior.

@ -40,36 +40,36 @@ and [adds defenses against additional vulnerabilities](https://aws.amazon.com/bl
 AWS recommends adopting IMDSv2 and disabling IMDSv1 [by configuring the Instance Metadata Service on the EC2 instances](https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/configuring-instance-metadata-service.html).

 By default *rabbitmq-aws* will attempt to use IMDSv2 first and will fallback to use IMDSv1 if calls to IMDSv2 fail. This behavior can be overridden
-by setting the ``aws.prefer_imdsv2`` setting to ``false``.
+by setting the `aws.prefer_imdsv2` setting to `false`.

 ### Environment Variables

 As with the AWS CLI, the following environment variables can be used to provide
 configuration or to impact configuration behavior:

- - ``AWS_DEFAULT_PROFILE``
- - ``AWS_DEFAULT_REGION``
- - ``AWS_CONFIG_FILE``
- - ``AWS_SHARED_CREDENTIALS_FILE``
- - ``AWS_ACCESS_KEY_ID``
- - ``AWS_SECRET_ACCESS_KEY``
+ - `AWS_DEFAULT_PROFILE`
+ - `AWS_DEFAULT_REGION`
+ - `AWS_CONFIG_FILE`
+ - `AWS_SHARED_CREDENTIALS_FILE`
+ - `AWS_ACCESS_KEY_ID`
+ - `AWS_SECRET_ACCESS_KEY`

 ## API Functions

-  Method                                | Description
- ---------------------------------------|--------------------------------------------------------------------------------------------
- ``rabbitmq_aws:set_region/1``          | Manually specify the AWS region to make requests to.
- ``rabbitmq_aws:set_credentials/2``     | Manually specify the request credentials to use.
- ``rabbitmq_aws:refresh_credentials/0`` | Refresh the credentials from the environment, filesystem, or EC2 Instance Metadata Service.
- ``rabbitmq_aws:ensure_imdsv2_token_valid/0`` | Make sure EC2 IMDSv2 token is active and valid.
- ``rabbitmq_aws:api_get_request/2``           | Perform an AWS service API request.
- ``rabbitmq_aws:get/2``                 | Perform a GET request to the API specifying the service and request path.
- ``rabbitmq_aws:get/3``                 | Perform a GET request specifying the service, path, and headers.
- ``rabbitmq_aws:post/4``                | Perform a POST request specifying the service, path, headers, and body.
- ``rabbitmq_aws:request/5``             | Perform a request specifying the service, method, path, headers, and body.
- ``rabbitmq_aws:request/6``             | Perform a request specifying the service, method, path, headers, body, and ``httpc:http_options().``
- ``rabbitmq_aws:request/7``             | Perform a request specifying the service, method, path, headers, body,  ``httpc:http_options()``, and override the API endpoint. 
- 
+  Method                                    | Description
+ -------------------------------------------|--------------------------------------------------------------------------------------------
+ `rabbitmq_aws:set_region/1`                | Manually specify the AWS region to make requests to.
+ `rabbitmq_aws:set_credentials/2`           | Manually specify the request credentials to use.
+ `rabbitmq_aws:refresh_credentials/0`       | Refresh the credentials from the environment, filesystem, or EC2 Instance Metadata Service.
+ `rabbitmq_aws:ensure_imdsv2_token_valid/0` | Make sure EC2 IMDSv2 token is active and valid.
+ `rabbitmq_aws:get/2`                       | Perform a GET request to the API specifying the service and request path.
+ `rabbitmq_aws:get/3`                       | Perform a GET request specifying the service, path, and headers.
+ `rabbitmq_aws:post/4`                      | Perform a POST request specifying the service, path, headers, and body.
+ `rabbitmq_aws:request/5`                   | Perform a request specifying the service, method, path, headers, and body.
+ `rabbitmq_aws:request/6`                   | Perform a request specifying the service, method, path, headers, body, and `httpc:http_options().`
+ `rabbitmq_aws:request/7`                   | Perform a request specifying the service, method, path, headers, body,  `httpc:http_options()`, and override the API endpoint.
+ `rabbitmq_aws:api_get_request/2`           | Perform an AWS service API request with retries.
+ `rabbitmq_aws:api_post_request/2`          | Perform an AWS service API request with retries.

 ## Example Usage

@ -80,8 +80,7 @@ you're using the EC2 Instance Metadata Service for credentials:
 application:start(rabbitmq_aws).
 {ok, {Headers, Response}} = rabbitmq_aws:get("ec2","/?Action=DescribeTags&Version=2015-10-01").
 ```
-
-To configure credentials, invoke ``rabbitmq_aws:set_credentials/2``:
+To configure credentials, invoke `rabbitmq_aws:set_credentials/2`:

 ```erlang
 application:start(rabbitmq_aws).
--- a/deps/rabbitmq_aws/priv/schema/rabbitmq_aws.schema
+++ b/deps/rabbitmq_aws/priv/schema/rabbitmq_aws.schema
@ -15,5 +15,5 @@
 %% When false, EC2 IMDSv1 will be used first and no attempt will be made to use EC2 IMDSv2.
 %% See https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/configuring-instance-metadata-service.html.

-{mapping, "aws.prefer_imdsv2", "rabbit.aws_prefer_imdsv2",
+{mapping, "aws.prefer_imdsv2", "rabbitmq_aws.aws_prefer_imdsv2",
    [{datatype, {enum, [true, false]}}]}.
--- a/deps/rabbitmq_aws/src/rabbitmq_aws.erl
+++ b/deps/rabbitmq_aws/src/rabbitmq_aws.erl
@ -15,10 +15,12 @@
    refresh_credentials/0,
    request/5, request/6, request/7,
    set_credentials/2,
+    set_credentials/3,
    has_credentials/0,
    set_region/1,
    ensure_imdsv2_token_valid/0,
-    api_get_request/2
+    api_get_request/2,
+    api_post_request/4
 ]).

 %% gen-server exports
@ -158,6 +160,12 @@ set_credentials(NewState) ->
 set_credentials(AccessKey, SecretAccessKey) ->
    gen_server:call(rabbitmq_aws, {set_credentials, AccessKey, SecretAccessKey}).

+-spec set_credentials(access_key(), secret_access_key(), security_token()) -> ok.
+%% @doc Manually set the access credentials with session token for requests.
+%% @end
+set_credentials(AccessKey, SecretAccessKey, SessionToken) ->
+    gen_server:call(rabbitmq_aws, {set_credentials, AccessKey, SecretAccessKey, SessionToken}).
+
 -spec set_region(Region :: string()) -> ok.
 %% @doc Manually set the AWS region to perform API requests to.
 %% @end
@ -224,6 +232,14 @@ handle_msg({set_credentials, AccessKey, SecretAccessKey}, State) ->
        expiration = undefined,
        error = undefined
    }};
+handle_msg({set_credentials, AccessKey, SecretAccessKey, SessionToken}, State) ->
+    {reply, ok, State#state{
+        access_key = AccessKey,
+        secret_access_key = SecretAccessKey,
+        security_token = SessionToken,
+        expiration = undefined,
+        error = undefined
+    }};
 handle_msg({set_credentials, NewState}, State) ->
    {reply, ok, State#state{
        access_key = NewState#state.access_key,
@ -607,8 +623,10 @@ ensure_credentials_valid() ->
    case has_credentials(State) of
        true ->
            case expired_credentials(State#state.expiration) of
-                true -> refresh_credentials(State);
-                _ -> ok
+                true ->
+                    refresh_credentials(State);
+                _ ->
+                    ok
            end;
        _ ->
            refresh_credentials(State)
@ -618,19 +636,42 @@ ensure_credentials_valid() ->
 %% @doc Invoke an API call to an AWS service.
 %% @end
 api_get_request(Service, Path) ->
-    ?LOG_DEBUG("Invoking AWS request {Service: ~tp; Path: ~tp}...", [Service, Path]),
-    api_get_request_with_retries(Service, Path, ?MAX_RETRIES, ?LINEAR_BACK_OFF_MILLIS).
+    ?LOG_DEBUG("invoking AWS get request {Service: ~tp; Path: ~tp}...", [Service, Path]),
+    api_request_with_retries(Service, get, Path, "", [],
+                             ?MAX_RETRIES, ?LINEAR_BACK_OFF_MILLIS).

-spec api_get_request_with_retries(string(), path(), integer(), integer()) ->
+-spec api_post_request(
+    Service :: string(),
+    Path :: path(),
+    Body :: body(),
+    Headers :: headers()
+) -> result().
+%% @doc Perform a HTTP Post request to the AWS API for the specified service. The
+%%      response will automatically be decoded if it is either in JSON or XML
+%%      format.
+%% @end
+api_post_request(Service, Path, Body, Headers) ->
+    ?LOG_DEBUG("invoking AWS post request {Service: ~tp; Path: ~tp}...", [Service, Path]),
+    api_request_with_retries(Service, post, Path, Body, Headers,
+                             ?MAX_RETRIES, ?LINEAR_BACK_OFF_MILLIS).
+
+-spec api_request_with_retries(
+        Service :: string(),
+        Method :: method(),
+        Path :: path(),
+        Body :: body(),
+        Headers :: headers(),
+        Retries :: integer(),
+        WaitTime :: integer()) ->
    {'ok', list()} | {'error', term()}.
 %% @doc Invoke an API call to an AWS service with retries.
 %% @end
-api_get_request_with_retries(_, _, 0, _) ->
-    ?LOG_WARNING("Request to AWS service has failed after ~b retries", [?MAX_RETRIES]),
+api_request_with_retries(_, _, _, _, _, 0, _) ->
+    ?LOG_ERROR("Request to AWS service has failed after ~b retries", [?MAX_RETRIES]),
    {error, "AWS service is unavailable"};
-api_get_request_with_retries(Service, Path, Retries, WaitTimeBetweenRetries) ->
-    ensure_credentials_valid(),
-    case get(Service, Path) of
+api_request_with_retries(Service, Method, Path, Body, Headers, Retries, WaitTime) ->
+    ok = ensure_credentials_valid(),
+    case request(Service, Method, Path, Body, Headers) of
        {ok, {_Headers, Payload}} ->
            ?LOG_DEBUG("AWS request: ~ts~nResponse: ~tp", [Path, Payload]),
            {ok, Payload};
@ -645,6 +686,6 @@ api_get_request_with_retries(Service, Path, Retries, WaitTimeBetweenRetries) ->
                    ok
            end,
            ?LOG_WARNING("Will retry AWS request, remaining retries: ~b", [Retries]),
-            timer:sleep(WaitTimeBetweenRetries),
-            api_get_request_with_retries(Service, Path, Retries - 1, WaitTimeBetweenRetries)
+            timer:sleep(WaitTime),
+            api_request_with_retries(Service, Method, Path, Body, Headers, Retries - 1, WaitTime)
    end.
--- a/deps/rabbitmq_aws/src/rabbitmq_aws_config.erl
+++ b/deps/rabbitmq_aws/src/rabbitmq_aws_config.erl
@ -134,7 +134,8 @@ credentials(Profile) ->
    lookup_credentials(
        Profile,
        os:getenv("AWS_ACCESS_KEY_ID"),
-        os:getenv("AWS_SECRET_ACCESS_KEY")
+        os:getenv("AWS_SECRET_ACCESS_KEY"),
+        os:getenv("AWS_SESSION_TOKEN")
    ).

 -spec region() -> {ok, string()}.
@ -452,32 +453,39 @@ instance_id_url() ->
 -spec lookup_credentials(
    Profile :: string(),
    AccessKey :: string() | false,
-    SecretKey :: string() | false
+    SecretKey :: string() | false,
+    SessionToken :: string() | false
 ) ->
    security_credentials().
 %% @doc Return the access key and secret access key if they are set in
 %%      environment variables, otherwise lookup the credentials from the config
 %%      file for the specified profile.
 %% @end
-lookup_credentials(Profile, false, _) ->
+lookup_credentials(Profile, false, _, _) ->
    lookup_credentials_from_config(
        Profile,
        value(Profile, aws_access_key_id),
-        value(Profile, aws_secret_access_key)
+        value(Profile, aws_secret_access_key),
+        value(Profile, aws_session_token)
    );
-lookup_credentials(Profile, _, false) ->
+lookup_credentials(Profile, _, false, _) ->
    lookup_credentials_from_config(
        Profile,
        value(Profile, aws_access_key_id),
-        value(Profile, aws_secret_access_key)
+        value(Profile, aws_secret_access_key),
+        value(Profile, aws_session_token)
    );
-lookup_credentials(_, AccessKey, SecretKey) ->
-    {ok, AccessKey, SecretKey, undefined, undefined}.
+lookup_credentials(_, AccessKey, SecretKey, SessionToken) ->
+    case SessionToken of
+        false -> {ok, AccessKey, SecretKey, undefined, undefined};
+        SessionToken -> {ok, AccessKey, SecretKey, undefined, SessionToken}
+    end.

 -spec lookup_credentials_from_config(
    Profile :: string(),
    access_key() | {error, Reason :: atom()},
-    secret_access_key() | {error, Reason :: atom()}
+    secret_access_key() | {error, Reason :: atom()},
+    security_token() | {error, Reason :: atom()}
 ) ->
    security_credentials().
 %% @doc Return the access key and secret access key if they are set in
@ -485,10 +493,13 @@ lookup_credentials(_, AccessKey, SecretKey) ->
 %%      not exist or the profile is not set or the values are not set in the
 %%      profile, look up the values in the shared credentials file
 %% @end
-lookup_credentials_from_config(Profile, {error, _}, _) ->
+lookup_credentials_from_config(Profile, {error, _}, _, _) ->
    lookup_credentials_from_file(Profile, credentials_file_data());
-lookup_credentials_from_config(_, AccessKey, SecretKey) ->
-    {ok, AccessKey, SecretKey, undefined, undefined}.
+lookup_credentials_from_config(_, AccessKey, SecretKey, SessionToken) ->
+    case SessionToken of
+        {error, _} -> {ok, AccessKey, SecretKey, undefined, undefined};
+        SessionToken -> {ok, AccessKey, SecretKey, undefined, SessionToken}
+    end.

 -spec lookup_credentials_from_file(
    Profile :: string(),
@ -518,22 +529,24 @@ lookup_credentials_from_section(undefined) ->
 lookup_credentials_from_section(Credentials) ->
    AccessKey = proplists:get_value(aws_access_key_id, Credentials, undefined),
    SecretKey = proplists:get_value(aws_secret_access_key, Credentials, undefined),
-    lookup_credentials_from_proplist(AccessKey, SecretKey).
+    SessionToken = proplists:get_value(aws_session_token, Credentials, undefined),
+    lookup_credentials_from_proplist(AccessKey, SecretKey, SessionToken).

 -spec lookup_credentials_from_proplist(
    AccessKey :: access_key(),
-    SecretAccessKey :: secret_access_key()
+    SecretAccessKey :: secret_access_key(),
+    SessionToken :: security_token()
 ) ->
    security_credentials().
 %% @doc Process the contents of the Credentials proplists checking if the
 %%      access key and secret access key are both set.
 %% @end
-lookup_credentials_from_proplist(undefined, _) ->
+lookup_credentials_from_proplist(undefined, _, _) ->
    lookup_credentials_from_instance_metadata();
-lookup_credentials_from_proplist(_, undefined) ->
+lookup_credentials_from_proplist(_, undefined, _) ->
    lookup_credentials_from_instance_metadata();
-lookup_credentials_from_proplist(AccessKey, SecretKey) ->
-    {ok, AccessKey, SecretKey, undefined, undefined}.
+lookup_credentials_from_proplist(AccessKey, SecretKey, SessionToken) ->
+    {ok, AccessKey, SecretKey, undefined, SessionToken}.

 -spec lookup_credentials_from_instance_metadata() ->
    security_credentials().
@ -773,7 +786,7 @@ load_imdsv2_token() ->
 %% @doc Return headers used for instance metadata service requests.
 %% @end
 instance_metadata_request_headers() ->
-    case application:get_env(rabbit, aws_prefer_imdsv2) of
+    case application:get_env(rabbitmq_aws, aws_prefer_imdsv2) of
        {ok, false} ->
            [];
        %% undefined or {ok, true}
--- a/deps/rabbitmq_aws/test/config_schema_SUITE_data/rabbitmq_aws.snippets
+++ b/deps/rabbitmq_aws/test/config_schema_SUITE_data/rabbitmq_aws.snippets
@ -1,15 +1,17 @@
 [
    {rabbitmq_aws_prefer_imdsv2_false,
        "aws.prefer_imdsv2 = false",
-        [{rabbit, [
+        [{rabbitmq_aws, [
            {aws_prefer_imdsv2, false}
-        ]}],
+         ]}
+        ],
        [rabbitmq_aws]},

    {rabbitmq_aws_prefer_imdsv2_true,
        "aws.prefer_imdsv2 = true",
-        [{rabbit, [
+        [{rabbitmq_aws, [
            {aws_prefer_imdsv2, true}
-        ]}],
+         ]}
+        ],
        [rabbitmq_aws]}
 ].
--- a/deps/rabbitmq_aws/test/rabbitmq_aws_config_tests.erl
+++ b/deps/rabbitmq_aws/test/rabbitmq_aws_config_tests.erl
@ -135,6 +135,15 @@ credentials_test_() ->
                    rabbitmq_aws_config:credentials()
                )
            end},
+            {"from environment variables with session token", fun() ->
+                os:putenv("AWS_ACCESS_KEY_ID", "Sésame"),
+                os:putenv("AWS_SECRET_ACCESS_KEY", "ouvre-toi"),
+                os:putenv("AWS_SESSION_TOKEN", "session42"),
+                ?assertEqual(
+                    {ok, "Sésame", "ouvre-toi", undefined, "session42"},
+                    rabbitmq_aws_config:credentials()
+                )
+            end},
            {"from config file with default profile", fun() ->
                setup_test_config_env_var(),
                ?assertEqual(
@ -187,6 +196,13 @@ credentials_test_() ->
                    rabbitmq_aws_config:credentials("development")
                )
            end},
+            {"from credentials file with session token", fun() ->
+                setup_test_credentials_env_var(),
+                ?assertEqual(
+                    {ok, "foo3", "bar3", undefined, "session42"},
+                    rabbitmq_aws_config:credentials("with-session-token")
+                )
+            end},
            {"from credentials file with bad profile", fun() ->
                setup_test_credentials_env_var(),
                meck:expect(rabbitmq_aws, ensure_imdsv2_token_valid, 0, undefined),
--- a/deps/rabbitmq_aws/test/rabbitmq_aws_tests.erl
+++ b/deps/rabbitmq_aws/test/rabbitmq_aws_tests.erl
@ -16,6 +16,7 @@ init_test_() ->
        end,
        [
            {"ok", fun() ->
+                os:unsetenv("AWS_SESSION_TOKEN"),
                os:putenv("AWS_ACCESS_KEY_ID", "Sésame"),
                os:putenv("AWS_SECRET_ACCESS_KEY", "ouvre-toi"),
                {ok, Pid} = rabbitmq_aws:start_link(),
@ -604,7 +605,7 @@ api_get_request_test_() ->
                {ok, Pid} = rabbitmq_aws:start_link(),
                rabbitmq_aws:set_region("us-east-1"),
                rabbitmq_aws:set_credentials(State),
-                Result = rabbitmq_aws:api_get_request_with_retries("AWS", "API", 3, 1),
+                Result = rabbitmq_aws:api_request_with_retries("AWS", get, "API", "", [], 3, 1),
                ok = gen_server:stop(Pid),
                ?assertEqual({error, "AWS service is unavailable"}, Result),
                meck:validate(httpc)
@ -637,7 +638,7 @@ api_get_request_test_() ->
                {ok, Pid} = rabbitmq_aws:start_link(),
                rabbitmq_aws:set_region("us-east-1"),
                rabbitmq_aws:set_credentials(State),
-                Result = rabbitmq_aws:api_get_request_with_retries("AWS", "API", 3, 1),
+                Result = rabbitmq_aws:api_request_with_retries("AWS", get, "API", "", [], 3, 1),
                ok = gen_server:stop(Pid),
                ?assertEqual({ok, [{"data", "value"}]}, Result),
                meck:validate(httpc)
--- a/deps/rabbitmq_aws/test/test_aws_credentials.ini
+++ b/deps/rabbitmq_aws/test/test_aws_credentials.ini
@ -6,6 +6,11 @@ aws_secret_access_key=bar1
 aws_access_key_id=foo2
 aws_secret_access_key=bar2

+[with-session-token]
+aws_access_key_id=foo3
+aws_secret_access_key=bar3
+aws_session_token=session42
+
 [only-key]
 aws_access_key_id = foo3