feat(amazonq): Migrate existing SSO connections to the LSP identity server (#7298)

## Problem
Existing SSO connections on old auth, are not migrated automatically to
the LSP Identity Server. This means that users would be signed out upon
migration.

## Solution
Add a migration functionality that checks for existing SSO connections,
and migrate if necessary:
* Check inside the `client.ts`, before we attempt to restore any
connections in the new auth
* Add `migrateSsoConnectionToLsp` method to `AuthUtil`, which:
* Fetches environment `memento`, and looks for the `auth.profiles` key
* Checks if there is an `sso` profile with `amazonQScopes`. If yes:
migrate, if not: nothing to migrate
   * [migration]: Update the SSO profile with Flare
* [migration]: Construct the existing (`from`) filenames and the Flare
(`to`) filenames for the registration cache file and the SSO token cache
file
   * [migration]: rename the files
* [migration]: set the `auth.profiles` key to undefined, so migration is
skipped next time


## Testing
* Covered with unit tests that confirms the actual file system
operations
* Manual E2E testing, verifying that there is no sign-out happening with
this setup
* We will verify the mechanism through bug bashing

---

- 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.

Author: opieter-aws

packages/amazonq/src/lsp/client.ts

@@ -112,6 +112,7 @@ export async function startLanguageServer(
     await validateNodeExe(executable, resourcePaths.lsp, argv, logger)
 
     // Options to control the language client
+    const clientName = 'AmazonQ-For-VSCode'
     const clientOptions: LanguageClientOptions = {
         // Register the server for json documents
         documentSelector,
@@ -136,7 +137,7 @@ export async function startLanguageServer(
                     name: env.appName,
                     version: version,
                     extension: {
-                        name: 'AmazonQ-For-VSCode',
+                        name: clientName,
                         version: extensionVersion,
                     },
                     clientId: getClientId(globals.globalState),
@@ -193,6 +194,15 @@ export async function startLanguageServer(
     async function initializeAuth(client: LanguageClient) {
         AuthUtil.create(new auth2.LanguageClientAuth(client, clientId, encryptionKey))
 
+        // Migrate SSO connections from old Auth to the LSP identity server
+        // This function only migrates connections once
+        // This call can be removed once all/most users have updated to the latest AmazonQ version
+        try {
+            await AuthUtil.instance.migrateSsoConnectionToLsp(clientName)
+        } catch (e) {
+            getLogger().error(`Error while migration SSO connection to Amazon Q LSP: ${e}`)
+        }
+
         /** All must be setup before {@link AuthUtil.restore} otherwise they may not trigger when expected */
         AuthUtil.instance.regionProfileManager.onDidChangeRegionProfile(async () => {
             void pushConfigUpdate(client, {

packages/amazonq/test/unit/codewhisperer/util/authUtil.test.ts

@@ -5,10 +5,12 @@
 
 import assert from 'assert'
 import * as sinon from 'sinon'
-import { AuthUtil } from 'aws-core-vscode/codewhisperer'
-import { createTestAuthUtil } from 'aws-core-vscode/test'
-import { constants } from 'aws-core-vscode/auth'
+import * as path from 'path'
+import { AuthUtil, amazonQScopes } from 'aws-core-vscode/codewhisperer'
+import { createTestAuthUtil, TestFolder } from 'aws-core-vscode/test'
+import { constants, cache } from 'aws-core-vscode/auth'
 import { auth2 } from 'aws-core-vscode/auth'
+import { mementoUtils, fs } from 'aws-core-vscode/shared'
 
 describe('AuthUtil', async function () {
     let auth: any
@@ -193,4 +195,145 @@ describe('AuthUtil', async function () {
             assert.ok(clearCacheSpy.called)
         })
     })
+
+    describe('migrateSsoConnectionToLsp', function () {
+        let memento: any
+        let cacheDir: string
+        let fromRegistrationFile: string
+        let fromTokenFile: string
+
+        const validProfile = {
+            type: 'sso',
+            startUrl: 'https://test2.com',
+            ssoRegion: 'us-east-1',
+            scopes: amazonQScopes,
+            metadata: {
+                connectionState: 'valid',
+            },
+        }
+
+        beforeEach(async function () {
+            memento = {
+                get: sinon.stub(),
+                update: sinon.stub().resolves(),
+            }
+            cacheDir = (await TestFolder.create()).path
+
+            sinon.stub(mementoUtils, 'getEnvironmentSpecificMemento').returns(memento)
+            sinon.stub(cache, 'getCacheDir').returns(cacheDir)
+
+            fromTokenFile = cache.getTokenCacheFile(cacheDir, 'profile1')
+            const registrationKey = {
+                startUrl: validProfile.startUrl,
+                region: validProfile.ssoRegion,
+                scopes: amazonQScopes,
+            }
+            fromRegistrationFile = cache.getRegistrationCacheFile(cacheDir, registrationKey)
+
+            const registrationData = { test: 'registration' }
+            const tokenData = { test: 'token' }
+
+            await fs.writeFile(fromRegistrationFile, JSON.stringify(registrationData))
+            await fs.writeFile(fromTokenFile, JSON.stringify(tokenData))
+        })
+
+        afterEach(async function () {
+            sinon.restore()
+        })
+
+        it('migrates valid SSO connection', async function () {
+            memento.get.returns({ profile1: validProfile })
+
+            const updateProfileStub = sinon.stub((auth as any).session, 'updateProfile').resolves()
+
+            await auth.migrateSsoConnectionToLsp('test-client')
+
+            assert.ok(updateProfileStub.calledOnce)
+            assert.ok(memento.update.calledWith('auth.profiles', undefined))
+
+            const files = await fs.readdir(cacheDir)
+            assert.strictEqual(files.length, 2) // Should have both the token and registration file
+
+            // Verify file contents were preserved
+            const newFiles = files.map((f) => path.join(cacheDir, f[0]))
+            for (const file of newFiles) {
+                const content = await fs.readFileText(file)
+                const parsed = JSON.parse(content)
+                assert.ok(parsed.test === 'registration' || parsed.test === 'token')
+            }
+        })
+
+        it('does not migrate if no matching SSO profile exists', async function () {
+            const mockProfiles = {
+                'test-profile': {
+                    type: 'iam',
+                    startUrl: 'https://test.com',
+                    ssoRegion: 'us-east-1',
+                },
+            }
+            memento.get.returns(mockProfiles)
+
+            await auth.migrateSsoConnectionToLsp('test-client')
+
+            // Assert that the file names have not updated
+            const files = await fs.readdir(cacheDir)
+            assert.ok(files.length === 2)
+            assert.ok(await fs.exists(fromRegistrationFile))
+            assert.ok(await fs.exists(fromTokenFile))
+            assert.ok(!memento.update.called)
+        })
+
+        it('migrates only profile with matching scopes', async function () {
+            const mockProfiles = {
+                profile1: validProfile,
+                profile2: {
+                    type: 'sso',
+                    startUrl: 'https://test.com',
+                    ssoRegion: 'us-east-1',
+                    scopes: ['different:scope'],
+                    metadata: {
+                        connectionState: 'valid',
+                    },
+                },
+            }
+            memento.get.returns(mockProfiles)
+
+            const updateProfileStub = sinon.stub((auth as any).session, 'updateProfile').resolves()
+
+            await auth.migrateSsoConnectionToLsp('test-client')
+
+            assert.ok(updateProfileStub.calledOnce)
+            assert.ok(memento.update.calledWith('auth.profiles', undefined))
+            assert.deepStrictEqual(updateProfileStub.firstCall.args[0], {
+                startUrl: validProfile.startUrl,
+                region: validProfile.ssoRegion,
+                scopes: validProfile.scopes,
+            })
+        })
+
+        it('uses valid connection state when multiple profiles exist', async function () {
+            const mockProfiles = {
+                profile2: {
+                    ...validProfile,
+                    metadata: {
+                        connectionState: 'invalid',
+                    },
+                },
+                profile1: validProfile,
+            }
+            memento.get.returns(mockProfiles)
+
+            const updateProfileStub = sinon.stub((auth as any).session, 'updateProfile').resolves()
+
+            await auth.migrateSsoConnectionToLsp('test-client')
+
+            assert.ok(
+                updateProfileStub.calledWith({
+                    startUrl: validProfile.startUrl,
+                    region: validProfile.ssoRegion,
+                    scopes: validProfile.scopes,
+                })
+            )
+        })
+    })
 })

packages/core/src/auth/index.ts

@@ -23,6 +23,7 @@ export { Auth } from './auth'
 export { CredentialsStore } from './credentials/store'
 export { LoginManager } from './deprecated/loginManager'
 export * as constants from './sso/constants'
+export * as cache from './sso/cache'
 export * as authUtils from './utils'
 export * as auth2 from './auth2'
 export * as SsoAccessTokenProvider from './sso/ssoAccessTokenProvider'

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

@@ -126,7 +126,7 @@ export function getTokenCache(directory = getCacheDir()): KeyedCache<SsoAccess>
     return mapCache(cache, read, write)
 }
 
-function getTokenCacheFile(ssoCacheDir: string, key: string): string {
+export function getTokenCacheFile(ssoCacheDir: string, key: string): string {
     const encoded = encodeURI(key)
     // Per the spec: 'SSO Login Token Flow' the access token must be
     // cached as the SHA1 hash of the bytes of the UTF-8 encoded
@@ -145,7 +145,7 @@ function getTokenCacheFile(ssoCacheDir: string, key: string): string {
     return path.join(ssoCacheDir, `${hashedKey}.json`)
 }
 
-function getRegistrationCacheFile(ssoCacheDir: string, key: RegistrationKey): string {
+export function getRegistrationCacheFile(ssoCacheDir: string, key: RegistrationKey): string {
     const hash = (startUrl: string, scopes: string[]) => {
         const shasum = crypto.createHash('sha256')
         shasum.update(startUrl)

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

@@ -6,6 +6,9 @@
 import * as vscode from 'vscode'
 import * as localizedText from '../../shared/localizedText'
 import * as nls from 'vscode-nls'
+import { fs } from '../../shared/fs/fs'
+import * as path from 'path'
+import * as crypto from 'crypto'
 import { ToolkitError } from '../../shared/errors'
 import { AmazonQPromptSettings } from '../../shared/settings'
 import {
@@ -16,6 +19,9 @@ import {
     TelemetryMetadata,
     scopesSsoAccountAccess,
     hasScopes,
+    SsoProfile,
+    StoredProfile,
+    hasExactScopes,
 } from '../../auth/connection'
 import { getLogger } from '../../shared/logger/logger'
 import { Commands } from '../../shared/vscode/commands2'
@@ -30,6 +36,8 @@ import { builderIdStartUrl, internalStartUrl } from '../../auth/sso/constants'
 import { VSCODE_EXTENSION_ID } from '../../shared/extensions'
 import { RegionProfileManager } from '../region/regionProfileManager'
 import { AuthFormId } from '../../login/webview/vue/types'
+import { getEnvironmentSpecificMemento } from '../../shared/utilities/mementos'
+import { getCacheDir, getRegistrationCacheFile, getTokenCacheFile } from '../../auth/sso/cache'
 import { notifySelectDeveloperProfile } from '../region/utils'
 import { once } from '../../shared/utilities/functionUtils'
 
@@ -367,4 +375,81 @@ export class AuthUtil implements IAuthProvider {
 
         return authIds
     }
+
+    /**
+     * Migrates existing SSO connections to the LSP identity server by updating the cache files
+     *
+     * @param clientName - The client name to use for the new registration cache file
+     * @returns A Promise that resolves when the migration is complete
+     * @throws Error if file operations fail during migration
+     */
+    async migrateSsoConnectionToLsp(clientName: string) {
+        const memento = getEnvironmentSpecificMemento()
+        const key = 'auth.profiles'
+        const profiles: { readonly [id: string]: StoredProfile } | undefined = memento.get(key)
+
+        let toImport: SsoProfile | undefined
+        let profileId: string | undefined
+
+        if (!profiles) {
+            return
+        } else {
+            getLogger().info(`codewhisperer: checking for old SSO connections`)
+            for (const [id, p] of Object.entries(profiles)) {
+                if (p.type === 'sso' && hasExactScopes(p.scopes ?? [], amazonQScopes)) {
+                    toImport = p
+                    profileId = id
+                    if (p.metadata.connectionState === 'valid') {
+                        break
+                    }
+                }
+            }
+
+            if (toImport && profileId) {
+                getLogger().info(`codewhisperer: migrating SSO connection to LSP identity server...`)
+
+                const registrationKey = {
+                    startUrl: toImport.startUrl,
+                    region: toImport.ssoRegion,
+                    scopes: amazonQScopes,
+                }
+
+                await this.session.updateProfile(registrationKey)
+
+                const cacheDir = getCacheDir()
+
+                const hash = (str: string) => {
+                    const hasher = crypto.createHash('sha1')
+                    return hasher.update(str).digest('hex')
+                }
+                const filePath = (str: string) => {
+                    return path.join(cacheDir, hash(str) + '.json')
+                }
+
+                const fromRegistrationFile = getRegistrationCacheFile(cacheDir, registrationKey)
+                const toRegistrationFile = filePath(
+                    JSON.stringify({
+                        region: toImport.ssoRegion,
+                        startUrl: toImport.startUrl,
+                        tool: clientName,
+                    })
+                )
+
+                const fromTokenFile = getTokenCacheFile(cacheDir, profileId)
+                const toTokenFile = filePath(this.profileName)
+
+                try {
+                    await fs.rename(fromRegistrationFile, toRegistrationFile)
+                    await fs.rename(fromTokenFile, toTokenFile)
+                    getLogger().debug('Successfully renamed registration and token files')
+                } catch (err) {
+                    getLogger().error(`Failed to rename files during migration: ${err}`)
+                    throw err
+                }
+
+                await memento.update(key, undefined)
+                getLogger().info(`codewhisperer: successfully migrated SSO connection to LSP identity server`)
+            }
+        }
+    }
 }

packages/core/src/shared/index.ts

@@ -74,3 +74,4 @@ export * as BaseLspInstaller from './lsp/baseLspInstaller'
 export * as collectionUtil from './utilities/collectionUtils'
 export * from './datetime'
 export * from './performance/marks'
+export * as mementoUtils from './utilities/mementos'