fix(auth): automatically refresh credentials #2618

Problem:
Credentials are not refreshed automatically, mostly because of caching but also
because not every source of credentials is "refreshable".

Solution:
Hide all the auth logic behind a simple interface then implement a provider that
is compatible with the SDK v2.

This implementation may prompt the user on expiration.

Author: JadenSimon

.changes/next-release/Bug Fix-7b4ceaca-4a12-48ec-97e8-8e8edba2159f.json

@@ -0,0 +1,4 @@
+{
+	"type": "Bug Fix",
+	"description": "Credential profiles that do not require user-input now correctly refresh when expired"
+}

src/credentials/credentialsStore.ts

@@ -61,11 +61,11 @@ export class CredentialsStore {
         let credentials = await this.getCredentials(credentialsId)
 
         if (!credentials) {
-            credentials = await this.setCredentials(credentialsId, credentialsProvider)
+            credentials = await this.consumeProvider(credentialsId, credentialsProvider)
         } else if (credentialsProvider.getHashCode() !== credentials.credentialsHashCode) {
             getLogger().verbose(`Using updated credentials: ${asString(credentialsId)}`)
             this.invalidateCredentials(credentialsId)
-            credentials = await this.setCredentials(credentialsId, credentialsProvider)
+            credentials = await this.consumeProvider(credentialsId, credentialsProvider)
         }
 
         return credentials
@@ -78,7 +78,14 @@ export class CredentialsStore {
         delete this.credentialsCache[asString(credentialsId)]
     }
 
-    private async setCredentials(
+    public async setCredentials(credentials: AWS.Credentials, provider: CredentialsProvider): Promise<void> {
+        this.credentialsCache[asString(provider.getCredentialsId())] = {
+            credentials,
+            credentialsHashCode: provider.getHashCode(),
+        }
+    }
+
+    private async consumeProvider(
         credentialsId: CredentialsId,
         credentialsProvider: CredentialsProvider
     ): Promise<CachedCredentials> {

src/credentials/loginManager.ts

@@ -3,13 +3,24 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import * as vscode from 'vscode'
 import globals from '../shared/extensionGlobals'
+
+import * as nls from 'vscode-nls'
+const localize = nls.loadMessageBundle()
+
+import * as vscode from 'vscode'
 import { CancellationError } from '../shared/utilities/timeoutUtils'
 import { AwsContext } from '../shared/awsContext'
 import { getAccountId } from '../shared/credentials/accountId'
 import { getLogger } from '../shared/logger'
-import { recordAwsValidateCredentials, recordVscodeActiveRegions, Result } from '../shared/telemetry/telemetry'
+import {
+    CredentialSourceId,
+    CredentialType,
+    recordAwsRefreshCredentials,
+    recordAwsValidateCredentials,
+    recordVscodeActiveRegions,
+    Result,
+} from '../shared/telemetry/telemetry'
 import { CredentialsStore } from './credentialsStore'
 import { CredentialsSettings, notifyUserInvalidCredentials } from './credentialsUtilities'
 import {
@@ -24,9 +35,9 @@ import { getIdeProperties, isCloud9 } from '../shared/extensionUtilities'
 import { SharedCredentialsProvider } from './providers/sharedCredentialsProvider'
 import { showViewLogsMessage } from '../shared/utilities/messages'
 import { isAutomation } from '../shared/vscode/env'
-
-import * as nls from 'vscode-nls'
-const localize = nls.loadMessageBundle()
+import { Credentials } from '@aws-sdk/types'
+import { ToolkitError } from '../shared/toolkitError'
+import * as localizedText from '../shared/localizedText'
 
 export class LoginManager {
     private readonly defaultCredentialsRegion = 'us-east-1'
@@ -51,27 +62,25 @@ export class LoginManager {
         let telemetryResult: Result = 'Failed'
 
         try {
-            provider = await CredentialsProviderManager.getInstance().getCredentialsProvider(args.providerId)
-            if (!provider) {
-                throw new Error(`Could not find Credentials Provider for ${asString(args.providerId)}`)
-            }
+            provider = await getProvider(args.providerId)
 
-            const storedCredentials = await this.store.upsertCredentials(args.providerId, provider)
-            if (!storedCredentials) {
+            const credentials = (await this.store.upsertCredentials(args.providerId, provider))?.credentials
+            if (!credentials) {
                 throw new Error(`No credentials found for id ${asString(args.providerId)}`)
             }
 
             const credentialsRegion = provider.getDefaultRegion() ?? this.defaultCredentialsRegion
-            const accountId = await getAccountId(storedCredentials.credentials, credentialsRegion)
+            const accountId = await getAccountId(credentials, credentialsRegion)
             if (!accountId) {
                 throw new Error('Could not determine Account Id for credentials')
             }
             recordVscodeActiveRegions({ value: (await this.awsContext.getExplorerRegions()).length })
 
+            this.awsContext.credentialsShim = createCredentialsShim(this.store, args.providerId, credentials)
             await this.awsContext.setCredentials({
-                credentials: storedCredentials.credentials,
-                credentialsId: asString(args.providerId),
+                credentials,
                 accountId: accountId,
+                credentialsId: asString(args.providerId),
                 defaultRegion: provider.getDefaultRegion(),
             })
 
@@ -91,6 +100,7 @@ export class LoginManager {
 
             await this.logout()
             this.store.invalidateCredentials(args.providerId)
+            this.awsContext.credentialsShim = undefined
             return false
         } finally {
             const credType = provider?.getTelemetryType()
@@ -233,3 +243,119 @@ function tryMakeCredentialsProviderId(credentials: string): CredentialsId | unde
         return undefined
     }
 }
+
+async function getProvider(id: CredentialsId): Promise<CredentialsProvider> {
+    const provider = await CredentialsProviderManager.getInstance().getCredentialsProvider(id)
+    if (!provider) {
+        throw new Error(`Could not find Credentials Provider for ${asString(id)}`)
+    }
+    return provider
+}
+
+/**
+ * The Toolkit implementation has a good amount of custom logic (SSO, source profiles, etc.)
+ * that was written to fill feature gaps in both AWS SDK V2 and V3. This code works imperfectly with
+ * pre-existing SDK refresh logic, leading to users experiencing issues with credentials expiring and
+ * forcing them to re-select a profile to refresh.
+ *
+ * So, this interface sits between everything else to mediate the refreshes. Adding a thin interface
+ * is preferred over bulking up existing ones since it allows for clean(-ish) refactors against the
+ * credentials subsystem. The pure data structure shape that existed previously was better, but it
+ * just wouldn't be able to support refreshes on its own.
+ */
+export interface CredentialsShim {
+    /**
+     * Fetches credentials, attempting a refresh if needed.
+     */
+    get: () => Promise<Credentials>
+
+    /**
+     * Removes the stored credentials and performs a refresh, allowing for prompts.
+     *
+     * Calling this function while a refresh is still pending returns the already pending promise.
+     */
+    refresh: () => Promise<Credentials>
+}
+
+/**
+ * Collapses a single {@link CredentialsProvider} (referenced by id) into something a bit simpler.
+ *
+ * We don't pass in a provider directly since {@link CredentialsProviderManager} is the true
+ * source of credential state, at least as far as the Toolkit is concerned.
+ */
+function createCredentialsShim(
+    store: CredentialsStore,
+    providerId: CredentialsId,
+    creds: Credentials
+): CredentialsShim {
+    interface State {
+        credentials: Promise<Credentials>
+        pendingRefresh: Promise<Credentials>
+    }
+
+    const state: Partial<State> = { credentials: Promise.resolve(creds) }
+
+    async function refresh(): Promise<Credentials> {
+        let result: Result = 'Failed'
+        let credentialType: CredentialType | undefined
+        let credentialSourceId: CredentialSourceId | undefined
+
+        try {
+            getLogger().debug(`credentials: refreshing provider: ${asString(providerId)}`)
+
+            const provider = await getProvider(providerId)
+            const formatProviderId = () => asString(provider.getCredentialsId())
+
+            credentialType = provider.getTelemetryType()
+            credentialSourceId = credentialsProviderToTelemetryType(provider.getProviderType())
+
+            if (!provider.canAutoConnect()) {
+                const message = localize('aws.credentials.expired', 'Credentials are expired or invalid, login again?')
+                const resp = await vscode.window.showInformationMessage(message, localizedText.yes, localizedText.no)
+
+                if (resp === localizedText.no) {
+                    throw new ToolkitError('User cancelled login', { cancelled: true })
+                }
+            }
+
+            const credentials = await provider.getCredentials()
+            store.setCredentials(credentials, provider)
+            getLogger().debug(`credentials: refresh succeeded for: ${formatProviderId()}`)
+            result = 'Succeeded'
+
+            return credentials
+        } catch (error) {
+            if (error instanceof ToolkitError && error.cancelled) {
+                result = 'Cancelled'
+            } else {
+                showViewLogsMessage(`Failed to refresh credentials: ${(error as any)?.message}`)
+            }
+
+            state.credentials = undefined
+            store.invalidateCredentials(providerId)
+            globals.awsContext.credentialsShim = undefined
+            globals.awsContext.setCredentials(undefined, true)
+
+            throw error
+        } finally {
+            recordAwsRefreshCredentials({
+                result,
+                passive: true,
+                credentialType,
+                credentialSourceId,
+            })
+        }
+    }
+
+    const shim = {
+        get: () => (state.credentials ??= shim.refresh()),
+        refresh: () => {
+            const clear = () => (state.pendingRefresh = undefined)
+            state.credentials = state.pendingRefresh ??= refresh().finally(clear)
+
+            return state.credentials
+        },
+    }
+
+    return shim
+}

src/credentials/providers/ssoCredentialProvider.ts

@@ -5,7 +5,7 @@
 
 import * as vscode from 'vscode'
 import { Credentials } from '@aws-sdk/types'
-import { SSO } from '@aws-sdk/client-sso'
+import { SSO, UnauthorizedException } from '@aws-sdk/client-sso'
 import { getLogger } from '../../shared/logger'
 import { SsoAccessTokenProvider } from '../sso/ssoAccessTokenProvider'
 import * as nls from 'vscode-nls'
@@ -23,19 +23,24 @@ export class SsoCredentialProvider {
     public async refreshCredentials(): Promise<Credentials> {
         try {
             const accessToken = await this.ssoAccessTokenProvider.accessToken()
-            const roleCredentials = await this.ssoClient.getRoleCredentials({
-                accountId: this.ssoAccount,
-                roleName: this.ssoRole,
-                accessToken: accessToken.accessToken,
-            })
+            const roleCredentials = await this.ssoClient
+                .getRoleCredentials({
+                    accountId: this.ssoAccount,
+                    roleName: this.ssoRole,
+                    accessToken: accessToken.accessToken,
+                })
+                .then(resp => resp.roleCredentials)
+
+            const expiration = roleCredentials?.expiration ? new Date(roleCredentials.expiration) : undefined
 
             return {
-                accessKeyId: roleCredentials.roleCredentials!.accessKeyId!,
-                secretAccessKey: roleCredentials.roleCredentials!.secretAccessKey!,
-                sessionToken: roleCredentials.roleCredentials?.sessionToken,
+                accessKeyId: roleCredentials!.accessKeyId!,
+                secretAccessKey: roleCredentials!.secretAccessKey!,
+                sessionToken: roleCredentials?.sessionToken,
+                expiration,
             }
         } catch (err) {
-            if ((err as { code: string }).code === 'UnauthorizedException') {
+            if (err instanceof UnauthorizedException) {
                 this.ssoAccessTokenProvider.invalidate()
             }
             vscode.window.showErrorMessage(

src/shared/awsClientBuilder.ts

@@ -3,7 +3,8 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import { Request, AWSError } from 'aws-sdk'
+import { Request, AWSError, Credentials } from 'aws-sdk'
+import { CredentialsOptions } from 'aws-sdk/lib/credentials'
 import { ServiceConfigurationOptions } from 'aws-sdk/lib/service'
 import { env, version } from 'vscode'
 import { AwsContext } from './awsContext'
@@ -60,11 +61,7 @@ export interface AWSClientBuilder {
 }
 
 export class DefaultAWSClientBuilder implements AWSClientBuilder {
-    private readonly _awsContext: AwsContext
-
-    public constructor(awsContext: AwsContext) {
-        this._awsContext = awsContext
-    }
+    public constructor(private readonly awsContext: AwsContext) {}
 
     public async createAwsService<T extends AWS.Service>(
         type: new (o: ServiceConfigurationOptions) => T,
@@ -78,7 +75,46 @@ export class DefaultAWSClientBuilder implements AWSClientBuilder {
         delete opt.onRequestSetup
 
         if (!opt.credentials) {
-            opt.credentials = await this._awsContext.getCredentials()
+            const shim = this.awsContext.credentialsShim
+
+            if (!shim) {
+                throw new Error('Toolkit is not logged-in.')
+            }
+
+            opt.credentials = new (class extends Credentials {
+                public constructor() {
+                    // The class doesn't like being instantiated with empty creds
+                    super({ accessKeyId: '???', secretAccessKey: '???' })
+                }
+
+                public override get(callback: (err?: AWSError) => void): void {
+                    // Always try to fetch the latest creds first, attempting a refresh if needed
+                    // A 'passive' refresh is attempted first, before trying an 'active' one if certain criteria are met
+                    shim.get()
+                        .then(creds => {
+                            this.loadCreds(creds)
+                            this.needsRefresh() ? this.refresh(callback) : callback()
+                        })
+                        .catch(callback)
+                }
+
+                public override refresh(callback: (err?: AWSError) => void): void {
+                    shim.refresh()
+                        .then(creds => {
+                            this.loadCreds(creds)
+                            callback()
+                        })
+                        .catch(callback)
+                }
+
+                private loadCreds(creds: CredentialsOptions & { expiration?: Date }) {
+                    this.expired = false
+                    this.accessKeyId = creds.accessKeyId
+                    this.secretAccessKey = creds.secretAccessKey
+                    this.sessionToken = creds.sessionToken ?? this.sessionToken
+                    this.expireTime = creds.expiration ?? this.expireTime
+                }
+            })()
         }
 
         if (!opt.region && region) {

src/shared/awsContext.ts

@@ -8,6 +8,7 @@ import * as AWS from '@aws-sdk/types'
 import { regionSettingKey } from './constants'
 import { getLogger } from '../shared/logger'
 import { ClassToInterfaceType } from './utilities/tsUtils'
+import { CredentialsShim } from '../credentials/loginManager'
 
 export interface AwsContextCredentials {
     readonly credentials: AWS.Credentials
@@ -39,6 +40,7 @@ const DEFAULT_REGION = 'us-east-1'
 export class DefaultAwsContext implements AwsContext {
     public readonly onDidChangeContext: vscode.Event<ContextChangeEventsArgs>
     private readonly _onDidChangeContext: vscode.EventEmitter<ContextChangeEventsArgs>
+    private shim?: CredentialsShim
 
     // the collection of regions the user has expressed an interest in working with in
     // the current workspace
@@ -54,6 +56,14 @@ export class DefaultAwsContext implements AwsContext {
         this.explorerRegions = persistedRegions || []
     }
 
+    public get credentialsShim(): CredentialsShim | undefined {
+        return this.shim
+    }
+
+    public set credentialsShim(shim: CredentialsShim | undefined) {
+        this.shim = shim
+    }
+
     /**
      * Sets the credentials to be used by the Toolkit.
      * Passing in undefined represents that there are no active credentials.
@@ -74,7 +84,12 @@ export class DefaultAwsContext implements AwsContext {
      * @description Gets the Credentials currently used by the Toolkit.
      */
     public async getCredentials(): Promise<AWS.Credentials | undefined> {
-        return this.currentCredentials?.credentials
+        return (
+            this.shim?.get().catch(error => {
+                getLogger().warn(`credentials: failed to retrieve latest credentials: ${error.message}`)
+                return undefined
+            }) ?? this.currentCredentials?.credentials
+        )
     }
 
     // returns the configured profile, if any