-
Notifications
You must be signed in to change notification settings - Fork 1.2k
/
StorageBearerTokenChallengeAuthenticationPolicy.ts
351 lines (314 loc) · 10.8 KB
/
StorageBearerTokenChallengeAuthenticationPolicy.ts
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
// Copyright (c) Microsoft Corporation.
// Licensed under the MIT license.
import { AccessToken, GetTokenOptions, TokenCredential, URLBuilder } from "@azure/core-http";
import {
BaseRequestPolicy,
RequestPolicy,
RequestPolicyFactory,
RequestPolicyOptions,
} from "@azure/core-http";
import { HttpOperationResponse } from "@azure/core-http";
import { WebResourceLike } from "@azure/core-http";
import { delay } from "@azure/core-http";
/**
* A set of constants used internally when processing requests.
*/
const Constants = {
DefaultScope: "/.default",
/**
* Defines constants for use with HTTP headers.
*/
HeaderConstants: {
/**
* The Authorization header.
*/
AUTHORIZATION: "authorization",
},
};
// #region Access Token Cycler
/**
* A function that gets a promise of an access token and allows providing
* options.
*
* @param options - the options to pass to the underlying token provider
*/
type AccessTokenGetter = (options: GetTokenOptions) => Promise<AccessToken>;
interface TokenCyclerOptions {
/**
* The window of time before token expiration during which the token will be
* considered unusable due to risk of the token expiring before sending the
* request.
*
* This will only become meaningful if the refresh fails for over
* (refreshWindow - forcedRefreshWindow) milliseconds.
*/
forcedRefreshWindowInMs: number;
/**
* Interval in milliseconds to retry failed token refreshes.
*/
retryIntervalInMs: number;
/**
* The window of time before token expiration during which
* we will attempt to refresh the token.
*/
refreshWindowInMs: number;
}
// Default options for the cycler if none are provided
const DEFAULT_CYCLER_OPTIONS: TokenCyclerOptions = {
forcedRefreshWindowInMs: 1000, // Force waiting for a refresh 1s before the token expires
retryIntervalInMs: 3000, // Allow refresh attempts every 3s
refreshWindowInMs: 1000 * 60 * 2, // Start refreshing 2m before expiry
};
/**
* Converts an an unreliable access token getter (which may resolve with null)
* into an AccessTokenGetter by retrying the unreliable getter in a regular
* interval.
*
* @param getAccessToken - a function that produces a promise of an access
* token that may fail by returning null
* @param retryIntervalInMs - the time (in milliseconds) to wait between retry
* attempts
* @param timeoutInMs - the timestamp after which the refresh attempt will fail,
* throwing an exception
* @returns - a promise that, if it resolves, will resolve with an access token
*/
async function beginRefresh(
getAccessToken: () => Promise<AccessToken | null>,
retryIntervalInMs: number,
timeoutInMs: number
): Promise<AccessToken> {
// This wrapper handles exceptions gracefully as long as we haven't exceeded
// the timeout.
async function tryGetAccessToken(): Promise<AccessToken | null> {
if (Date.now() < timeoutInMs) {
try {
return await getAccessToken();
} catch {
return null;
}
} else {
const finalToken = await getAccessToken();
// Timeout is up, so throw if it's still null
if (finalToken === null) {
throw new Error("Failed to refresh access token.");
}
return finalToken;
}
}
let token: AccessToken | null = await tryGetAccessToken();
while (token === null) {
await delay(retryIntervalInMs);
token = await tryGetAccessToken();
}
return token;
}
/**
* Creates a token cycler from a credential, scopes, and optional settings.
*
* A token cycler represents a way to reliably retrieve a valid access token
* from a TokenCredential. It will handle initializing the token, refreshing it
* when it nears expiration, and synchronizes refresh attempts to avoid
* concurrency hazards.
*
* @param credential - the underlying TokenCredential that provides the access
* token
* @param scopes - the scopes to request authorization for
* @param tokenCyclerOptions - optionally override default settings for the cycler
*
* @returns - a function that reliably produces a valid access token
*/
function createTokenCycler(
credential: TokenCredential,
scopes: string | string[],
tokenCyclerOptions?: Partial<TokenCyclerOptions>
): AccessTokenGetter {
let refreshWorker: Promise<AccessToken> | null = null;
let token: AccessToken | null = null;
const options = {
...DEFAULT_CYCLER_OPTIONS,
...tokenCyclerOptions,
};
/**
* This little holder defines several predicates that we use to construct
* the rules of refreshing the token.
*/
const cycler = {
/**
* Produces true if a refresh job is currently in progress.
*/
get isRefreshing(): boolean {
return refreshWorker !== null;
},
/**
* Produces true if the cycler SHOULD refresh (we are within the refresh
* window and not already refreshing)
*/
get shouldRefresh(): boolean {
return (
!cycler.isRefreshing &&
(token?.expiresOnTimestamp ?? 0) - options.refreshWindowInMs < Date.now()
);
},
/**
* Produces true if the cycler MUST refresh (null or nearly-expired
* token).
*/
get mustRefresh(): boolean {
return (
token === null || token.expiresOnTimestamp - options.forcedRefreshWindowInMs < Date.now()
);
},
};
/**
* Starts a refresh job or returns the existing job if one is already
* running.
*/
function refresh(getTokenOptions: GetTokenOptions): Promise<AccessToken> {
if (!cycler.isRefreshing) {
// We bind `scopes` here to avoid passing it around a lot
const tryGetAccessToken = (): Promise<AccessToken | null> =>
credential.getToken(scopes, getTokenOptions);
// Take advantage of promise chaining to insert an assignment to `token`
// before the refresh can be considered done.
refreshWorker = beginRefresh(
tryGetAccessToken,
options.retryIntervalInMs,
// If we don't have a token, then we should timeout immediately
token?.expiresOnTimestamp ?? Date.now()
)
.then((_token) => {
refreshWorker = null;
token = _token;
return token;
})
.catch((reason) => {
// We also should reset the refresher if we enter a failed state. All
// existing awaiters will throw, but subsequent requests will start a
// new retry chain.
refreshWorker = null;
token = null;
throw reason;
});
}
return refreshWorker as Promise<AccessToken>;
}
return async (tokenOptions: GetTokenOptions): Promise<AccessToken> => {
//
// Simple rules:
// - If we MUST refresh, then return the refresh task, blocking
// the pipeline until a token is available.
// - If we SHOULD refresh, then run refresh but don't return it
// (we can still use the cached token).
// - Return the token, since it's fine if we didn't return in
// step 1.
//
if (cycler.mustRefresh) return refresh(tokenOptions);
if (cycler.shouldRefresh) {
refresh(tokenOptions);
}
return token as AccessToken;
};
}
/**
* We will retrieve the challenge only if the response status code was 401,
* and if the response contained the header "WWW-Authenticate" with a non-empty value.
*/
function getChallenge(response: HttpOperationResponse): string | undefined {
const challenge = response.headers.get("WWW-Authenticate");
if (response.status === 401 && challenge) {
return challenge;
}
return;
}
/**
* Challenge structure
*/
interface Challenge {
authorization_uri: string;
resource_id: string;
}
/**
* Converts: `Bearer a="b" c="d"`.
* Into: `[ { a: 'b', c: 'd' }]`.
*
* @internal
*/
function parseChallenge(challenge: string): any {
const bearerChallenge = challenge.slice("Bearer ".length);
const challengeParts = `${bearerChallenge.trim()} `.split(" ").filter((x) => x);
const keyValuePairs = challengeParts.map((keyValue) =>
(([key, value]) => ({ [key]: value }))(keyValue.trim().split("="))
);
// Key-value pairs to plain object:
return keyValuePairs.reduce((a, b) => ({ ...a, ...b }), {});
}
// #endregion
/**
* Creates a new factory for a RequestPolicy that applies a bearer token to
* the requests' `Authorization` headers.
*
* @param credential - The TokenCredential implementation that can supply the bearer token.
* @param scopes - The scopes for which the bearer token applies.
*/
export function storageBearerTokenChallengeAuthenticationPolicy(
credential: TokenCredential,
scopes: string | string[]
): RequestPolicyFactory {
// This simple function encapsulates the entire process of reliably retrieving the token
let getToken = createTokenCycler(credential, scopes);
class StorageBearerTokenChallengeAuthenticationPolicy extends BaseRequestPolicy {
public constructor(nextPolicy: RequestPolicy, options: RequestPolicyOptions) {
super(nextPolicy, options);
}
public async sendRequest(webResource: WebResourceLike): Promise<HttpOperationResponse> {
if (!webResource.url.toLowerCase().startsWith("https://")) {
throw new Error(
"Bearer token authentication is not permitted for non-TLS protected (non-https) URLs."
);
}
const getTokenInternal = getToken;
const token = (
await getTokenInternal({
abortSignal: webResource.abortSignal,
tracingOptions: {
tracingContext: webResource.tracingContext,
},
})
).token;
webResource.headers.set(Constants.HeaderConstants.AUTHORIZATION, `Bearer ${token}`);
const response = await this._nextPolicy.sendRequest(webResource);
if (response?.status === 401) {
const challenge = getChallenge(response);
if (challenge) {
const challengeInfo: Challenge = parseChallenge(challenge);
const challengeScopes = challengeInfo.resource_id + Constants.DefaultScope;
const parsedAuthUri = URLBuilder.parse(challengeInfo.authorization_uri);
const pathSegments = parsedAuthUri.getPath()!.split("/");
const tenantId = pathSegments[1];
const getTokenForChallenge = createTokenCycler(credential, challengeScopes);
const tokenForChallenge = (
await getTokenForChallenge({
abortSignal: webResource.abortSignal,
tracingOptions: {
tracingContext: webResource.tracingContext,
},
tenantId: tenantId,
})
).token;
getToken = getTokenForChallenge;
webResource.headers.set(
Constants.HeaderConstants.AUTHORIZATION,
`Bearer ${tokenForChallenge}`
);
return this._nextPolicy.sendRequest(webResource);
}
}
return response;
}
}
return {
create: (nextPolicy: RequestPolicy, options: RequestPolicyOptions) => {
return new StorageBearerTokenChallengeAuthenticationPolicy(nextPolicy, options);
},
};
}