feat(core): Add new auth class using flare identity server for AmazonQ

Author: opieter-aws

packages/core/src/auth/auth2.ts

@@ -0,0 +1,330 @@
+/*!
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import * as vscode from 'vscode'
+import * as jose from 'jose'
+import {
+    GetSsoTokenParams,
+    getSsoTokenRequestType,
+    GetSsoTokenResult,
+    IamIdentityCenterSsoTokenSource,
+    InvalidateSsoTokenParams,
+    invalidateSsoTokenRequestType,
+    ProfileKind,
+    UpdateProfileParams,
+    updateProfileRequestType,
+    SsoTokenChangedParams,
+    ssoTokenChangedRequestType,
+    AwsBuilderIdSsoTokenSource,
+    UpdateCredentialsParams,
+    AwsErrorCodes,
+    SsoTokenSourceKind,
+    listProfilesRequestType,
+    ListProfilesResult,
+    UpdateProfileResult,
+    InvalidateSsoTokenResult,
+    AuthorizationFlowKind,
+    CancellationToken,
+    CancellationTokenSource,
+    bearerCredentialsDeleteNotificationType,
+    bearerCredentialsUpdateRequestType,
+    SsoTokenChangedKind,
+} from '@aws/language-server-runtimes/protocol'
+import { LanguageClient } from 'vscode-languageclient'
+import { getLogger } from '../shared/logger/logger'
+import { ToolkitError } from '../shared/errors'
+import { useDeviceFlow } from './sso/ssoAccessTokenProvider'
+
+export const AuthStates = {
+    NOT_CONNECTED: 'notConnected',
+    CONNECTED: 'connected',
+    EXPIRED: 'expired',
+} as const
+export type AuthState = (typeof AuthStates)[keyof typeof AuthStates]
+
+export type AuthStateEvent = { id: string; state: AuthState | 'refreshed' }
+
+export const LoginTypes = {
+    SSO: 'sso',
+    IAM: 'iam',
+} as const
+export type LoginType = (typeof LoginTypes)[keyof typeof LoginTypes]
+
+interface BaseLogin {
+    readonly loginType: LoginType
+}
+
+export type Login = SsoLogin // TODO: add IamLogin type when supported
+
+export type TokenSource = IamIdentityCenterSsoTokenSource | AwsBuilderIdSsoTokenSource
+
+/**
+ * Handles auth requests to the Identity Server in the Amazon Q LSP.
+ */
+export class LanguageClientAuth {
+    constructor(
+        private readonly client: LanguageClient,
+        private readonly clientName: string,
+        public readonly encryptionKey: Buffer
+    ) {}
+
+    getSsoToken(
+        tokenSource: TokenSource,
+        login: boolean = false,
+        cancellationToken?: CancellationToken
+    ): Promise<GetSsoTokenResult> {
+        return this.client.sendRequest(
+            getSsoTokenRequestType.method,
+            {
+                clientName: this.clientName,
+                source: tokenSource,
+                options: {
+                    loginOnInvalidToken: login,
+                    authorizationFlow: useDeviceFlow() ? AuthorizationFlowKind.DeviceCode : AuthorizationFlowKind.Pkce,
+                },
+            } satisfies GetSsoTokenParams,
+            cancellationToken
+        )
+    }
+
+    updateProfile(
+        profileName: string,
+        startUrl: string,
+        region: string,
+        scopes: string[]
+    ): Promise<UpdateProfileResult> {
+        return this.client.sendRequest(updateProfileRequestType.method, {
+            profile: {
+                kinds: [ProfileKind.SsoTokenProfile],
+                name: profileName,
+                settings: {
+                    region,
+                    sso_session: profileName,
+                },
+            },
+            ssoSession: {
+                name: profileName,
+                settings: {
+                    sso_region: region,
+                    sso_start_url: startUrl,
+                    sso_registration_scopes: scopes,
+                },
+            },
+        } satisfies UpdateProfileParams)
+    }
+
+    listProfiles() {
+        return this.client.sendRequest(listProfilesRequestType.method, {}) as Promise<ListProfilesResult>
+    }
+
+    /**
+     * Returns a profile by name along with its linked sso_session.
+     * Does not currently exist as an API in the Identity Service.
+     */
+    async getProfile(profileName: string) {
+        const response = await this.listProfiles()
+        const profile = response.profiles.find((profile) => profile.name === profileName)
+        const ssoSession = profile?.settings?.sso_session
+            ? response.ssoSessions.find((session) => session.name === profile!.settings!.sso_session)
+            : undefined
+
+        return { profile, ssoSession }
+    }
+
+    updateBearerToken(request: UpdateCredentialsParams) {
+        return this.client.sendRequest(bearerCredentialsUpdateRequestType.method, request)
+    }
+
+    deleteBearerToken() {
+        return this.client.sendNotification(bearerCredentialsDeleteNotificationType.method)
+    }
+
+    invalidateSsoToken(tokenId: string) {
+        return this.client.sendRequest(invalidateSsoTokenRequestType.method, {
+            ssoTokenId: tokenId,
+        } satisfies InvalidateSsoTokenParams) as Promise<InvalidateSsoTokenResult>
+    }
+
+    registerSsoTokenChangedHandler(ssoTokenChangedHandler: (params: SsoTokenChangedParams) => any) {
+        this.client.onNotification(ssoTokenChangedRequestType.method, ssoTokenChangedHandler)
+    }
+}
+
+/**
+ * Manages an SSO connection.
+ */
+export class SsoLogin implements BaseLogin {
+    readonly loginType = LoginTypes.SSO
+    private readonly eventEmitter = new vscode.EventEmitter<AuthStateEvent>()
+
+    // Cached information from the identity server for easy reference
+    private ssoTokenId: string | undefined
+    private connectionState: AuthState = AuthStates.NOT_CONNECTED
+    private _data: { startUrl: string; region: string } | undefined
+
+    private cancellationToken: CancellationTokenSource | undefined
+
+    constructor(
+        public readonly profileName: string,
+        private readonly lspAuth: LanguageClientAuth
+    ) {
+        lspAuth.registerSsoTokenChangedHandler((params: SsoTokenChangedParams) => this.ssoTokenChangedHandler(params))
+    }
+
+    get data() {
+        return this._data
+    }
+
+    async login(opts: { startUrl: string; region: string; scopes: string[] }) {
+        await this.updateProfile(opts)
+        return this._getSsoToken(true)
+    }
+
+    async reauthenticate() {
+        if (this.connectionState === AuthStates.NOT_CONNECTED) {
+            throw new ToolkitError('Cannot reauthenticate when not connected.')
+        }
+        return this._getSsoToken(true)
+    }
+
+    async logout() {
+        if (this.ssoTokenId) {
+            await this.lspAuth.invalidateSsoToken(this.ssoTokenId)
+        }
+        this.updateConnectionState(AuthStates.NOT_CONNECTED)
+        this._data = undefined
+        // TODO: DeleteProfile api in Identity Service (this doesn't exist yet)
+    }
+
+    async updateProfile(opts: { startUrl: string; region: string; scopes: string[] }) {
+        await this.lspAuth.updateProfile(this.profileName, opts.startUrl, opts.region, opts.scopes)
+        this._data = {
+            startUrl: opts.startUrl,
+            region: opts.region,
+        }
+    }
+
+    /**
+     * Restore the connection state and connection details to memory, if they exist.
+     */
+    async restore() {
+        const sessionData = await this.lspAuth.getProfile(this.profileName)
+        const ssoSession = sessionData?.ssoSession?.settings
+        if (ssoSession?.sso_region && ssoSession?.sso_start_url) {
+            this._data = {
+                startUrl: ssoSession.sso_start_url,
+                region: ssoSession.sso_region,
+            }
+        }
+
+        try {
+            await this._getSsoToken(false)
+        } catch (err) {
+            getLogger().error('Restoring connection failed: %s', err)
+        }
+    }
+
+    /**
+     * Cancels running active login flows.
+     */
+    cancelLogin() {
+        this.cancellationToken?.cancel()
+        this.cancellationToken?.dispose()
+        this.cancellationToken = undefined
+    }
+
+    /**
+     * Returns a decrypted access token and a payload to send to the `updateCredentials` API provided by
+     * the Amazon Q LSP.
+     */
+    async getToken() {
+        const response = await this._getSsoToken(false)
+        const decryptedKey = await jose.compactDecrypt(response.ssoToken.accessToken, this.lspAuth.encryptionKey)
+        return {
+            token: decryptedKey.plaintext.toString().replaceAll('"', ''),
+            updateCredentialsParams: response.updateCredentialsParams,
+        }
+    }
+
+    /**
+     * Returns the response from `getToken` LSP API and sets the connection state based on the errors/result
+     * of the call.
+     */
+    private async _getSsoToken(login: boolean) {
+        let response: GetSsoTokenResult
+        this.cancellationToken = new CancellationTokenSource()
+
+        try {
+            response = await this.lspAuth.getSsoToken(
+                {
+                    /**
+                     * Note that we do not use SsoTokenSourceKind.AwsBuilderId here.
+                     * This is because it does not leave any state behind on disk, so
+                     * we cannot infer that a builder ID connection exists via the
+                     * Identity Server alone.
+                     */
+                    kind: SsoTokenSourceKind.IamIdentityCenter,
+                    profileName: this.profileName,
+                } satisfies IamIdentityCenterSsoTokenSource,
+                login,
+                this.cancellationToken.token
+            )
+        } catch (err: any) {
+            switch (err.data?.awsErrorCode) {
+                case AwsErrorCodes.E_CANCELLED:
+                case AwsErrorCodes.E_SSO_SESSION_NOT_FOUND:
+                case AwsErrorCodes.E_PROFILE_NOT_FOUND:
+                case AwsErrorCodes.E_INVALID_SSO_TOKEN:
+                    this.updateConnectionState(AuthStates.NOT_CONNECTED)
+                    break
+                case AwsErrorCodes.E_CANNOT_REFRESH_SSO_TOKEN:
+                    this.updateConnectionState(AuthStates.EXPIRED)
+                    break
+                // TODO: implement when identity server emits E_NETWORK_ERROR, E_FILESYSTEM_ERROR
+                // case AwsErrorCodes.E_NETWORK_ERROR:
+                // case AwsErrorCodes.E_FILESYSTEM_ERROR:
+                //     // do stuff, probably nothing at all
+                //     break
+                default:
+                    getLogger().error('SsoLogin: unknown error when requesting token: %s', err)
+                    break
+            }
+            throw err
+        } finally {
+            this.cancellationToken?.dispose()
+            this.cancellationToken = undefined
+        }
+
+        this.ssoTokenId = response.ssoToken.id
+        this.updateConnectionState(AuthStates.CONNECTED)
+        return response
+    }
+
+    getConnectionState() {
+        return this.connectionState
+    }
+
+    onDidChangeConnectionState(handler: (e: AuthStateEvent) => any) {
+        return this.eventEmitter.event(handler)
+    }
+
+    private updateConnectionState(state: AuthState) {
+        if (this.connectionState !== state) {
+            this.eventEmitter.fire({ id: this.profileName, state })
+        }
+        this.connectionState = state
+    }
+
+    private ssoTokenChangedHandler(params: SsoTokenChangedParams) {
+        if (params.ssoTokenId === this.ssoTokenId) {
+            if (params.kind === SsoTokenChangedKind.Expired) {
+                this.updateConnectionState(AuthStates.EXPIRED)
+                return
+            } else if (params.kind === SsoTokenChangedKind.Refreshed) {
+                this.eventEmitter.fire({ id: this.profileName, state: 'refreshed' })
+            }
+        }
+    }
+}

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

@@ -289,17 +289,7 @@ export abstract class SsoAccessTokenProvider {
         profile: Pick<SsoProfile, 'startUrl' | 'region' | 'scopes' | 'identifier'>,
         cache = getCache(),
         oidc: OidcClient = OidcClient.create(profile.region),
-        reAuthState?: ReAuthState,
-        useDeviceFlow: () => boolean = () => {
-            /**
-             * Device code flow is neccessary when:
-             * 1. We are in a workspace connected through ssh (codecatalyst, etc)
-             * 2. We are connected to a remote backend through the web browser (code server, openshift dev spaces)
-             *
-             * Since we are unable to serve the final authorization page
-             */
-            return getExtRuntimeContext().extensionHost === 'remote'
-        }
+        reAuthState?: ReAuthState
     ) {
         if (DevSettings.instance.get('webAuth', false) && getExtRuntimeContext().extensionHost === 'webworker') {
             return new WebAuthorization(profile, cache, oidc, reAuthState)
@@ -400,6 +390,17 @@ function getSessionDuration(id: string) {
     return creationDate !== undefined ? globals.clock.Date.now() - creationDate : undefined
 }
 
+export function useDeviceFlow(): boolean {
+    /**
+     * Device code flow is neccessary when:
+     * 1. We are in a workspace connected through ssh (codecatalyst, etc)
+     * 2. We are connected to a remote backend through the web browser (code server, openshift dev spaces)
+     *
+     * Since we are unable to serve the final authorization page
+     */
+    return getExtRuntimeContext().extensionHost === 'remote'
+}
+
 /**
  *  SSO "device code" flow (RFC: https://tools.ietf.org/html/rfc8628)
  *    1. Get a client id (SSO-OIDC identifier, formatted per RFC6749).