fix(lib-dynamodb): fix use of log filters in conjunction with DynamDBDocumentClient (#4249)

* fix(lib-dynamodb): fix use of log filters in conjunction with DynamoDBDocumentClient

* fix(lib-dynamodb): add tests in middleware-logger

Author: kuhe

Diff for: lib/lib-dynamodb/src/baseCommand/DynamoDBDocumentClientCommand.spec.ts

@@ -14,7 +14,7 @@ class AnyCommand extends DynamoDBDocumentClientCommand<{}, {}, {}, {}, {}> {
   protected readonly clientCommand = {
     middlewareStack: {
       argCaptor: this.argCaptor,
-      add(fn, config) {
+      addRelativeTo(fn, config) {
         this.argCaptor.push([fn, config]);
       },
     },
@@ -41,7 +41,8 @@ describe("DynamoDBDocumentClientCommand", () => {
       expect(options).toEqual({
         name: "DocumentMarshall",
         override: true,
-        step: "initialize",
+        relation: "before",
+        toMiddleware: "serializerMiddleware",
       });
     }
     {
@@ -50,7 +51,8 @@ describe("DynamoDBDocumentClientCommand", () => {
       expect(options).toEqual({
         name: "DocumentUnmarshall",
         override: true,
-        step: "deserialize",
+        relation: "before",
+        toMiddleware: "deserializerMiddleware",
       });
     }
   });

Diff for: lib/lib-dynamodb/src/baseCommand/DynamoDBDocumentClientCommand.ts

@@ -3,6 +3,7 @@ import {
   DeserializeHandler,
   DeserializeHandlerArguments,
   DeserializeHandlerOutput,
+  HandlerExecutionContext,
   InitializeHandler,
   InitializeHandlerArguments,
   InitializeHandlerOutput,
@@ -29,35 +30,64 @@ export abstract class DynamoDBDocumentClientCommand<
 
   public abstract middlewareStack: MiddlewareStack<Input | BaseInput, Output | BaseOutput>;
 
+  private static defaultLogFilterOverrides = {
+    overrideInputFilterSensitiveLog(...args: any[]) {},
+    overrideOutputFilterSensitiveLog(...args: any[]) {},
+  };
+
   protected addMarshallingMiddleware(configuration: DynamoDBDocumentClientResolvedConfig): void {
     const { marshallOptions, unmarshallOptions } = configuration.translateConfig || {};
 
-    this.clientCommand.middlewareStack.add(
-      (next: InitializeHandler<Input | BaseInput, Output | BaseOutput>) =>
+    this.clientCommand.middlewareStack.addRelativeTo(
+      (next: InitializeHandler<Input | BaseInput, Output | BaseOutput>, context: HandlerExecutionContext) =>
         async (
           args: InitializeHandlerArguments<Input | BaseInput>
         ): Promise<InitializeHandlerOutput<Output | BaseOutput>> => {
           args.input = marshallInput(this.input, this.inputKeyNodes, marshallOptions);
+          context.dynamoDbDocumentClientOptions =
+            context.dynamoDbDocumentClientOptions || DynamoDBDocumentClientCommand.defaultLogFilterOverrides;
+
+          const input = args.input;
+          context.dynamoDbDocumentClientOptions.overrideInputFilterSensitiveLog = () => {
+            return context.inputFilterSensitiveLog?.(input);
+          };
           return next(args);
         },
       {
         name: "DocumentMarshall",
-        step: "initialize",
+        relation: "before",
+        toMiddleware: "serializerMiddleware",
         override: true,
       }
     );
-    this.clientCommand.middlewareStack.add(
-      (next: DeserializeHandler<Input | BaseInput, Output | BaseOutput>) =>
+    this.clientCommand.middlewareStack.addRelativeTo(
+      (next: DeserializeHandler<Input | BaseInput, Output | BaseOutput>, context: HandlerExecutionContext) =>
         async (
           args: DeserializeHandlerArguments<Input | BaseInput>
         ): Promise<DeserializeHandlerOutput<Output | BaseOutput>> => {
           const deserialized = await next(args);
+
+          /**
+           * The original filter function is based on the shape of the
+           * base DynamoDB type, whereas the returned output will be
+           * unmarshalled. Therefore the filter log needs to be modified
+           * to act on the original output structure.
+           */
+          const output = deserialized.output;
+          context.dynamoDbDocumentClientOptions =
+            context.dynamoDbDocumentClientOptions || DynamoDBDocumentClientCommand.defaultLogFilterOverrides;
+
+          context.dynamoDbDocumentClientOptions.overrideOutputFilterSensitiveLog = () => {
+            return context.outputFilterSensitiveLog?.(output);
+          };
+
           deserialized.output = unmarshallOutput(deserialized.output, this.outputKeyNodes, unmarshallOptions);
           return deserialized;
         },
       {
         name: "DocumentUnmarshall",
-        step: "deserialize",
+        relation: "before",
+        toMiddleware: "deserializerMiddleware",
         override: true,
       }
     );

Diff for: packages/middleware-logger/src/loggerMiddleware.spec.ts

@@ -1,4 +1,4 @@
-import { Logger, MiddlewareStack } from "@aws-sdk/types";
+import { HandlerExecutionContext, Logger, MiddlewareStack } from "@aws-sdk/types";
 
 import { getLoggerPlugin, loggerMiddleware, loggerMiddlewareOptions } from "./loggerMiddleware";
 
@@ -114,6 +114,55 @@ describe("loggerMiddleware", () => {
       });
     });
 
+    it("should use override log filters for DynamoDBDocumentClient if present", async () => {
+      mockNext.mockResolvedValueOnce(mockResponse);
+
+      const logger = { info: jest.fn() } as unknown as Logger;
+      const clientName = "mockClientName";
+      const commandName = "mockCommandName";
+
+      const mockInputLog = { inputKey: "inputKey", inputSensitiveKey: "SENSITIVE_VALUE" };
+      const inputFilterSensitiveLog = jest.fn().mockReturnValueOnce(mockInputLog);
+      const mockOutputLog = { outputKey: "outputKey", outputSensitiveKey: "SENSITIVE_VALUE" };
+      const outputFilterSensitiveLog = jest.fn().mockReturnValueOnce(mockOutputLog);
+
+      const context: HandlerExecutionContext = {
+        logger,
+        clientName,
+        commandName,
+        dynamoDbDocumentClientOptions: {
+          overrideInputFilterSensitiveLog: inputFilterSensitiveLog,
+          overrideOutputFilterSensitiveLog: outputFilterSensitiveLog,
+        },
+        inputFilterSensitiveLog() {
+          throw new Error("should not be called");
+        },
+        outputFilterSensitiveLog() {
+          throw new Error("should not be called");
+        },
+      };
+
+      const response = await loggerMiddleware()(mockNext, context)(mockArgs);
+      expect(mockNext).toHaveBeenCalledTimes(1);
+      expect(response).toStrictEqual(mockResponse);
+
+      expect(inputFilterSensitiveLog).toHaveBeenCalledTimes(1);
+      expect(inputFilterSensitiveLog).toHaveBeenCalledWith(mockArgs.input);
+
+      const { $metadata, ...outputWithoutMetadata } = mockOutput;
+      expect(outputFilterSensitiveLog).toHaveBeenCalledTimes(1);
+      expect(outputFilterSensitiveLog).toHaveBeenCalledWith(outputWithoutMetadata);
+
+      expect(logger.info).toHaveBeenCalledTimes(1);
+      expect(logger.info).toHaveBeenCalledWith({
+        clientName,
+        commandName,
+        input: mockInputLog,
+        output: mockOutputLog,
+        metadata: $metadata,
+      });
+    });
+
     it("header x-amzn-request-id as requestId if x-amzn-requestid is not present", async () => {
       const requestIdBackup = "requestIdBackup";
       const customResponse = {

Diff for: packages/middleware-logger/src/loggerMiddleware.ts

@@ -16,21 +16,30 @@ export const loggerMiddleware =
     context: HandlerExecutionContext
   ): InitializeHandler<any, Output> =>
   async (args: InitializeHandlerArguments<any>): Promise<InitializeHandlerOutput<Output>> => {
-    const { clientName, commandName, inputFilterSensitiveLog, logger, outputFilterSensitiveLog } = context;
-
     const response = await next(args);
+    const {
+      clientName,
+      commandName,
+      logger,
+      inputFilterSensitiveLog,
+      outputFilterSensitiveLog,
+      dynamoDbDocumentClientOptions = {},
+    } = context;
+
+    const { overrideInputFilterSensitiveLog, overrideOutputFilterSensitiveLog } = dynamoDbDocumentClientOptions;
 
     if (!logger) {
       return response;
     }
 
     if (typeof logger.info === "function") {
       const { $metadata, ...outputWithoutMetadata } = response.output;
+
       logger.info({
         clientName,
         commandName,
-        input: inputFilterSensitiveLog(args.input),
-        output: outputFilterSensitiveLog(outputWithoutMetadata),
+        input: (overrideInputFilterSensitiveLog ?? inputFilterSensitiveLog)(args.input),
+        output: (overrideOutputFilterSensitiveLog ?? outputFilterSensitiveLog)(outputWithoutMetadata),
         metadata: $metadata,
       });
     }

Diff for: packages/types/src/middleware.ts

@@ -407,6 +407,14 @@ export interface HandlerExecutionContext {
    */
   authSchemes?: AuthScheme[];
 
+  /**
+   * Used by DynamoDbDocumentClient.
+   */
+  dynamoDbDocumentClientOptions?: Partial<{
+    overrideInputFilterSensitiveLog(...args: any[]): string | void;
+    overrideOutputFilterSensitiveLog(...args: any[]): string | void;
+  }>;
+
   [key: string]: any;
 }