[core-http] Token refresher fix: Refresh in the background, keep returning valid token

sdk/core/core-http/CHANGELOG.md

@@ -2,6 +2,7 @@
 
 ## 1.2.4 (Unreleased)
 
+- Updated the `bearerTokenAuthenticationPolicy` to refresh tokens only when they're about to expire and not multiple times before. This fixes the issue: [13369](https://github.com/Azure/azure-sdk-for-js/issues/13369).
 
 ## 1.2.3 (2021-02-04)
 

sdk/core/core-http/src/credentials/accessTokenRefresher.ts

@@ -21,9 +21,13 @@ export class AccessTokenRefresher {
    * that we are ready for a new refresh.
    */
   public isReady(): boolean {
-    // We're only ready for a new refresh if the required milliseconds have passed.
     return (
-      !this.lastCalled || Date.now() - this.lastCalled > this.requiredMillisecondsBeforeNewRefresh
+      // If no refresh has happened, we allow new refreshes to happen.
+      !this.lastCalled ||
+      // If a request is currently active, we allow new refreshes to happen, since the promise gets reused.
+      !!this.promise ||
+      // Otherwise we're ready for a new refresh if the required milliseconds have passed.
+      Date.now() - this.lastCalled >= this.requiredMillisecondsBeforeNewRefresh
     );
   }
 

sdk/core/core-http/src/policies/bearerTokenAuthenticationPolicy.ts

@@ -1,7 +1,7 @@
 // Copyright (c) Microsoft Corporation.
 // Licensed under the MIT license.
 
-import { TokenCredential, GetTokenOptions } from "@azure/core-auth";
+import { TokenCredential, GetTokenOptions, AccessToken } from "@azure/core-auth";
 import {
   BaseRequestPolicy,
   RequestPolicy,
@@ -54,6 +54,8 @@ export function bearerTokenAuthenticationPolicy(
  *
  */
 export class BearerTokenAuthenticationPolicy extends BaseRequestPolicy {
+  private token: AccessToken | undefined;
+
   /**
    * Creates a new BearerTokenAuthenticationPolicy object.
    *
@@ -83,34 +85,36 @@ export class BearerTokenAuthenticationPolicy extends BaseRequestPolicy {
         spanOptions: webResource.spanOptions
       }
     });
-    webResource.headers.set(Constants.HeaderConstants.AUTHORIZATION, `Bearer ${token}`);
+    if (token) {
+      webResource.headers.set(Constants.HeaderConstants.AUTHORIZATION, `Bearer ${token}`);
+    }
     return this._nextPolicy.sendRequest(webResource);
   }
 
-  /**
-   * Attempts a token update if any other time related conditionals have been reached based on the tokenRefresher class.
-   */
-  private async updateTokenIfNeeded(options: GetTokenOptions): Promise<void> {
-    if (this.tokenRefresher.isReady()) {
-      const accessToken = await this.tokenRefresher.refresh(options);
-      this.tokenCache.setCachedToken(accessToken);
+  private async getToken(options: GetTokenOptions): Promise<string | undefined> {
+    // We reset the cached token in a time window before it expires,
+    // after that point, we retry the refresh of the token only if the token refresher is ready.
+
+    let token = this.tokenCache.getCachedToken();
+    let refreshPromise: Promise<AccessToken | undefined> | undefined;
+    if (!token && this.tokenRefresher.isReady()) {
+      refreshPromise = this.tokenRefresher.refresh(options);
     }
-  }
 
-  private async getToken(options: GetTokenOptions): Promise<string | undefined> {
-    let accessToken = this.tokenCache.getCachedToken();
-    if (accessToken === undefined) {
-      // Waiting for the next refresh only if the cache is unable to retrieve the access token,
-      // which means that it has expired, or it has never been set.
-      accessToken = await this.tokenRefresher.refresh(options);
-      this.tokenCache.setCachedToken(accessToken);
-    } else {
-      // If we still have a cached access token,
-      // And any other time related conditionals have been reached based on the tokenRefresher class,
-      // then attempt to refresh without waiting.
-      this.updateTokenIfNeeded(options);
+    // We don't use the local copy if the token cache has a token that has an expiration date greater than the local copy.
+    if (this.token && this.token.expiresOnTimestamp > (token?.expiresOnTimestamp || 0)) {
+      token = this.token;
     }
 
-    return accessToken ? accessToken.token : undefined;
+    // If we don't have a token, or if the token has expired, we wait for the refresh promise
+    if (!token || token.expiresOnTimestamp < Date.now()) {
+      token = await refreshPromise;
+    }
+
+    // We save the token we were able to retrieve.
+    this.tokenCache.setCachedToken(token);
+    this.token = token;
+
+    return token ? token.token : undefined;
   }
 }

sdk/core/core-http/test/policies/bearerTokenAuthenticationPolicyTests.ts

@@ -2,21 +2,24 @@
 // Licensed under the MIT license.
 
 import { assert } from "chai";
-import { fake, createSandbox } from "sinon";
+import Sinon, { fake, createSandbox } from "sinon";
 import { OperationSpec } from "../../src/operationSpec";
-import { TokenCredential, GetTokenOptions, AccessToken } from "@azure/core-auth";
+import { TokenCredential, AccessToken } from "@azure/core-auth";
 import { RequestPolicy, RequestPolicyOptions } from "../../src/policies/requestPolicy";
 import { Constants } from "../../src/util/constants";
 import { HttpOperationResponse } from "../../src/httpOperationResponse";
 import { HttpHeaders } from "../../src/httpHeaders";
 import { WebResource } from "../../src/webResource";
 import { BearerTokenAuthenticationPolicy } from "../../src/policies/bearerTokenAuthenticationPolicy";
-import {
-  ExpiringAccessTokenCache,
-  TokenRefreshBufferMs
-} from "../../src/credentials/accessTokenCache";
+import { ExpiringAccessTokenCache } from "../../src/credentials/accessTokenCache";
 import { AccessTokenRefresher } from "../../src/credentials/accessTokenRefresher";
 
+// Token is refreshed one millisecond BEFORE it expires.
+const beforeTokenExpiresInMs = 1000;
+
+// To avoid many refresh requests, we make new refresh requests only after:
+const tokenRefreshIntervalInMs = 500;
+
 describe("BearerTokenAuthenticationPolicy", function() {
   const mockPolicy: RequestPolicy = {
     sendRequest(request: WebResource): Promise<HttpOperationResponse> {
@@ -28,6 +31,15 @@ describe("BearerTokenAuthenticationPolicy", function() {
     }
   };
 
+  let clock: Sinon.SinonFakeTimers;
+
+  beforeEach(() => {
+    clock = Sinon.useFakeTimers(Date.now());
+  });
+  afterEach(() => {
+    clock.restore();
+  });
+
   it("correctly adds an Authentication header with the Bearer token", async function() {
     const mockToken = "token";
     const tokenScopes = ["scope1", "scope2"];
@@ -53,32 +65,11 @@ describe("BearerTokenAuthenticationPolicy", function() {
     );
   });
 
-  it("refreshes access tokens when they expire", async () => {
-    const now = Date.now();
-    const refreshCred1 = new MockRefreshAzureCredential(now);
-    const refreshCred2 = new MockRefreshAzureCredential(now + TokenRefreshBufferMs);
-    const notRefreshCred1 = new MockRefreshAzureCredential(now + TokenRefreshBufferMs + 5000);
-
-    const credentialsToTest: [MockRefreshAzureCredential, number, string][] = [
-      [refreshCred1, 2, "Refresh credential 1"],
-      [refreshCred2, 2, "Refresh credential 2"],
-      [notRefreshCred1, 1, "Not refresh credential 1"]
-    ];
-
-    const request = createRequest();
-    for (const [credentialToTest, expectedCalls, message] of credentialsToTest) {
-      const policy = createBearerTokenPolicy("testscope", credentialToTest);
-      await policy.sendRequest(request);
-      await policy.sendRequest(request);
-      assert.strictEqual(credentialToTest.authCount, expectedCalls, `${message} failed`);
-    }
-  });
-
   it("tests that AccessTokenRefresher is working", async function() {
     const now = Date.now();
     const credentialToTest = new MockRefreshAzureCredential(now);
     const request = createRequest();
-    const policy = createBearerTokenPolicy("testscope", credentialToTest);
+    const policy = createBearerTokenPolicy("test-scope", credentialToTest);
     await policy.sendRequest(request);
 
     const sandbox = createSandbox();
@@ -88,6 +79,135 @@ describe("BearerTokenAuthenticationPolicy", function() {
     assert.strictEqual(credentialToTest.authCount, 2);
   });
 
+  it("refreshes the token on initial request", async () => {
+    const expiresOn = Date.now() + 1000 * 60; // One minute later.
+    const credential = new MockRefreshAzureCredential(expiresOn);
+    const request = createRequest();
+    const policy = createBearerTokenPolicy("test-scope", credential);
+
+    await policy.sendRequest(request);
+    assert.strictEqual(credential.authCount, 1);
+  });
+
+  it("refreshes the token during the refresh window", async () => {
+    const expireDelayMs = 5000;
+    let tokenExpiration = Date.now() + expireDelayMs;
+    const credential = new MockRefreshAzureCredential(tokenExpiration);
+    const request = createRequest();
+    const policy = createBearerTokenPolicy("test-scope", credential);
+
+    // The token is cached and remains cached for a bit.
+    await policy.sendRequest(request);
+    await policy.sendRequest(request);
+    assert.strictEqual(credential.authCount, 1);
+
+    // The token will remain cached until tokenExpiration - testTokenRefreshBufferMs, so in (5000 - 1000) milliseconds.
+
+    // For safe measure, we test the token is still cached a second earlier than the refresh expectation.
+    clock.tick(expireDelayMs - beforeTokenExpiresInMs - 1000);
+    await policy.sendRequest(request);
+    assert.strictEqual(credential.authCount, 1);
+
+    // The new token will last 5 seconds again.
+    tokenExpiration = Date.now() + expireDelayMs;
+    credential.expiresOnTimestamp = tokenExpiration;
+
+    // Now we wait until it expires:
+    clock.tick(1000);
+    await policy.sendRequest(request);
+    assert.strictEqual(credential.authCount, 2);
+  });
+
+  it("access token refresher should prevent multiple initial getToken requests to break", async () => {
+    const expireDelayMs = 5000;
+    const startTime = Date.now();
+    const tokenExpiration = startTime + expireDelayMs;
+    const getTokenDelay = 100;
+    const credential = new MockRefreshAzureCredential(tokenExpiration, getTokenDelay, clock);
+    const request = createRequest();
+    const policy = createBearerTokenPolicy("test-scope", credential);
+
+    // Now we send some requests.
+    const promises = [
+      policy.sendRequest(request),
+      policy.sendRequest(request),
+      policy.sendRequest(request)
+    ];
+    // Now we wait until they're all resolved.
+    for (const promise of promises) {
+      await promise;
+    }
+    assert.strictEqual(credential.authCount, 1, "The first authentication should have happened");
+  });
+
+  it("credential errors should bubble up", async () => {
+    const expireDelayMs = 5000;
+    const startTime = Date.now();
+    const tokenExpiration = startTime + expireDelayMs;
+    const getTokenDelay = 100;
+    const credential = new MockRefreshAzureCredential(tokenExpiration, getTokenDelay, clock);
+    const request = createRequest();
+    const policy = createBearerTokenPolicy("test-scope", credential);
+
+    credential.shouldThrow = true;
+
+    let error: Error | undefined;
+    try {
+      await policy.sendRequest(request);
+    } catch (e) {
+      error = e;
+    }
+    assert.equal(error?.message, "Failed to retrieve the token");
+
+    assert.strictEqual(
+      credential.authCount,
+      1,
+      "The first authentication attempt should have happened"
+    );
+  });
+
+  it("access token refresher should prevent refreshers to happen too fast while the token is about to expire", async () => {
+    const expireDelayMs = 5000;
+    const startTime = Date.now();
+    const tokenExpiration = startTime + expireDelayMs;
+    const getTokenDelay = 100;
+    const credential = new MockRefreshAzureCredential(tokenExpiration, getTokenDelay, clock);
+    const request = createRequest();
+    const policy = createBearerTokenPolicy("test-scope", credential);
+
+    await policy.sendRequest(request);
+    assert.strictEqual(credential.authCount, 1, "The first authentication should have happened");
+
+    clock.tick(expireDelayMs - beforeTokenExpiresInMs); // Until we start refreshing the token
+
+    // Now we send some requests.
+    const promises = [
+      policy.sendRequest(request),
+      policy.sendRequest(request),
+      policy.sendRequest(request)
+    ];
+
+    // Now we wait until they're all resolved.
+    for (const promise of promises) {
+      await promise;
+    }
+
+    // Only getTokenDelay should have passed, and only one refresh should have happened.
+    assert.strictEqual(
+      credential.authCount,
+      2,
+      "authCode should have been called once during the refresh time"
+    );
+
+    const exceptionMessage =
+      "the total time passed should be in the refresh room, plus the many getTokens that have happened so far";
+    assert.equal(
+      expireDelayMs - beforeTokenExpiresInMs + getTokenDelay * 2,
+      Date.now() - startTime,
+      exceptionMessage
+    );
+  });
+
   function createRequest(operationSpec?: OperationSpec): WebResource {
     const request = new WebResource();
     request.operationSpec = operationSpec;
@@ -101,25 +221,34 @@ describe("BearerTokenAuthenticationPolicy", function() {
     return new BearerTokenAuthenticationPolicy(
       mockPolicy,
       new RequestPolicyOptions(),
-      new ExpiringAccessTokenCache(),
-      new AccessTokenRefresher(credential, scopes)
+      new ExpiringAccessTokenCache(beforeTokenExpiresInMs),
+      new AccessTokenRefresher(credential, scopes, tokenRefreshIntervalInMs)
     );
   }
 });
 
 class MockRefreshAzureCredential implements TokenCredential {
-  private _expiresOnTimestamp: number;
   public authCount = 0;
+  public shouldThrow: boolean = false;
 
-  constructor(expiresOnTimestamp: number) {
-    this._expiresOnTimestamp = expiresOnTimestamp;
-  }
+  constructor(
+    public expiresOnTimestamp: number,
+    public getTokenDelay?: number,
+    public clock?: sinon.SinonFakeTimers
+  ) {}
 
-  public getToken(
-    _scopes: string | string[],
-    _options?: GetTokenOptions
-  ): Promise<AccessToken | null> {
+  public async getToken(): Promise<AccessToken> {
     this.authCount++;
-    return Promise.resolve({ token: "mocktoken", expiresOnTimestamp: this._expiresOnTimestamp });
+
+    if (this.shouldThrow) {
+      throw new Error("Failed to retrieve the token");
+    }
+
+    // Allowing getToken to take a while
+    if (this.getTokenDelay && this.clock) {
+      this.clock.tick(this.getTokenDelay);
+    }
+
+    return { token: "mock-token", expiresOnTimestamp: this.expiresOnTimestamp };
   }
 }