refactor: unify middleware definitions

Author: Hweinstock

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

@@ -36,6 +36,7 @@ import { StandardRetryStrategy, defaultRetryDecider } from '@smithy/middleware-r
 import { AuthenticationFlow } from './model'
 import { toSnakeCase } from '../../shared/utilities/textUtilities'
 import { getUserAgent, withTelemetryContext } from '../../shared/telemetry/util'
+import { defaultDeserializeMiddleware, finalizeLoggingMiddleware } from '../../shared/awsClientBuilderV3'
 
 export class OidcClient {
     public constructor(
@@ -249,52 +250,6 @@ export class SsoClient {
 }
 
 function addLoggingMiddleware(client: SSOOIDCClient) {
-    client.middlewareStack.add(
-        (next, context) => (args) => {
-            if (HttpRequest.isInstance(args.request)) {
-                const { hostname, path } = args.request
-                const input = partialClone(
-                    // TODO: Fix
-                    args.input as unknown as Record<string, unknown>,
-                    3,
-                    ['clientSecret', 'accessToken', 'refreshToken'],
-                    '[omitted]'
-                )
-                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 = partialClone(
-                    // TODO: Fix
-                    result.output as unknown as Record<string, unknown>,
-                    3,
-                    ['clientSecret', 'accessToken', 'refreshToken'],
-                    '[omitted]'
-                )
-                getLogger().debug('API response (%s %s): %O', hostname, path, output)
-            }
-
-            return result
-        },
-        { step: 'deserialize' }
-    )
+    client.middlewareStack.add(finalizeLoggingMiddleware, { step: 'finalizeRequest' })
+    client.middlewareStack.add(defaultDeserializeMiddleware, { step: 'deserialize' })
 }

packages/core/src/shared/awsClientBuilderV3.ts

@@ -42,6 +42,7 @@ import { partialClone } from './utilities/collectionUtils'
 import { selectFrom } from './utilities/tsUtils'
 import { once } from './utilities/functionUtils'
 import { isWeb } from './extensionGlobals'
+import { AuthorizationPendingException } from '@aws-sdk/client-sso-oidc'
 
 export type AwsClientConstructor<C> = new (o: AwsClientOptions) => C
 export type AwsCommandConstructor<CommandInput extends object, Command extends AwsCommand<CommandInput, object>> = new (
@@ -175,8 +176,8 @@ export class AWSClientBuilderV3 {
         }
 
         const service = new serviceOptions.serviceClient(opt)
-        service.middlewareStack.add(telemetryMiddleware, { step: 'deserialize' })
-        service.middlewareStack.add(loggingMiddleware, { step: 'finalizeRequest' })
+        service.middlewareStack.add(defaultDeserializeMiddleware, { step: 'deserialize' })
+        service.middlewareStack.add(finalizeLoggingMiddleware, { step: 'finalizeRequest' })
         service.middlewareStack.add(getEndpointMiddleware(serviceOptions.settings), { step: 'build' })
 
         if (keepAlive) {
@@ -211,20 +212,13 @@ export function recordErrorTelemetry(err: Error, serviceName?: string) {
     })
 }
 
-function logAndThrow(e: any, serviceId: string, errorMessageAppend: string): never {
-    if (e instanceof Error) {
-        recordErrorTelemetry(e, serviceId)
-        getLogger().error('API Response %s: %O', errorMessageAppend, e)
-    }
-    throw e
-}
-
-const telemetryMiddleware: DeserializeMiddleware<any, any> =
+export const defaultDeserializeMiddleware: DeserializeMiddleware<any, any> =
     (next: DeserializeHandler<any, any>, context: HandlerExecutionContext) => async (args: any) =>
-        emitOnRequest(next, context, args)
+        onDeserialize(next, context, args)
 
-const loggingMiddleware: FinalizeRequestMiddleware<any, any> = (next: FinalizeHandler<any, any>) => async (args: any) =>
-    logOnRequest(next, args)
+export const finalizeLoggingMiddleware: FinalizeRequestMiddleware<any, any> =
+    (next: FinalizeHandler<any, any>) => async (args: any) =>
+        logOnFinalize(next, args)
 
 function getEndpointMiddleware(settings: DevSettings = DevSettings.instance): BuildMiddleware<any, any> {
     return (next: BuildHandler<any, any>, context: HandlerExecutionContext) => async (args: any) =>
@@ -234,32 +228,36 @@ function getEndpointMiddleware(settings: DevSettings = DevSettings.instance): Bu
 const keepAliveMiddleware: BuildMiddleware<any, any> = (next: BuildHandler<any, any>) => async (args: any) =>
     addKeepAliveHeader(next, args)
 
-export async function emitOnRequest(next: DeserializeHandler<any, any>, context: HandlerExecutionContext, args: any) {
+export async function onDeserialize(next: DeserializeHandler<any, any>, context: HandlerExecutionContext, args: any) {
     if (!HttpResponse.isInstance(args.request)) {
         return next(args)
     }
-    const serviceId = getServiceId(context as object)
     const { hostname, path } = args.request
+    const serviceId = getServiceId(context as object)
     const logTail = `(${hostname} ${path})`
     try {
         const result = await next(args)
         if (HttpResponse.isInstance(result.response)) {
-            // TODO: omit credentials / sensitive info from the telemetry.
-            const output = partialClone(result.output, 3)
+            const output = partialClone(result.output, 3, ['clientSecret', 'accessToken', 'refreshToken'], '[omitted]')
             getLogger().debug(`API Response %s: %O`, logTail, output)
         }
         return result
-    } catch (e: any) {
-        logAndThrow(e, serviceId, logTail)
+    } catch (e: unknown) {
+        if (e instanceof Error && !(e instanceof AuthorizationPendingException)) {
+            const err = { ...e }
+            delete err['stack']
+            recordErrorTelemetry(err, serviceId)
+            getLogger().error('API Response %s: %O', logTail, err)
+        }
+        throw e
     }
 }
 
-export async function logOnRequest(next: FinalizeHandler<any, any>, args: any) {
+export async function logOnFinalize(next: FinalizeHandler<any, any>, args: any) {
     const request = args.request
     if (HttpRequest.isInstance(args.request)) {
         const { hostname, path } = request
-        // TODO: omit credentials / sensitive info from the logs.
-        const input = partialClone(args.input, 3)
+        const input = partialClone(args.input, 3, ['clientSecret', 'accessToken', 'refreshToken'], '[omitted]')
         getLogger().debug(`API Request (%s %s): %O`, hostname, path, input)
     }
     return next(args)

packages/core/src/test/shared/awsClientBuilderV3.test.ts

@@ -14,9 +14,9 @@ import {
     AWSClientBuilderV3,
     AwsClientOptions,
     AwsCommand,
-    emitOnRequest,
+    onDeserialize,
     getServiceId,
-    logOnRequest,
+    logOnFinalize,
     overwriteEndpoint,
     recordErrorTelemetry,
 } from '../../shared/awsClientBuilderV3'
@@ -237,7 +237,7 @@ describe('AwsClientBuilderV3', function () {
         })
 
         it('logs messages on request', async function () {
-            await logOnRequest((_: any) => _, args as any)
+            await logOnFinalize((_: any) => _, args as any)
             assertLogsContainAllOf(['testHost', 'testPath'], false, 'debug')
         })
 
@@ -246,7 +246,7 @@ describe('AwsClientBuilderV3', function () {
                 throw new Error('test error')
             }
             await telemetry.vscode_executeCommand.run(async (span) => {
-                await assert.rejects(emitOnRequest(next, context, args))
+                await assert.rejects(onDeserialize(next, context, args))
             })
             assertLogsContain('test error', false, 'error')
             assertTelemetry('vscode_executeCommand', { requestServiceType: 'foo' })
@@ -257,7 +257,7 @@ describe('AwsClientBuilderV3', function () {
                 return response
             }
             await telemetry.vscode_executeCommand.run(async (span) => {
-                assert.deepStrictEqual(await emitOnRequest(next, context, args), response)
+                assert.deepStrictEqual(await onDeserialize(next, context, args), response)
             })
             assertLogsContainAllOf(['testHost', 'testPath'], false, 'debug')
             assert.throws(() => assertTelemetry('vscode_executeCommand', { requestServiceType: 'foo' }))