Skip to content

Commit 0d6d35e

Browse files
authored
fix(amazonq): fix cache watcher on token file to avoid logout after SSO connection migration (#7345)
## Problem The current cache file watcher is triggered on every `*.json` file change in the cache folder. When migrating old auth SSO connections to Flare, there is a race condition. If the old connection is migrated (meaning that the old token file deletes) after the file watcher is instantiated, the file watcher `delete` event will trigger and logout the user ## Solution The Flare auth cache file watcher will only listen to events for the Flare token .json file. This way, when we delete the old token .json file, there is no logout event triggered --- - 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 534e7a4 commit 0d6d35e

File tree

5 files changed

+50
-31
lines changed

5 files changed

+50
-31
lines changed

packages/amazonq/src/util/clearCache.ts

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -32,10 +32,9 @@ async function clearCache() {
3232
return
3333
}
3434

35-
// SSO cache persists on disk, this should indirectly delete it
36-
const conn = AuthUtil.instance.conn
37-
if (conn) {
38-
await AuthUtil.instance.auth.deleteConnection(conn)
35+
// SSO cache persists on disk, this should log out
36+
if (AuthUtil.instance.isConnected()) {
37+
await AuthUtil.instance.logout()
3938
}
4039

4140
await globals.globalState.clear()

packages/amazonq/test/unit/amazonq/apps/inline/recommendationService.test.ts

Lines changed: 14 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -9,14 +9,19 @@ import { Position, CancellationToken, InlineCompletionItem } from 'vscode'
99
import assert from 'assert'
1010
import { RecommendationService } from '../../../../../src/app/inline/recommendationService'
1111
import { SessionManager } from '../../../../../src/app/inline/sessionManager'
12-
import { createMockDocument } from 'aws-core-vscode/test'
12+
import { createMockDocument, createTestAuthUtil } from 'aws-core-vscode/test'
1313
import { LineTracker } from '../../../../../src/app/inline/stateTracker/lineTracker'
1414
import { InlineGeneratingMessage } from '../../../../../src/app/inline/inlineGeneratingMessage'
1515

1616
describe('RecommendationService', () => {
1717
let languageClient: LanguageClient
1818
let sendRequestStub: sinon.SinonStub
1919
let sandbox: sinon.SinonSandbox
20+
let sessionManager: SessionManager
21+
let lineTracker: LineTracker
22+
let activeStateController: InlineGeneratingMessage
23+
let service: RecommendationService
24+
2025
const mockDocument = createMockDocument()
2126
const mockPosition = { line: 0, character: 0 } as Position
2227
const mockContext = { triggerKind: 1, selectedCompletionInfo: undefined }
@@ -29,19 +34,22 @@ describe('RecommendationService', () => {
2934
insertText: 'ItemTwo',
3035
} as InlineCompletionItem
3136
const mockPartialResultToken = 'some-random-token'
32-
const sessionManager = new SessionManager()
33-
const lineTracker = new LineTracker()
34-
const activeStateController = new InlineGeneratingMessage(lineTracker)
35-
const service = new RecommendationService(sessionManager, activeStateController)
3637

37-
beforeEach(() => {
38+
beforeEach(async () => {
3839
sandbox = sinon.createSandbox()
3940

4041
sendRequestStub = sandbox.stub()
4142

4243
languageClient = {
4344
sendRequest: sendRequestStub,
4445
} as unknown as LanguageClient
46+
47+
await createTestAuthUtil()
48+
49+
sessionManager = new SessionManager()
50+
lineTracker = new LineTracker()
51+
activeStateController = new InlineGeneratingMessage(lineTracker)
52+
service = new RecommendationService(sessionManager, activeStateController)
4553
})
4654

4755
afterEach(() => {

packages/core/src/auth/auth2.ts

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -41,7 +41,8 @@ import { LanguageClient } from 'vscode-languageclient'
4141
import { getLogger } from '../shared/logger/logger'
4242
import { ToolkitError } from '../shared/errors'
4343
import { useDeviceFlow } from './sso/ssoAccessTokenProvider'
44-
import { getCacheFileWatcher } from './sso/cache'
44+
import { getCacheDir, getCacheFileWatcher, getFlareCacheFileName } from './sso/cache'
45+
import { VSCODE_EXTENSION_ID } from '../shared/extensions'
4546

4647
export const notificationTypes = {
4748
updateBearerToken: new RequestType<UpdateCredentialsParams, ResponseMessage, Error>(
@@ -77,7 +78,7 @@ export type TokenSource = IamIdentityCenterSsoTokenSource | AwsBuilderIdSsoToken
7778
* Handles auth requests to the Identity Server in the Amazon Q LSP.
7879
*/
7980
export class LanguageClientAuth {
80-
readonly #ssoCacheWatcher = getCacheFileWatcher()
81+
readonly #ssoCacheWatcher = getCacheFileWatcher(getCacheDir(), getFlareCacheFileName(VSCODE_EXTENSION_ID.amazonq))
8182

8283
constructor(
8384
private readonly client: LanguageClient,

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

Lines changed: 18 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -45,8 +45,8 @@ export function getCache(directory = getCacheDir()): SsoCache {
4545
}
4646
}
4747

48-
export function getCacheFileWatcher(directory = getCacheDir()) {
49-
const watcher = vscode.workspace.createFileSystemWatcher(new vscode.RelativePattern(directory, '*.json'))
48+
export function getCacheFileWatcher(directory = getCacheDir(), file = '*.json') {
49+
const watcher = vscode.workspace.createFileSystemWatcher(new vscode.RelativePattern(directory, file))
5050
globals.context.subscriptions.push(watcher)
5151
return watcher
5252
}
@@ -158,3 +158,19 @@ export function getRegistrationCacheFile(ssoCacheDir: string, key: RegistrationK
158158
const suffix = `${key.region}${key.scopes && key.scopes.length > 0 ? `-${hash(key.startUrl, key.scopes)}` : ''}`
159159
return path.join(ssoCacheDir, `aws-toolkit-vscode-client-id-${suffix}.json`)
160160
}
161+
162+
/**
163+
* Returns the cache file name that Flare identity server uses for SSO token and registration
164+
*
165+
* @param key - The key to use for the new registration cache file.
166+
* See https://github.com/aws/language-servers/blob/c10819ea2c25ce564c75fb43a6792f3c919b757a/server/aws-lsp-identity/src/sso/cache/fileSystemSsoCache.ts
167+
* @returns File name of the Flare cache file
168+
*/
169+
export function getFlareCacheFileName(key: string) {
170+
const hash = (str: string) => {
171+
const hasher = crypto.createHash('sha1')
172+
return hasher.update(str).digest('hex')
173+
}
174+
175+
return hash(key) + '.json'
176+
}

packages/core/src/codewhisperer/util/authUtil.ts

Lines changed: 12 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,6 @@ import * as localizedText from '../../shared/localizedText'
88
import * as nls from 'vscode-nls'
99
import { fs } from '../../shared/fs/fs'
1010
import * as path from 'path'
11-
import * as crypto from 'crypto'
1211
import { ToolkitError } from '../../shared/errors'
1312
import { AmazonQPromptSettings } from '../../shared/settings'
1413
import {
@@ -37,7 +36,7 @@ import { VSCODE_EXTENSION_ID } from '../../shared/extensions'
3736
import { RegionProfileManager } from '../region/regionProfileManager'
3837
import { AuthFormId } from '../../login/webview/vue/types'
3938
import { getEnvironmentSpecificMemento } from '../../shared/utilities/mementos'
40-
import { getCacheDir, getRegistrationCacheFile, getTokenCacheFile } from '../../auth/sso/cache'
39+
import { getCacheDir, getFlareCacheFileName, getRegistrationCacheFile, getTokenCacheFile } from '../../auth/sso/cache'
4140
import { notifySelectDeveloperProfile } from '../region/utils'
4241
import { once } from '../../shared/utilities/functionUtils'
4342

@@ -278,6 +277,7 @@ export class AuthUtil implements IAuthProvider {
278277
}
279278

280279
private async cacheChangedHandler(event: cacheChangedEvent) {
280+
getLogger().debug(`Auth: Cache change event received: ${event}`)
281281
if (event === 'delete') {
282282
await this.logout()
283283
} else if (event === 'create') {
@@ -427,25 +427,20 @@ export class AuthUtil implements IAuthProvider {
427427

428428
const cacheDir = getCacheDir()
429429

430-
const hash = (str: string) => {
431-
const hasher = crypto.createHash('sha1')
432-
return hasher.update(str).digest('hex')
433-
}
434-
const filePath = (str: string) => {
435-
return path.join(cacheDir, hash(str) + '.json')
436-
}
437-
438430
const fromRegistrationFile = getRegistrationCacheFile(cacheDir, registrationKey)
439-
const toRegistrationFile = filePath(
440-
JSON.stringify({
441-
region: toImport.ssoRegion,
442-
startUrl: toImport.startUrl,
443-
tool: clientName,
444-
})
431+
const toRegistrationFile = path.join(
432+
cacheDir,
433+
getFlareCacheFileName(
434+
JSON.stringify({
435+
region: toImport.ssoRegion,
436+
startUrl: toImport.startUrl,
437+
tool: clientName,
438+
})
439+
)
445440
)
446441

447442
const fromTokenFile = getTokenCacheFile(cacheDir, profileId)
448-
const toTokenFile = filePath(this.profileName)
443+
const toTokenFile = path.join(cacheDir, getFlareCacheFileName(this.profileName))
449444

450445
try {
451446
await fs.rename(fromRegistrationFile, toRegistrationFile)

0 commit comments

Comments
 (0)