auth: add telemetry for why the user needed to reauth (aws#5243)

* util: Add in-memory cache utils interface

Problem:

There is a common use for giving a key and getting back a cached
value. We have an existing implementation named KeyedCache, but it
utilizes Promises, leaving it open to race conditions.

Solution:

Create a synchronous interface similar to KeyedCache.

Signed-off-by: Nikolas Komonen <[email protected]>

* auth: createToken accepts args related to reauthentication

Problem:

SsoAccessTokenProvider.createToken() is used for both authentication and reauthentication, but
there isn't currently any way to identify the reason. This leaves a gap
in our telemetry.

Solution:

Created an object CreateTokenArgs which will hold args to be used by
createToken(). Now we can pass data to other parts of SsoAccessTokenProvider
that may want this data.

Most of the changes are adding in this new argument to createToken(). The notable
change is in auth.ts#Auth.reauthenticat() where we now pass isReAuth in the createToken() call.

Signed-off-by: Nikolas Komonen <[email protected]>

* auth: add telemetry for why the user needed to reauth

Problem:

In the reauth metric (aws_loginWithBrowser + `isReAuth: true`) we didn't know why
the user needed to reauth. It could have been due to regular expiration or some unexpected
error.

Solution:

- When we invalidate an SSO connection, record the reason
  in ReAuthState. Then next time we reauth we can look at this state to see why.
- The ReAuthState is in the SsoAccessToken provider since invalidation happens
  in this class itself and not in Auth. The SsoAccessTokenProvider.invalidate()
  is called to invalidate a connection and knows why the connection was invalidated.

Signed-off-by: Nikolas Komonen <[email protected]>

* refactor: remove hacky telemetry for reauth

Now that we've changed how we pass the isReAuth and reAuthReason
values to the final telemetry metric, we can get rid of the hack
we used to pass these values in Auth.reauthenticate().

This commit is removing redundant code. SsoAccessTokenProvider.withBrowserLoginTelemetry()
essentially has all the information it needs and doesn't need Auth.reauthenticate()
to use `record()` to externally setup values.

Signed-off-by: Nikolas Komonen <[email protected]>

* run formatter

Signed-off-by: Nikolas Komonen <[email protected]>

* refactor: fix PR comments

- Add some comments regarding why we call 'MapSync' with 'sync' in the name
- Rename 'MemoryMap' to 'NestedMap'
- rename some of the Map interface functions to be closer to the ES Map interface

Signed-off-by: Nikolas Komonen <[email protected]>

* refactor: merge new reauth state/telemetry tests with existing tests

Signed-off-by: Nikolas Komonen <[email protected]>

* rename: asKey() -> hash()

Signed-off-by: Nikolas Komonen <[email protected]>

---------

Signed-off-by: Nikolas Komonen <[email protected]>

Author: nkomonen-amazon

packages/core/src/auth/auth.ts

@@ -182,13 +182,7 @@ export class Auth implements AuthService, ConnectionManager {
         const profile = this.store.getProfileOrThrow(id)
         if (profile.type === 'sso') {
             const provider = this.getSsoTokenProvider(id, profile)
-
-            // We cannot easily set isReAuth inside the createToken() call,
-            // so we need to set it here.
-            await telemetry.aws_loginWithBrowser.run(async (span) => {
-                span.record({ isReAuth: true, credentialStartUrl: profile.startUrl })
-                await this.authenticate(id, () => provider.createToken(), shouldInvalidate)
-            })
+            await this.authenticate(id, () => provider.createToken({ isReAuth: true }), shouldInvalidate)
 
             return this.getSsoConnection(id, profile)
         } else {
@@ -386,7 +380,7 @@ export class Auth implements AuthService, ConnectionManager {
             throw new ToolkitError('Auth: Cannot force expire an IAM connection')
         }
         const provider = this.getSsoTokenProvider(conn.id, profile)
-        await provider.invalidate()
+        await provider.invalidate('devModeManualExpiration')
         // updates the state of the connection
         await this.refreshConnectionState(conn)
     }
@@ -530,7 +524,8 @@ export class Auth implements AuthService, ConnectionManager {
 
             // XXX: never drop tokens in a dev environment, unless you are Amazon Q!
             if (getCodeCatalystDevEnvId() === undefined || isAmazonQ()) {
-                await provider.invalidate()
+                // TODO: Require invalidateConnection() to require a reason and then pass it in to the following instead
+                await provider.invalidate('explicitInvalidation')
             }
         } else if (profile.type === 'iam') {
             globals.loginManager.store.invalidateCredentials(fromString(id))

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

@@ -217,7 +217,7 @@ export class SsoClient {
     private async handleError(error: unknown): Promise<never> {
         if (error instanceof SSOServiceException && isClientFault(error) && error.name !== 'ForbiddenException') {
             getLogger().warn(`credentials (sso): invalidating stored token: ${error.message}`)
-            await this.provider.invalidate()
+            await this.provider.invalidate(`ssoClient:${error.name}`)
         }
 
         throw error

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

@@ -29,7 +29,7 @@ import {
     isNetworkError,
 } from '../../shared/errors'
 import { getLogger } from '../../shared/logger'
-import { AwsLoginWithBrowser, AwsRefreshCredentials, Metric, telemetry } from '../../shared/telemetry/telemetry'
+import { AwsLoginWithBrowser, AwsRefreshCredentials, telemetry } from '../../shared/telemetry/telemetry'
 import { indent, toBase64URL } from '../../shared/utilities/textUtilities'
 import { AuthSSOServer } from './server'
 import { CancellationError, sleep } from '../../shared/utilities/timeoutUtils'
@@ -41,6 +41,7 @@ import { isRemoteWorkspace, isWebWorkspace } from '../../shared/vscode/env'
 import { showInputBox } from '../../shared/ui/inputPrompter'
 import { DevSettings } from '../../shared/settings'
 import { onceChanged } from '../../shared/utilities/functionUtils'
+import { NestedMap } from '../../shared/utilities/map'
 
 export const authenticationPath = 'sso/authenticated'
 
@@ -67,16 +68,18 @@ export abstract class SsoAccessTokenProvider {
     public constructor(
         protected readonly profile: Pick<SsoProfile, 'startUrl' | 'region' | 'scopes' | 'identifier'>,
         protected readonly cache = getCache(),
-        protected readonly oidc: OidcClient = OidcClient.create(profile.region)
+        protected readonly oidc: OidcClient = OidcClient.create(profile.region),
+        protected readonly reAuthState: ReAuthState = ReAuthState.instance
     ) {}
 
-    public async invalidate(): Promise<void> {
+    public async invalidate(reason: string): Promise<void> {
         getLogger().info(`SsoAccessTokenProvider: invalidate token and registration`)
         // Use allSettled() instead of all() to ensure all clear() calls are resolved.
         await Promise.allSettled([
             this.cache.token.clear(this.tokenCacheKey, 'SsoAccessTokenProvider.invalidate()'),
             this.cache.registration.clear(this.registrationCacheKey, 'SsoAccessTokenProvider.invalidate()'),
         ])
+        this.reAuthState.set(this.profile, { reAuthReason: `invalidate():${reason}` })
     }
 
     public async getToken(): Promise<SsoToken | undefined> {
@@ -99,23 +102,30 @@ export abstract class SsoAccessTokenProvider {
 
             return refreshed.token
         } else {
-            await this.invalidate()
+            await this.invalidate('allCacheExpired')
         }
     }
 
-    public async createToken(): Promise<SsoToken> {
-        const access = await this.runFlow()
+    public async createToken(args?: CreateTokenArgs): Promise<SsoToken> {
+        const access = await this.runFlow(args)
         const identity = this.tokenCacheKey
         await this.cache.token.save(identity, access)
         await globals.globalState.setSsoSessionCreationDate(this.tokenCacheKey, new globals.clock.Date())
 
         return { ...access.token, identity }
     }
 
-    private async runFlow() {
+    private async runFlow(args?: CreateTokenArgs) {
         const registration = await this.getValidatedClientRegistration()
         try {
-            return await this.authorize(registration)
+            const result = await this.authorize(registration, args)
+
+            // Authentication in the browser is successfully done, so the reauth reason is now stale.
+            // We don't clear the reason on failure since we want to keep reporting it as the reason until
+            // reauth is a success.
+            this.reAuthState.delete(this.profile, 'reauth successful')
+
+            return result
         } catch (err) {
             if (err instanceof SSOOIDCServiceException && isClientFault(err)) {
                 await this.cache.registration.clear(
@@ -150,9 +160,10 @@ export abstract class SsoAccessTokenProvider {
             return refreshed
         } catch (err) {
             if (!isNetworkError(err)) {
+                const reason = getTelemetryReason(err)
                 telemetry.aws_refreshCredentials.emit({
                     result: getTelemetryResult(err),
-                    reason: getTelemetryReason(err),
+                    reason,
                     reasonDesc: getTelemetryReasonDesc(err),
                     requestId: getRequestId(err),
                     ...metric,
@@ -163,6 +174,10 @@ export abstract class SsoAccessTokenProvider {
                         this.tokenCacheKey,
                         `client fault: SSOOIDCServiceException: ${err.message}`
                     )
+                    // remember why refresh failed so next reauth flow will know why reauth is needed
+                    if (reason) {
+                        this.reAuthState.set(this.profile, { reAuthReason: `refresh:${reason}` })
+                    }
                 }
             }
 
@@ -185,33 +200,30 @@ export abstract class SsoAccessTokenProvider {
     /**
      * Wraps the given function with telemetry related to the browser login.
      */
-    protected withBrowserLoginTelemetry<T extends (...args: any[]) => any>(func: T): ReturnType<T> {
-        const run = (span: Metric<AwsLoginWithBrowser>) => {
+    protected withBrowserLoginTelemetry<T extends (...args: any[]) => any>(
+        func: T,
+        args?: CreateTokenArgs
+    ): ReturnType<T> {
+        return telemetry.aws_loginWithBrowser.run((span) => {
             span.record({
                 credentialStartUrl: this.profile.startUrl,
                 source: SsoAccessTokenProvider._authSource,
+                isReAuth: args?.isReAuth,
+                reAuthReason: args?.isReAuth ? this.reAuthState.get(this.profile).reAuthReason : undefined,
             })
 
             // Reset source in case there is a case where browser login was called but we forgot to set the source.
             // We don't want to attribute the wrong source.
             SsoAccessTokenProvider.authSource = 'unknown'
 
             return func()
-        }
-
-        // During certain flows, eg reauthentication, we are already running within a span (run())
-        // so we don't need to create a new one.
-        const span = telemetry.spans.find((s) => s.name === 'aws_loginWithBrowser')
-        if (span !== undefined) {
-            return run(span as unknown as Metric<AwsLoginWithBrowser>)
-        }
-
-        return telemetry.aws_loginWithBrowser.run((span) => {
-            return run(span)
         })
     }
 
-    protected abstract authorize(registration: ClientRegistration): Promise<{
+    protected abstract authorize(
+        registration: ClientRegistration,
+        args?: CreateTokenArgs
+    ): Promise<{
         token: SsoToken
         registration: ClientRegistration
         region: string
@@ -230,6 +242,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:
@@ -242,15 +255,24 @@ export abstract class SsoAccessTokenProvider {
         }
     ) {
         if (DevSettings.instance.get('webAuth', false) && isWebWorkspace()) {
-            return new WebAuthorization(profile, cache, oidc)
+            return new WebAuthorization(profile, cache, oidc, reAuthState)
         }
         if (useDeviceFlow()) {
-            return new DeviceFlowAuthorization(profile, cache, oidc)
+            return new DeviceFlowAuthorization(profile, cache, oidc, reAuthState)
         }
-        return new AuthFlowAuthorization(profile, cache, oidc)
+        return new AuthFlowAuthorization(profile, cache, oidc, reAuthState)
     }
 }
 
+/**
+ * Supplementary arguments for the create token flow. This data can be used
+ * for things like telemetry.
+ */
+export type CreateTokenArgs = {
+    /** true if the create token flow is for reauthentication */
+    isReAuth?: boolean
+}
+
 const backoffDelayMs = 5000
 async function pollForTokenWithProgress<T extends { requestId?: string }>(
     fn: () => Promise<T>,
@@ -356,14 +378,6 @@ function getSessionDuration(id: string) {
  *         - RefreshToken (optional)
  */
 export class DeviceFlowAuthorization extends SsoAccessTokenProvider {
-    constructor(
-        profile: Pick<SsoProfile, 'startUrl' | 'region' | 'scopes' | 'identifier'>,
-        cache = getCache(),
-        oidc: OidcClient = OidcClient.create(profile.region)
-    ) {
-        super(profile, cache, oidc)
-    }
-
     override async registerClient(): Promise<ClientRegistration> {
         const companyName = getIdeProperties().company
         return this.oidc.registerClient(
@@ -377,7 +391,8 @@ export class DeviceFlowAuthorization extends SsoAccessTokenProvider {
     }
 
     override async authorize(
-        registration: ClientRegistration
+        registration: ClientRegistration,
+        args?: CreateTokenArgs
     ): Promise<{ token: SsoToken; registration: ClientRegistration; region: string; startUrl: string }> {
         // This will NOT throw on expired clientId/Secret, but WILL throw on invalid clientId/Secret
         const authorization = await this.oidc.startDeviceAuthorization({
@@ -403,7 +418,7 @@ export class DeviceFlowAuthorization extends SsoAccessTokenProvider {
             )
         }
 
-        const token = this.withBrowserLoginTelemetry(() => openBrowserAndWaitUntilComplete())
+        const token = this.withBrowserLoginTelemetry(() => openBrowserAndWaitUntilComplete(), args)
 
         return this.formatToken(await token, registration)
     }
@@ -419,7 +434,7 @@ export class DeviceFlowAuthorization extends SsoAccessTokenProvider {
 
         // Clear cached if registration is expired
         if (cachedRegistration && isExpired(cachedRegistration)) {
-            await this.invalidate()
+            await this.invalidate('registrationExpired:DeviceCode')
         }
 
         return loadOr(this.cache.registration, cacheKey, () => this.registerClient())
@@ -463,14 +478,6 @@ export class DeviceFlowAuthorization extends SsoAccessTokenProvider {
  *               2. If there is a problem, server responds with `invalid_grant` error.
  */
 class AuthFlowAuthorization extends SsoAccessTokenProvider {
-    constructor(
-        profile: Pick<SsoProfile, 'startUrl' | 'region' | 'scopes' | 'identifier'>,
-        cache = getCache(),
-        oidc: OidcClient
-    ) {
-        super(profile, cache, oidc)
-    }
-
     override async registerClient(): Promise<ClientRegistration> {
         const companyName = getIdeProperties().company
         return this.oidc.registerClient(
@@ -489,7 +496,8 @@ class AuthFlowAuthorization extends SsoAccessTokenProvider {
     }
 
     override async authorize(
-        registration: ClientRegistration
+        registration: ClientRegistration,
+        args?: CreateTokenArgs
     ): Promise<{ token: SsoToken; registration: ClientRegistration; region: string; startUrl: string }> {
         const state = randomUUID()
         const authServer = AuthSSOServer.init(state)
@@ -531,7 +539,7 @@ class AuthFlowAuthorization extends SsoAccessTokenProvider {
                 telemetry.record({ requestId: res.requestId })
 
                 return res
-            })
+            }, args)
 
             return this.formatToken(token, registration)
         } finally {
@@ -560,7 +568,7 @@ class AuthFlowAuthorization extends SsoAccessTokenProvider {
 
         // Clear cached if registration is expired or it uses a deprecate auth version (device code)
         if (cachedRegistration && (isExpired(cachedRegistration) || isDeprecatedAuth(cachedRegistration))) {
-            await this.invalidate()
+            await this.invalidate('registrationExpired:AuthFlow')
         }
 
         return loadOr(this.cache.registration, cacheKey, () => this.registerClient())
@@ -575,14 +583,6 @@ class AuthFlowAuthorization extends SsoAccessTokenProvider {
 class WebAuthorization extends SsoAccessTokenProvider {
     private redirectUri = 'http://127.0.0.1:54321/oauth/callback'
 
-    constructor(
-        profile: Pick<SsoProfile, 'startUrl' | 'region' | 'scopes' | 'identifier'>,
-        cache = getCache(),
-        oidc: OidcClient
-    ) {
-        super(profile, cache, oidc)
-    }
-
     override async registerClient(): Promise<ClientRegistration> {
         const companyName = getIdeProperties().company
         return this.oidc.registerClient(
@@ -601,7 +601,8 @@ class WebAuthorization extends SsoAccessTokenProvider {
     }
 
     override async authorize(
-        registration: ClientRegistration
+        registration: ClientRegistration,
+        args?: CreateTokenArgs
     ): Promise<{ token: SsoToken; registration: ClientRegistration; region: string; startUrl: string }> {
         const state = randomUUID()
 
@@ -641,7 +642,7 @@ class WebAuthorization extends SsoAccessTokenProvider {
                 codeVerifier,
                 code: inputBox,
             })
-        })
+        }, args)
 
         return this.formatToken(token, registration)
     }
@@ -651,9 +652,46 @@ class WebAuthorization extends SsoAccessTokenProvider {
         const cachedRegistration = await this.cache.registration.load(cacheKey)
 
         if (cachedRegistration && (isExpired(cachedRegistration) || cachedRegistration.flow !== 'web auth code')) {
-            await this.invalidate()
+            await this.invalidate('registrationExpired:WebAuth')
         }
 
         return loadOr(this.cache.registration, cacheKey, () => this.registerClient())
     }
 }
+
+/**
+ * Remembers the reason an SSO session was put in to a "needs reauthentication" state.
+ * The current use is for telemetry. When the user reauths, we want {@link AwsLoginWithBrowser}
+ * to know why it needed to be reauthed.
+ *
+ * The flow is to use `set()` to remember why the user was put in to a reauth state,
+ * then upon the next reauth use `get()`. Finally, use `clear()` if the reauth is
+ * successful.
+ */
+export class ReAuthState extends NestedMap<ReAuthStateKey, ReAuthStateValue> {
+    static #instance: ReAuthState
+    static get instance() {
+        return (this.#instance ??= new ReAuthState())
+    }
+    protected constructor() {
+        super()
+    }
+
+    protected override hash(profile: ReAuthStateKey): string {
+        return profile.identifier ?? profile.startUrl
+    }
+
+    protected override get name(): string {
+        return ReAuthState.name
+    }
+
+    override get default(): ReAuthStateValue {
+        return { reAuthReason: undefined }
+    }
+}
+
+type ReAuthStateKey = Pick<SsoProfile, 'identifier' | 'startUrl'>
+type ReAuthStateValue = {
+    // the latest reason for why the connection was moved in to a "needs reauth" state
+    reAuthReason?: string
+}