fix(auth): more logging in OIDC client #3265

* Add extra logging to the OIDC client
* Log exceptions
* Add telemetry

Author: JadenSimon

package-lock.json

@@ -3468,7 +3468,7 @@
         "report": "nyc report --reporter=html --reporter=json"
     },
     "devDependencies": {
-        "@aws-toolkits/telemetry": "^1.0.100",
+        "@aws-toolkits/telemetry": "^1.0.120",
         "@cspotcode/source-map-support": "^0.8.1",
         "@sinonjs/fake-timers": "^8.1.0",
         "@types/adm-zip": "^0.4.34",

package.json

@@ -16,6 +16,7 @@ import {
     SSOServiceException,
 } from '@aws-sdk/client-sso'
 import {
+    AuthorizationPendingException,
     CreateTokenRequest,
     RegisterClientRequest,
     SSOOIDC,
@@ -28,6 +29,9 @@ import { getLogger } from '../../shared/logger'
 import { SsoAccessTokenProvider } from './ssoAccessTokenProvider'
 import { isClientFault } from '../../shared/errors'
 import { DevSettings } from '../../shared/settings'
+import { Client } from '@aws-sdk/smithy-client'
+import { HttpHandlerOptions } from '@aws-sdk/types'
+import { HttpRequest, HttpResponse } from '@aws-sdk/protocol-http'
 
 export class OidcClient {
     public constructor(private readonly client: SSOOIDC, private readonly clock: { Date: typeof Date }) {}
@@ -66,13 +70,13 @@ export class OidcClient {
     }
 
     public static create(region: string) {
-        return new this(
-            new SSOOIDC({
-                region,
-                endpoint: DevSettings.instance.get('endpoints', {})['ssooidc'],
-            }),
-            globals.clock
-        )
+        const client = new SSOOIDC({
+            region,
+            endpoint: DevSettings.instance.get('endpoints', {})['ssooidc'],
+        })
+
+        addLoggingMiddleware(client)
+        return new this(client, globals.clock)
     }
 }
 
@@ -174,3 +178,52 @@ export class SsoClient {
         )
     }
 }
+
+function omitIfPresent<T extends Record<string, unknown>>(obj: T, ...keys: string[]): T {
+    const objCopy = { ...obj }
+    for (const key of keys) {
+        if (key in objCopy) {
+            ;(objCopy as any)[key] = '[omitted]'
+        }
+    }
+    return objCopy
+}
+
+function addLoggingMiddleware(client: Client<HttpHandlerOptions, any, any, any>) {
+    client.middlewareStack.add(
+        (next, context) => args => {
+            if (HttpRequest.isInstance(args.request)) {
+                const { hostname, path } = args.request
+                const input = omitIfPresent(args.input, 'clientSecret', 'accessToken', 'refreshToken')
+                getLogger().debug('API request (%s %s): %O', hostname, path, input)
+            }
+            return next(args)
+        },
+        { step: 'finalizeRequest' }
+    )
+
+    client.middlewareStack.add(
+        (next, context) => async args => {
+            if (!HttpRequest.isInstance(args.request)) {
+                return next(args)
+            }
+
+            const { hostname, path } = args.request
+            const result = await next(args).catch(e => {
+                if (e instanceof Error && !(e instanceof AuthorizationPendingException)) {
+                    const err = { ...e }
+                    delete err['stack']
+                    getLogger().error('API response (%s %s): %O', hostname, path, err)
+                }
+                throw e
+            })
+            if (HttpResponse.isInstance(result.response)) {
+                const output = omitIfPresent(result.output, 'clientSecret', 'accessToken', 'refreshToken')
+                getLogger().debug('API response (%s %s): %O', hostname, path, output)
+            }
+
+            return result
+        },
+        { step: 'deserialize' }
+    )
+}

src/credentials/sso/clients.ts

@@ -9,13 +9,15 @@ const localize = nls.loadMessageBundle()
 import globals from '../../shared/extensionGlobals'
 import * as vscode from 'vscode'
 import { AuthorizationPendingException, SlowDownException, SSOOIDCServiceException } from '@aws-sdk/client-sso-oidc'
-import { openSsoPortalLink, SsoToken, ClientRegistration, isExpired, SsoProfile } from './model'
+import { openSsoPortalLink, SsoToken, ClientRegistration, isExpired, SsoProfile, builderIdStartUrl } from './model'
 import { getCache } from './cache'
 import { hasProps, hasStringProps, RequiredProps, selectFrom } from '../../shared/utilities/tsUtils'
 import { CancellationError, sleep } from '../../shared/utilities/timeoutUtils'
 import { OidcClient } from './clients'
 import { loadOr } from '../../shared/utilities/cacheUtils'
-import { isClientFault, ToolkitError } from '../../shared/errors'
+import { getTelemetryReason, getTelemetryResult, isClientFault, ToolkitError } from '../../shared/errors'
+import { getLogger } from '../../shared/logger'
+import { telemetry } from '../../shared/telemetry/telemetry'
 
 const clientRegistrationType = 'public'
 const deviceGrantType = 'urn:ietf:params:oauth:grant-type:device_code'
@@ -81,11 +83,7 @@ export class SsoAccessTokenProvider {
         if (data.registration && !isExpired(data.registration) && hasProps(data.token, 'refreshToken')) {
             const refreshed = await this.refreshToken(data.token, data.registration)
 
-            if (refreshed) {
-                await this.cache.token.save(this.tokenCacheKey, refreshed)
-            }
-
-            return refreshed?.token
+            return refreshed.token
         } else {
             await this.invalidate()
         }
@@ -95,6 +93,7 @@ export class SsoAccessTokenProvider {
         const access = await this.runFlow()
         const identity = (await identityProvider?.(access.token)) ?? this.tokenCacheKey
         await this.cache.token.save(identity, access)
+        await setSessionCreationDate(this.tokenCacheKey, new Date())
 
         return { ...access.token, identity }
     }
@@ -118,9 +117,19 @@ export class SsoAccessTokenProvider {
         try {
             const clientInfo = selectFrom(registration, 'clientId', 'clientSecret')
             const response = await this.oidc.createToken({ ...clientInfo, ...token, grantType: refreshGrantType })
+            const refreshed = this.formatToken(response, registration)
+            await this.cache.token.save(this.tokenCacheKey, refreshed)
 
-            return this.formatToken(response, registration)
+            return refreshed
         } catch (err) {
+            telemetry.aws_refreshCredentials.emit({
+                result: getTelemetryResult(err),
+                reason: getTelemetryReason(err),
+                sessionDuration: getSessionDuration(this.tokenCacheKey),
+                credentialType: 'bearerToken',
+                credentialSourceId: this.profile.startUrl === builderIdStartUrl ? 'awsId' : 'iamIdentityCenter',
+            })
+
             if (err instanceof SSOOIDCServiceException && isClientFault(err)) {
                 await this.cache.token.clear(this.tokenCacheKey)
             }
@@ -228,3 +237,25 @@ async function pollForTokenWithProgress<T>(
             ])
     )
 }
+
+const sessionCreationDateKey = '#sessionCreationDates'
+async function setSessionCreationDate(id: string, date: Date, memento = globals.context.globalState) {
+    try {
+        await memento.update(sessionCreationDateKey, {
+            ...memento.get(sessionCreationDateKey),
+            [id]: date.getTime(),
+        })
+    } catch (err) {
+        getLogger().verbose('auth: failed to set session creation date: %s', err)
+    }
+}
+
+function getSessionCreationDate(id: string, memento = globals.context.globalState): number | undefined {
+    return memento.get(sessionCreationDateKey, {} as Record<string, number>)[id]
+}
+
+function getSessionDuration(id: string, memento = globals.context.globalState) {
+    const creationDate = getSessionCreationDate(id, memento)
+
+    return creationDate !== undefined ? Date.now() - creationDate : undefined
+}

src/credentials/sso/ssoAccessTokenProvider.ts

@@ -59,24 +59,6 @@
             "description": "Called after attempting to install a local copy of a missing CLI",
             "metadata": [{ "type": "result" }, { "type": "cli" }]
         },
-        {
-            "name": "aws_refreshCredentials",
-            "description": "Emitted when credentials are automatically refreshed by the AWS SDK",
-            "passive": true,
-            "metadata": [
-                {
-                    "type": "result"
-                },
-                {
-                    "type": "credentialType",
-                    "required": false
-                },
-                {
-                    "type": "credentialSourceId",
-                    "required": false
-                }
-            ]
-        },
         {
             "name": "ssm_createDocument",
             "description": "An SSM Document is created locally",