From cb81b763d9b4ff8db67521f6b1944a387f9168c5 Mon Sep 17 00:00:00 2001
From: aeitzman <eitzman@google.com>
Date: Tue, 5 Mar 2024 11:57:19 -0800
Subject: [PATCH] fix: making aws request signer get a new session token each
 time security credentials are requested.

---
 src/auth/awsclient.ts  | 12 +++++-------
 test/test.awsclient.ts |  1 +
 2 files changed, 6 insertions(+), 7 deletions(-)

diff --git a/src/auth/awsclient.ts b/src/auth/awsclient.ts
index bd47c7c4..dc2e0eee 100644
--- a/src/auth/awsclient.ts
+++ b/src/auth/awsclient.ts
@@ -155,7 +155,7 @@ export class AwsClient extends BaseExternalAccountClient {
       // The credential config contains all the URLs by default but clients may be running this
       // where the metadata server is not available and returning the credentials through the environment.
       // Removing this check may break them.
-      if (this.shouldUseMetadataServer() && this.imdsV2SessionTokenUrl) {
+      if (!this.regionFromEnv && this.imdsV2SessionTokenUrl) {
         metadataHeaders['x-aws-ec2-metadata-token'] =
           await this.getImdsV2SessionToken();
       }
@@ -167,6 +167,10 @@ export class AwsClient extends BaseExternalAccountClient {
         if (this.securityCredentialsFromEnv) {
           return this.securityCredentialsFromEnv;
         }
+        if (this.imdsV2SessionTokenUrl) {
+          metadataHeaders['x-aws-ec2-metadata-token'] =
+            await this.getImdsV2SessionToken();
+        }
         // Since the role on a VM can change, we don't need to cache it.
         const roleName = await this.getAwsRoleName(metadataHeaders);
         // Temporary credentials typically last for several hours.
@@ -316,12 +320,6 @@ export class AwsClient extends BaseExternalAccountClient {
     return response.data;
   }
 
-  private shouldUseMetadataServer(): boolean {
-    // The metadata server must be used when either the AWS region or AWS security
-    // credentials cannot be retrieved through their defined environment variables.
-    return !this.regionFromEnv || !this.securityCredentialsFromEnv;
-  }
-
   private get regionFromEnv(): string | null {
     // The AWS region can be provided through AWS_REGION or AWS_DEFAULT_REGION.
     // Only one is required.
diff --git a/test/test.awsclient.ts b/test/test.awsclient.ts
index 9a99685c..800f670b 100644
--- a/test/test.awsclient.ts
+++ b/test/test.awsclient.ts
@@ -330,6 +330,7 @@ describe('AwsClient', () => {
             reqheaders: {'x-aws-ec2-metadata-token-ttl-seconds': '300'},
           })
             .put('/latest/api/token')
+            .twice()
             .reply(200, awsSessionToken)
         );