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

Author: opieter-aws

packages/amazonq/src/lsp/client.ts

@@ -109,6 +109,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,
@@ -133,7 +134,7 @@ export async function startLanguageServer(
                     name: env.appName,
                     version: version,
                     extension: {
-                        name: 'AmazonQ-For-VSCode',
+                        name: clientName,
                         version: '0.0.1',
                     },
                     clientId: crypto.randomUUID(),
@@ -177,6 +178,12 @@ export async function startLanguageServer(
     await client.onReady()
     AuthUtil.create(new auth2.LanguageClientAuth(client, clientId, encryptionKey))
 
+    try {
+        await AuthUtil.instance.migrateSsoConnectionToLsp(clientName)
+    } catch (e) {
+        client.error(`Error while migration SSO connection to Amazon Q LSP: ${e}`)
+    }
+
     // Request handler for when the server wants to know about the clients auth connnection. Must be registered before the initial auth init call
     client.onRequest<ConnectionMetadata, Error>(auth2.notificationTypes.getConnectionMetadata.method, () => {
         return {

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 * as fs from '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'
 
 const localize = nls.loadMessageBundle()
 
@@ -333,4 +341,79 @@ 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
+
+        getLogger().info(`codewhisperer: checking for old SSO connections`)
+        if (profiles) {
+            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 {
+                    fs.renameSync(fromRegistrationFile, toRegistrationFile)
+                    fs.renameSync(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`)
+            }
+        }
+    }
 }