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.

Author: sunfinite

deps/rabbitmq_aws/.editorconfig

@@ -12,4 +12,4 @@ indent_style = tab
 # 2 space indentation
 [{*.erl, *.hrl, *.md}]
 indent_style = space
-indent_size = 2
+indent_size = 4

deps/rabbitmq_aws/Makefile

@@ -7,9 +7,12 @@ define PROJECT_ENV
 []
 endef
 
+BUILD_DEPS = rabbit
+DEPS = rabbit_common
+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

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]}}]}.

deps/rabbitmq_aws/src/rabbitmq_aws.erl

@@ -15,6 +15,7 @@
     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,
@@ -158,6 +159,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 +231,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,

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,43 +453,53 @@ 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
 %%      for the specified profile in the config file, if it exists. If it does
 %%      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}

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]}
 ].

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),

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(),

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