[core-rest-pipeline] Improve robustness of tokenCycler (#29638)

### Packages impacted by this PR `core-rest-pipeline` `ts-http-runtime` ### Issues associated with this PR #29608 ### Describe the problem that is addressed by this PR When debugging a recent issue from Storage, it was reported that in the case of the tokenCycler being unable to retrieve a new access token from a CAE challenge, it would fall back to delaying for the maximum lifetime of the current (invalid) access token instead of failing immediately. This change improves the robustness of this scenario by invaliding the existing cached token when claims are received. ### Are there test cases added in this PR? _(If not, why?)_ Yes, I added a test and confirmed it was timing out before the fix was added.
Azure · May 13, 2024 · 72cae43 · 72cae43
1 parent fa4fd3a
commit 72cae43
Show file tree

Hide file tree

Showing 4 changed files with 171 additions and 7 deletions.
diff --git a/sdk/core/core-rest-pipeline/src/util/tokenCycler.ts b/sdk/core/core-rest-pipeline/src/util/tokenCycler.ts
@@ -201,13 +201,23 @@ export function createTokenCycler(
     //   step 1.
     //
 
+    const hasClaimChallenge = Boolean(tokenOptions.claims);
+    const tenantIdChanged = tenantId !== tokenOptions.tenantId;
+
+    if (hasClaimChallenge) {
+      // If we've received a claim, we know the existing token isn't valid
+      // We want to clear it so that that refresh worker won't use the old expiration time as a timeout
+      token = null;
+    }
+
     // If the tenantId passed in token options is different to the one we have
     // Or if we are in claim challenge and the token was rejected and a new access token need to be issued, we need to
     // refresh the token with the new tenantId or token.
-    const mustRefresh =
-      tenantId !== tokenOptions.tenantId || Boolean(tokenOptions.claims) || cycler.mustRefresh;
+    const mustRefresh = tenantIdChanged || hasClaimChallenge || cycler.mustRefresh;
 
-    if (mustRefresh) return refresh(scopes, tokenOptions);
+    if (mustRefresh) {
+      return refresh(scopes, tokenOptions);
+    }
 
     if (cycler.shouldRefresh) {
       refresh(scopes, tokenOptions);

diff --git a/sdk/core/core-rest-pipeline/test/bearerTokenAuthenticationPolicy.spec.ts b/sdk/core/core-rest-pipeline/test/bearerTokenAuthenticationPolicy.spec.ts
@@ -2,7 +2,12 @@
 // Licensed under the MIT license.
 
 import { AccessToken, TokenCredential } from "@azure/core-auth";
-import type { PipelinePolicy, PipelineResponse, SendRequest } from "../src/index.js";
+import type {
+  AuthorizeRequestOnChallengeOptions,
+  PipelinePolicy,
+  PipelineResponse,
+  SendRequest,
+} from "../src/index.js";
 import {
   bearerTokenAuthenticationPolicy,
   createHttpHeaders,
@@ -238,6 +243,75 @@ describe("BearerTokenAuthenticationPolicy", function () {
     );
   });
 
+  it("does not wait for token expiry after receiving an authentication challenge", async () => {
+    // tokens can live for a long time
+    const oneDayInMs = 24 * 60 * 60 * 1000;
+    const tokenExpiration = Date.now() + oneDayInMs;
+    const getToken = vi.fn<[], Promise<AccessToken | null>>();
+
+    // first time getToken is called to put the header on the initial request
+    getToken.mockResolvedValueOnce({
+      token: "mock-token",
+      expiresOnTimestamp: tokenExpiration,
+    });
+    // simulate failure of retriving the token, rejecting with an error would also work
+    // but returning null exercises a slightly different code path for better coverage
+    getToken.mockResolvedValueOnce(null);
+    const credential: TokenCredential = {
+      getToken,
+    };
+    const scopes = ["test-scope"];
+    const request = createPipelineRequest({ url: "https://example.com" });
+
+    async function authorizeRequestOnChallenge(
+      options: AuthorizeRequestOnChallengeOptions,
+    ): Promise<boolean> {
+      // this will trigger a second call into getToken, which should fail
+      // what we don't want is to wait for expiresOnTimestamp of the original credential before giving up
+      const token = await options.getAccessToken(scopes, {
+        claims: '{"access_token":{"nbf":{"essential":true, "value":"1603742800"}}}',
+      });
+      if (token) {
+        options.request.headers.set("Authorization", `Bearer ${token.token}`);
+        return true;
+      }
+      return false;
+    }
+
+    const policy = bearerTokenAuthenticationPolicy({
+      scopes,
+      credential,
+      challengeCallbacks: {
+        authorizeRequestOnChallenge,
+      },
+    });
+    const next = vi.fn<Parameters<SendRequest>, ReturnType<SendRequest>>();
+
+    // first response is an auth challenge
+    const challengeResponse: PipelineResponse = {
+      headers: createHttpHeaders({
+        "WWW-Authenticate": [
+          `Bearer authorization_uri="https://login.windows-ppe.net/", error="invalid_token"`,
+          `error_description="User session has been revoked"`,
+          `claims="eyJhY2Nlc3NfdG9rZW4iOnsibmJmIjp7ImVzc2VudGlhbCI6dHJ1ZSwgInZhbHVlIjoiMTYwMzc0MjgwMCJ9fX0="`,
+        ].join(", "),
+      }),
+      request,
+      status: 401,
+    };
+
+    next.mockResolvedValueOnce(challengeResponse);
+
+    let error: Error | undefined;
+    try {
+      await policy.sendRequest(request, next);
+    } catch (e: any) {
+      error = e;
+    }
+    assert.isDefined(error);
+    assert.equal(error?.message, "Failed to refresh access token.");
+  });
+
   function createBearerTokenPolicy(
     scopes: string | string[],
     credential: TokenCredential,

diff --git a/sdk/core/ts-http-runtime/src/util/tokenCycler.ts b/sdk/core/ts-http-runtime/src/util/tokenCycler.ts
@@ -201,13 +201,23 @@ export function createTokenCycler(
     //   step 1.
     //
 
+    const hasClaimChallenge = Boolean(tokenOptions.claims);
+    const tenantIdChanged = tenantId !== tokenOptions.tenantId;
+
+    if (hasClaimChallenge) {
+      // If we've received a claim, we know the existing token isn't valid
+      // We want to clear it so that that refresh worker won't use the old expiration time as a timeout
+      token = null;
+    }
+
     // If the tenantId passed in token options is different to the one we have
     // Or if we are in claim challenge and the token was rejected and a new access token need to be issued, we need to
     // refresh the token with the new tenantId or token.
-    const mustRefresh =
-      tenantId !== tokenOptions.tenantId || Boolean(tokenOptions.claims) || cycler.mustRefresh;
+    const mustRefresh = tenantIdChanged || hasClaimChallenge || cycler.mustRefresh;
 
-    if (mustRefresh) return refresh(scopes, tokenOptions);
+    if (mustRefresh) {
+      return refresh(scopes, tokenOptions);
+    }
 
     if (cycler.shouldRefresh) {
       refresh(scopes, tokenOptions);

diff --git a/sdk/core/ts-http-runtime/test/bearerTokenAuthenticationPolicy.spec.ts b/sdk/core/ts-http-runtime/test/bearerTokenAuthenticationPolicy.spec.ts
@@ -4,6 +4,7 @@
 import { describe, it, assert, expect, vi, beforeEach, afterEach } from "vitest";
 import { AccessToken, TokenCredential } from "../src/auth/tokenCredential.js";
 import {
+  AuthorizeRequestOnChallengeOptions,
   PipelinePolicy,
   PipelineResponse,
   SendRequest,
@@ -240,6 +241,75 @@ describe("BearerTokenAuthenticationPolicy", function () {
     );
   });
 
+  it("does not wait for token expiry after receiving an authentication challenge", async () => {
+    // tokens can live for a long time
+    const oneDayInMs = 24 * 60 * 60 * 1000;
+    const tokenExpiration = Date.now() + oneDayInMs;
+    const getToken = vi.fn<[], Promise<AccessToken | null>>();
+
+    // first time getToken is called to put the header on the initial request
+    getToken.mockResolvedValueOnce({
+      token: "mock-token",
+      expiresOnTimestamp: tokenExpiration,
+    });
+    // simulate failure of retriving the token, rejecting with an error would also work
+    // but returning null exercises a slightly different code path for better coverage
+    getToken.mockResolvedValueOnce(null);
+    const credential: TokenCredential = {
+      getToken,
+    };
+    const scopes = ["test-scope"];
+    const request = createPipelineRequest({ url: "https://example.com" });
+
+    async function authorizeRequestOnChallenge(
+      options: AuthorizeRequestOnChallengeOptions,
+    ): Promise<boolean> {
+      // this will trigger a second call into getToken, which should fail
+      // what we don't want is to wait for expiresOnTimestamp of the original credential before giving up
+      const token = await options.getAccessToken(scopes, {
+        claims: '{"access_token":{"nbf":{"essential":true, "value":"1603742800"}}}',
+      });
+      if (token) {
+        options.request.headers.set("Authorization", `Bearer ${token.token}`);
+        return true;
+      }
+      return false;
+    }
+
+    const policy = bearerTokenAuthenticationPolicy({
+      scopes,
+      credential,
+      challengeCallbacks: {
+        authorizeRequestOnChallenge,
+      },
+    });
+    const next = vi.fn<Parameters<SendRequest>, ReturnType<SendRequest>>();
+
+    // first response is an auth challenge
+    const challengeResponse: PipelineResponse = {
+      headers: createHttpHeaders({
+        "WWW-Authenticate": [
+          `Bearer authorization_uri="https://login.windows-ppe.net/", error="invalid_token"`,
+          `error_description="User session has been revoked"`,
+          `claims="eyJhY2Nlc3NfdG9rZW4iOnsibmJmIjp7ImVzc2VudGlhbCI6dHJ1ZSwgInZhbHVlIjoiMTYwMzc0MjgwMCJ9fX0="`,
+        ].join(", "),
+      }),
+      request,
+      status: 401,
+    };
+
+    next.mockResolvedValueOnce(challengeResponse);
+
+    let error: Error | undefined;
+    try {
+      await policy.sendRequest(request, next);
+    } catch (e: any) {
+      error = e;
+    }
+    assert.isDefined(error);
+    assert.equal(error?.message, "Failed to refresh access token.");
+  });
+
   function createBearerTokenPolicy(
     scopes: string | string[],
     credential: TokenCredential,