From 72cae436b7663838a363aa4ca00fc2bfd75c88b4 Mon Sep 17 00:00:00 2001 From: Jeff Fisher Date: Mon, 13 May 2024 17:48:24 -0500 Subject: [PATCH] [core-rest-pipeline] Improve robustness of tokenCycler (#29638) ### Packages impacted by this PR `core-rest-pipeline` `ts-http-runtime` ### Issues associated with this PR https://github.com/Azure/azure-sdk-for-js/issues/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. --- .../src/util/tokenCycler.ts | 16 +++- .../bearerTokenAuthenticationPolicy.spec.ts | 76 ++++++++++++++++++- .../ts-http-runtime/src/util/tokenCycler.ts | 16 +++- .../bearerTokenAuthenticationPolicy.spec.ts | 70 +++++++++++++++++ 4 files changed, 171 insertions(+), 7 deletions(-) diff --git a/sdk/core/core-rest-pipeline/src/util/tokenCycler.ts b/sdk/core/core-rest-pipeline/src/util/tokenCycler.ts index d048647bf873..1da185de3518 100644 --- 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 index 53d6a8b53e22..8097e90a304c 100644 --- 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>(); + + // 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 { + // 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, ReturnType>(); + + // 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 index aad1f68d7967..73b4e36954fe 100644 --- 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 index 7e24435de19a..392fc5f2a882 100644 --- 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>(); + + // 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 { + // 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, ReturnType>(); + + // 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,