From 2a17009931ef9ad46eb07805b96e634a47a9952f Mon Sep 17 00:00:00 2001
From: James Telfer <792299+jamestelfer@users.noreply.github.com>
Date: Thu, 19 Sep 2024 12:58:08 +1000
Subject: [PATCH] fix: set region after STS call

Ensures that the global STS endpoint is used.
---
 hooks/environment      | 14 +++++++-------
 tests/environment.bats | 26 ++++++++++++++++++++++++++
 2 files changed, 33 insertions(+), 7 deletions(-)

diff --git a/hooks/environment b/hooks/environment
index 3ae909f..19435d6 100755
--- a/hooks/environment
+++ b/hooks/environment
@@ -16,13 +16,6 @@ fi
 
 echo "~~~ :buildkite::key::aws: Requesting an OIDC token for AWS from Buildkite"
 
-region="${BUILDKITE_PLUGIN_AWS_ASSUME_ROLE_WITH_WEB_IDENTITY_REGION:-}"
-if [[ -n $region ]]; then
-  export AWS_REGION="$region"
-  export AWS_DEFAULT_REGION="$region"
-  echo "Using region: ${AWS_REGION}"
-fi
-
 BUILDKITE_OIDC_TOKEN="$(buildkite-agent oidc request-token --audience sts.amazonaws.com ${REQUEST_TOKEN_OPTIONAL_ARGS})"
 
 echo "~~~ :aws: Assuming role using OIDC token"
@@ -47,3 +40,10 @@ export AWS_SECRET_ACCESS_KEY="$(jq -r ".Credentials.SecretAccessKey" <<< "${RESP
 export AWS_SESSION_TOKEN="$(jq -r ".Credentials.SessionToken" <<< "${RESPONSE}")"
 
 echo "Assumed role: $(jq -r .AssumedRoleUser.AssumedRoleId <<< "${RESPONSE}")"
+
+region="${BUILDKITE_PLUGIN_AWS_ASSUME_ROLE_WITH_WEB_IDENTITY_REGION:-}"
+if [[ -n $region ]]; then
+  export AWS_REGION="$region"
+  export AWS_DEFAULT_REGION="$region"
+  echo "Using region: ${AWS_REGION}"
+fi
diff --git a/tests/environment.bats b/tests/environment.bats
index c46db87..58803f6 100644
--- a/tests/environment.bats
+++ b/tests/environment.bats
@@ -119,6 +119,32 @@ EOF
   unstub buildkite-agent
 }
 
+@test "region not used for STS call" {
+  export BUILDKITE_JOB_ID="job-uuid-42"
+  export BUILDKITE_PLUGIN_AWS_ASSUME_ROLE_WITH_WEB_IDENTITY_ROLE_ARN="role123"
+  export BUILDKITE_PLUGIN_AWS_ASSUME_ROLE_WITH_WEB_IDENTITY_REGION="eu-central-1"
+
+  stub buildkite-agent "oidc request-token --audience sts.amazonaws.com * : echo 'buildkite-oidc-token'"
+  stub aws "sts assume-role-with-web-identity --role-arn role123 --role-session-name buildkite-job-job-uuid-42 --web-identity-token buildkite-oidc-token : echo \"STS-REGION:[\${AWS_REGION-<not set>}]\" 1>&2; cat tests/sts.json"
+
+  run run_test_command $PWD/hooks/environment
+
+  assert_success
+  assert_output --partial "Using region: eu-central-1"
+  assert_output --partial "Role ARN: role123"
+  refute_output --partial "STS-REGION:[eu-central-1]"
+  assert_output --partial "STS-REGION:[<not set>]"
+
+  assert_output --partial "TESTRESULT:AWS_ACCESS_KEY_ID=access-key-id-value"
+  assert_output --partial "TESTRESULT:AWS_SECRET_ACCESS_KEY=secret-access-key-value"
+  assert_output --partial "TESTRESULT:AWS_SESSION_TOKEN=session-token-value"
+  assert_output --partial "TESTRESULT:AWS_REGION=eu-central-1"
+  assert_output --partial "TESTRESULT:AWS_DEFAULT_REGION=eu-central-1"
+
+  unstub aws
+  unstub buildkite-agent
+}
+
 @test "does not pass in a custom region" {
   export BUILDKITE_JOB_ID="job-uuid-42"
   export BUILDKITE_PLUGIN_AWS_ASSUME_ROLE_WITH_WEB_IDENTITY_ROLE_ARN="role123"