From 0e08fc58eb61bba431ab4f217f7f7ad3a7dce9df Mon Sep 17 00:00:00 2001
From: Daniel Bankhead <danielbankhead@google.com>
Date: Mon, 29 Jul 2024 13:20:34 -0700
Subject: [PATCH] feat: Group Concurrent Access Token Requests for Base
 External Clients (#1840)

---
 src/auth/baseexternalclient.ts  | 18 ++++++++++
 test/test.baseexternalclient.ts | 64 +++++++++++++++++++++++++++++++++
 2 files changed, 82 insertions(+)

diff --git a/src/auth/baseexternalclient.ts b/src/auth/baseexternalclient.ts
index bd234ec3..26436979 100644
--- a/src/auth/baseexternalclient.ts
+++ b/src/auth/baseexternalclient.ts
@@ -252,6 +252,11 @@ export abstract class BaseExternalAccountClient extends AuthClient {
    */
   protected cloudResourceManagerURL: URL | string;
   protected supplierContext: ExternalAccountSupplierContext;
+  /**
+   * A pending access token request. Used for concurrent calls.
+   */
+  #pendingAccessToken: Promise<CredentialsWithResponse> | null = null;
+
   /**
    * Instantiate a BaseExternalAccountClient instance using the provided JSON
    * object loaded from an external account credentials file.
@@ -545,6 +550,19 @@ export abstract class BaseExternalAccountClient extends AuthClient {
    * @return A promise that resolves with the fresh GCP access tokens.
    */
   protected async refreshAccessTokenAsync(): Promise<CredentialsWithResponse> {
+    // Use an existing access token request, or cache a new one
+    this.#pendingAccessToken =
+      this.#pendingAccessToken || this.#internalRefreshAccessTokenAsync();
+
+    try {
+      return await this.#pendingAccessToken;
+    } finally {
+      // clear pending access token for future requests
+      this.#pendingAccessToken = null;
+    }
+  }
+
+  async #internalRefreshAccessTokenAsync(): Promise<CredentialsWithResponse> {
     // Retrieve the external credential.
     const subjectToken = await this.retrieveSubjectToken();
     // Construct the STS credentials options.
diff --git a/test/test.baseexternalclient.ts b/test/test.baseexternalclient.ts
index a0974a85..71c4ec7b 100644
--- a/test/test.baseexternalclient.ts
+++ b/test/test.baseexternalclient.ts
@@ -326,6 +326,70 @@ describe('BaseExternalAccountClient', () => {
 
       scope.done();
     });
+
+    it('should not duplicate access token requests for concurrent requests', async () => {
+      const client = new TestExternalAccountClient(externalAccountOptionsNoUrl);
+      const RESPONSE_A = {
+        access_token: 'ACCESS_TOKEN',
+        issued_token_type: 'urn:ietf:params:oauth:token-type:access_token',
+        token_type: 'Bearer',
+        expires_in: ONE_HOUR_IN_SECS,
+        scope: 'scope1 scope2',
+      };
+
+      const RESPONSE_B = {
+        ...RESPONSE_A,
+        access_token: 'ACCESS_TOKEN_2',
+      };
+
+      const scope = mockStsTokenExchange([
+        {
+          statusCode: 200,
+          response: RESPONSE_A,
+          request: {
+            grant_type: 'urn:ietf:params:oauth:grant-type:token-exchange',
+            audience,
+            scope: 'https://www.googleapis.com/auth/cloud-platform',
+            requested_token_type:
+              'urn:ietf:params:oauth:token-type:access_token',
+            subject_token: 'subject_token_0',
+            subject_token_type: 'urn:ietf:params:oauth:token-type:jwt',
+          },
+        },
+        {
+          statusCode: 200,
+          response: RESPONSE_B,
+          request: {
+            grant_type: 'urn:ietf:params:oauth:grant-type:token-exchange',
+            audience,
+            scope: 'https://www.googleapis.com/auth/cloud-platform',
+            requested_token_type:
+              'urn:ietf:params:oauth:token-type:access_token',
+            subject_token: 'subject_token_1',
+            subject_token_type: 'urn:ietf:params:oauth:token-type:jwt',
+          },
+        },
+      ]);
+
+      // simulate 5 concurrent requests
+      const calls = [
+        client.getAccessToken(),
+        client.getAccessToken(),
+        client.getAccessToken(),
+        client.getAccessToken(),
+        client.getAccessToken(),
+      ];
+
+      for (const {token} of await Promise.all(calls)) {
+        assert.strictEqual(token, RESPONSE_A.access_token);
+      }
+
+      // this should be handled in a second request as the above were all awaited and we're forcing an expiration
+      client.eagerRefreshThresholdMillis = RESPONSE_A.expires_in * 1000;
+      assert((await client.getAccessToken()).token, RESPONSE_B.access_token);
+
+      scope.done();
+    });
   });
 
   describe('projectNumber', () => {