From 3a2f617d15631da1616044d3ce5ef44178463475 Mon Sep 17 00:00:00 2001
From: Trivikram Kamat <16024985+trivikr@users.noreply.github.com>
Date: Thu, 29 Oct 2020 13:26:57 -0700
Subject: [PATCH] fix: log requestId, extendedRequestId, cfId in metadata
 (#1640)

---
 packages/middleware-logger/package.json       |  1 +
 .../src/loggerMiddleware.spec.ts              | 65 +++++++++++++++----
 .../middleware-logger/src/loggerMiddleware.ts | 32 +++++----
 3 files changed, 71 insertions(+), 27 deletions(-)

diff --git a/packages/middleware-logger/package.json b/packages/middleware-logger/package.json
index 614092122fa6e..86fae97bb1803 100644
--- a/packages/middleware-logger/package.json
+++ b/packages/middleware-logger/package.json
@@ -23,6 +23,7 @@
     "tslib": "^1.8.0"
   },
   "devDependencies": {
+    "@aws-sdk/protocol-http": "1.0.0-rc.3",
     "@types/jest": "^26.0.4",
     "@types/node": "^10.0.0",
     "jest": "^26.1.0",
diff --git a/packages/middleware-logger/src/loggerMiddleware.spec.ts b/packages/middleware-logger/src/loggerMiddleware.spec.ts
index 0dd617de2c2a4..749dc97f952a5 100644
--- a/packages/middleware-logger/src/loggerMiddleware.spec.ts
+++ b/packages/middleware-logger/src/loggerMiddleware.spec.ts
@@ -32,37 +32,40 @@ describe("loggerMiddleware", () => {
   };
 
   const mockResponse = {
-    output: {
-      $metadata: {
-        statusCode: 200,
-        requestId: "requestId",
+    response: {
+      statusCode: 200,
+      headers: {
+        "x-amzn-requestid": "requestId",
+        "x-amz-id-2": "extendedRequestId",
+        "x-amz-cf-id": "cfId",
       },
+    },
+    output: {
       outputKey: "outputValue",
     },
   };
 
-  beforeEach(() => {
-    mockNext.mockResolvedValueOnce(mockResponse);
-  });
-
   afterEach(() => {
     jest.clearAllMocks();
   });
 
   it("returns without logging if context.logger is not defined", async () => {
+    mockNext.mockResolvedValueOnce(mockResponse);
     const response = await loggerMiddleware()(mockNext, {})(mockArgs);
     expect(mockNext).toHaveBeenCalledTimes(1);
     expect(response).toStrictEqual(mockResponse);
   });
 
   it("returns without logging if context.logger doesn't have info function", async () => {
+    mockNext.mockResolvedValueOnce(mockResponse);
     const logger = {} as Logger;
     const response = await loggerMiddleware()(mockNext, { logger })(mockArgs);
     expect(mockNext).toHaveBeenCalledTimes(1);
     expect(response).toStrictEqual(mockResponse);
   });
 
-  it("logs $metadata if context.logger has info function", async () => {
+  it("logs metadata if context.logger has info function", async () => {
+    mockNext.mockResolvedValueOnce(mockResponse);
     const logger = ({ info: jest.fn() } as unknown) as Logger;
 
     const context = {
@@ -75,11 +78,47 @@ describe("loggerMiddleware", () => {
 
     expect(logger.info).toHaveBeenCalledTimes(1);
 
-    const {
-      output: { $metadata },
-    } = mockResponse;
     expect(logger.info).toHaveBeenCalledWith({
-      $metadata,
+      metadata: {
+        statusCode: mockResponse.response.statusCode,
+        requestId: mockResponse.response.headers["x-amzn-requestid"],
+        extendedRequestId: mockResponse.response.headers["x-amz-id-2"],
+        cfId: mockResponse.response.headers["x-amz-cf-id"],
+      },
+    });
+  });
+
+  it("logs header x-amzn-request-id as requestId if x-amzn-requestid is not present", async () => {
+    const requestIdBackup = "requestIdBackup";
+    const customResponse = {
+      ...mockResponse,
+      response: {
+        ...mockResponse.response,
+        headers: {
+          "x-amzn-request-id": requestIdBackup,
+        },
+      },
+    };
+    mockNext.mockResolvedValueOnce(customResponse);
+    const logger = ({ info: jest.fn() } as unknown) as Logger;
+
+    const context = {
+      logger,
+    };
+
+    const response = await loggerMiddleware()(mockNext, context)(mockArgs);
+    expect(mockNext).toHaveBeenCalledTimes(1);
+    expect(response).toStrictEqual(customResponse);
+
+    expect(logger.info).toHaveBeenCalledTimes(1);
+
+    expect(logger.info).toHaveBeenCalledWith({
+      metadata: {
+        statusCode: customResponse.response.statusCode,
+        requestId: requestIdBackup,
+        extendedRequestId: undefined,
+        cfId: undefined,
+      },
     });
   });
 });
diff --git a/packages/middleware-logger/src/loggerMiddleware.ts b/packages/middleware-logger/src/loggerMiddleware.ts
index 704abfbf273de..42c3b36f74e6b 100644
--- a/packages/middleware-logger/src/loggerMiddleware.ts
+++ b/packages/middleware-logger/src/loggerMiddleware.ts
@@ -1,18 +1,21 @@
+import { HttpResponse } from "@aws-sdk/protocol-http";
 import {
   AbsoluteLocation,
-  BuildHandler,
-  BuildHandlerArguments,
-  BuildHandlerOptions,
-  BuildHandlerOutput,
   HandlerExecutionContext,
+  InitializeHandler,
+  InitializeHandlerArguments,
+  InitializeHandlerOptions,
+  InitializeHandlerOutput,
   MetadataBearer,
   Pluggable,
 } from "@aws-sdk/types";
 
 export const loggerMiddleware = () => <Output extends MetadataBearer = MetadataBearer>(
-  next: BuildHandler<any, Output>,
+  next: InitializeHandler<any, Output>,
   context: HandlerExecutionContext
-): BuildHandler<any, Output> => async (args: BuildHandlerArguments<any>): Promise<BuildHandlerOutput<Output>> => {
+): InitializeHandler<any, Output> => async (
+  args: InitializeHandlerArguments<any>
+): Promise<InitializeHandlerOutput<Output>> => {
   const { logger } = context;
 
   const response = await next(args);
@@ -21,25 +24,26 @@ export const loggerMiddleware = () => <Output extends MetadataBearer = MetadataB
     return response;
   }
 
-  const {
-    output: { $metadata },
-  } = response;
+  const httpResponse = response.response as HttpResponse;
 
-  // TODO: Populate custom metadata in https://github.com/aws/aws-sdk-js-v3/issues/1491#issuecomment-692174256
-  // $metadata will be removed in https://github.com/aws/aws-sdk-js-v3/issues/1490
   if (typeof logger.info === "function") {
     logger.info({
-      $metadata,
+      metadata: {
+        statusCode: httpResponse.statusCode,
+        requestId: httpResponse.headers["x-amzn-requestid"] ?? httpResponse.headers["x-amzn-request-id"],
+        extendedRequestId: httpResponse.headers["x-amz-id-2"],
+        cfId: httpResponse.headers["x-amz-cf-id"],
+      },
     });
   }
 
   return response;
 };
 
-export const loggerMiddlewareOptions: BuildHandlerOptions & AbsoluteLocation = {
+export const loggerMiddlewareOptions: InitializeHandlerOptions & AbsoluteLocation = {
   name: "loggerMiddleware",
   tags: ["LOGGER"],
-  step: "build",
+  step: "initialize",
 };
 
 // eslint-disable-next-line @typescript-eslint/no-unused-vars