Skip to content

Commit 182103f

Browse files
authored
fix(auth): prevent concurrent token refresh and adding necessary logs (#8294)
## Problem - There was a significant increase in call volume for the sso-oidc CreateToken operation from AmazonQ-For-VSCode clients. ## Solution - The root cause was identified as a race condition in concurrent asynchronous token refresh operations within the VSCode extension. When bearer tokens expired (after 1 hour), multiple simultaneous calls to getToken() independently triggered token refresh operations, resulting in duplicate OIDC API calls. - Implemented a "promise deduplication" pattern to prevent duplicate async OIDC calls. --- - Treat all work as PUBLIC. Private `feature/x` branches will not be squash-merged at release time. - Your code changes must meet the guidelines in [CONTRIBUTING.md](https://github.com/aws/aws-toolkit-vscode/blob/master/CONTRIBUTING.md#guidelines). - License: I confirm that my contribution is made under the terms of the Apache 2.0 license.
1 parent a450753 commit 182103f

File tree

1 file changed

+54
-2
lines changed

1 file changed

+54
-2
lines changed

packages/core/src/auth/sso/ssoAccessTokenProvider.ts

Lines changed: 54 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -59,6 +59,15 @@ export abstract class SsoAccessTokenProvider {
5959
private static logIfChanged = onceChanged((s: string) => getLogger().info(s))
6060
private readonly className = 'SsoAccessTokenProvider'
6161

62+
/**
63+
* Prevents concurrent token refresh operations.
64+
* Maps tokenCacheKey to an in-flight refresh promise.
65+
*/
66+
private static refreshPromises = new Map<
67+
string,
68+
Promise<{ token: SsoToken; registration: ClientRegistration; region: string; startUrl: string }>
69+
>()
70+
6271
public static set authSource(val: string) {
6372
SsoAccessTokenProvider._authSource = val
6473
}
@@ -108,15 +117,43 @@ export abstract class SsoAccessTokenProvider {
108117
true
109118
)
110119
)
120+
111121
if (!data || !isExpired(data.token)) {
122+
getLogger().debug('Auth: token is valid, returning cached token (key=%s)', this.tokenCacheKey)
112123
return data?.token
113124
}
114125

126+
getLogger().info(
127+
`Auth: bearer token expired (expires at ${data.token.expiresAt}), attempting refresh (key=${this.tokenCacheKey})`
128+
)
129+
115130
if (data.registration && !isExpired(data.registration) && hasProps(data.token, 'refreshToken')) {
116-
const refreshed = await this.refreshToken(data.token, data.registration)
131+
getLogger().debug(`Auth: refresh token available, calling refreshToken() (key=${this.tokenCacheKey})`)
132+
// Check if a refresh is already in progress for this token
133+
const existingRefresh = SsoAccessTokenProvider.refreshPromises.get(this.tokenCacheKey)
134+
if (existingRefresh) {
135+
getLogger().debug(
136+
'SsoAccessTokenProvider: Token refresh already in progress, waiting for existing refresh'
137+
)
138+
const refreshed = await existingRefresh
139+
return refreshed.token
140+
}
141+
142+
// Start a new refresh and store the promise
143+
const refreshPromise = this.refreshToken(data.token, data.registration)
144+
SsoAccessTokenProvider.refreshPromises.set(this.tokenCacheKey, refreshPromise)
117145

118-
return refreshed.token
146+
try {
147+
const refreshed = await refreshPromise
148+
return refreshed.token
149+
} finally {
150+
// Clean up the promise from the map once complete (success or failure)
151+
SsoAccessTokenProvider.refreshPromises.delete(this.tokenCacheKey)
152+
}
119153
} else {
154+
getLogger().warn(
155+
`getToken: cannot refresh - registration expired or no refresh token available (key=${this.tokenCacheKey})`
156+
)
120157
await this.invalidate('allCacheExpired')
121158
}
122159
}
@@ -172,10 +209,18 @@ export abstract class SsoAccessTokenProvider {
172209

173210
try {
174211
const clientInfo = selectFrom(registration, 'clientId', 'clientSecret')
212+
getLogger().debug(`Auth refreshToken: calling OIDC createToken API (key=${this.tokenCacheKey})`)
175213
const response = await this.oidc.createToken({ ...clientInfo, ...token, grantType: refreshGrantType })
214+
215+
getLogger().debug(`Auth refreshToken: got response, now saving to cache...`)
216+
176217
const refreshed = this.formatToken(response, registration)
218+
getLogger().debug(`refreshToken: saving refreshed token to cache (key=${this.tokenCacheKey})`)
177219
await this.cache.token.save(this.tokenCacheKey, refreshed)
178220

221+
getLogger().info(
222+
`Auth refreshToken: token refresh successful (key=${this.tokenCacheKey}, new expiry=${response.expiresAt})`
223+
)
179224
telemetry.aws_refreshCredentials.emit({
180225
result: 'Succeeded',
181226
requestId: response.requestId,
@@ -184,6 +229,10 @@ export abstract class SsoAccessTokenProvider {
184229

185230
return refreshed
186231
} catch (err) {
232+
getLogger().error(
233+
`Auth refreshToken: token refresh failed (key=${this.tokenCacheKey}): ${getErrorMsg(err as unknown as Error)}`
234+
)
235+
187236
if (err instanceof DiskCacheError) {
188237
/**
189238
* Background:
@@ -197,6 +246,9 @@ export abstract class SsoAccessTokenProvider {
197246
* to the logs where the error was logged. Hopefully they can use this information to fix the issue,
198247
* or at least hint for them to provide the logs in a bug report.
199248
*/
249+
getLogger().warn(
250+
`Auth refreshToken: DiskCacheError during refresh, not invalidating session (key=${this.tokenCacheKey})`
251+
)
200252
void DiskCacheErrorMessage.instance.showMessageThrottled(err)
201253
} else if (!isNetworkError(err)) {
202254
const reason = getTelemetryReason(err)

0 commit comments

Comments
 (0)