Optionally capture Lambda request and response payloads (#222)

* feat: Allow users to capture and tag request and response payloads.

* fix: remove commented console

* feat: switch to funciton.request and function.response

* feat: Max depth 5 for tagging payloads

* feat: Update env var name after product feedback. Update filename to be consistent

* fix: Vscode will update imports but not save the file if it's open, nice

* fix: Fix renamed import

* feat: set maxDepth to 10

Author: astuyve

.prettierrc

@@ -1,5 +1,6 @@
 {
   "printWidth": 120,
   "trailingComma": "all",
-  "arrowParens": "always"
+  "arrowParens": "always",
+  "semi": true
 }

src/index.ts

@@ -17,12 +17,14 @@ import {
   setLogger,
   promisifiedHandler,
   logDebug,
+  tagObject,
 } from "./utils";
 export { TraceHeaders } from "./trace";
 import { extractHTTPStatusCodeTag } from "./trace/trigger";
 
 export const apiKeyEnvVar = "DD_API_KEY";
 export const apiKeyKMSEnvVar = "DD_KMS_API_KEY";
+export const captureLambdaPayloadEnvVar = "DD_CAPTURE_LAMBDA_PAYLOAD";
 export const siteURLEnvVar = "DD_SITE";
 export const logLevelEnvVar = "DD_LOG_LEVEL";
 export const logForwardingEnvVar = "DD_FLUSH_TO_LOG";
@@ -60,6 +62,7 @@ export const defaultConfig: Config = {
   apiKey: "",
   apiKeyKMS: "",
   autoPatchHTTP: true,
+  captureLambdaPayload: false,
   debugLogging: false,
   enhancedMetrics: true,
   forceWrap: false,
@@ -139,6 +142,10 @@ export function datadog<TEvent, TResult>(
               // Store the status tag in the listener to send to Xray on invocation completion
               traceListener.triggerTags["http.status_code"] = statusCode;
               if (traceListener.currentSpan) {
+                if (finalConfig.captureLambdaPayload) {
+                  tagObject(traceListener.currentSpan, "function.request", localEvent);
+                  tagObject(traceListener.currentSpan, "function.response", localResult);
+                }
                 traceListener.currentSpan.setTag("http.status_code", statusCode);
               }
             }
@@ -176,7 +183,7 @@ export function datadog<TEvent, TResult>(
  * Sends a Distribution metric asynchronously to the Datadog API.
  * @param name The name of the metric to send.
  * @param value The value of the metric
- * @param metricTime The timesamp associated with this metric data point.
+ * @param metricTime The timestamp associated with this metric data point.
  * @param tags The tags associated with the metric. Should be of the format "tag:value".
  */
 export function sendDistributionMetricWithDate(name: string, value: number, metricTime: Date, ...tags: string[]) {
@@ -260,6 +267,11 @@ function getConfig(userConfig?: Partial<Config>): Config {
     config.mergeDatadogXrayTraces = result === "true";
   }
 
+  if (userConfig === undefined || userConfig.captureLambdaPayload === undefined) {
+    const result = getEnvValue(captureLambdaPayloadEnvVar, "false").toLocaleLowerCase();
+    config.captureLambdaPayload = result === "true";
+  }
+
   return config;
 }
 

src/trace/listener.spec.ts

@@ -45,7 +45,12 @@ jest.mock("./trace-context-service", () => {
 });
 
 describe("TraceListener", () => {
-  const defaultConfig = { autoPatchHTTP: true, mergeDatadogXrayTraces: false, injectLogContext: false };
+  const defaultConfig = {
+    autoPatchHTTP: true,
+    captureLambdaPayload: false,
+    mergeDatadogXrayTraces: false,
+    injectLogContext: false,
+  };
   const context = {
     invokedFunctionArn: "arn:aws:lambda:us-east-1:123456789101:function:my-lambda",
     awsRequestId: "1234",

src/trace/listener.ts

@@ -26,6 +26,10 @@ export interface TraceConfig {
    * @default true.
    */
   autoPatchHTTP: boolean;
+  /**
+   * Whether to capture the lambda payload and response in Datadog.
+   */
+  captureLambdaPayload: boolean;
   /**
    * Whether to automatically patch console.log with Datadog's tracing ids.
    */

src/utils/index.ts

@@ -3,3 +3,4 @@ export { wrap, promisifiedHandler } from "./handler";
 export { Timer } from "./timer";
 export { logError, logDebug, Logger, setLogLevel, setLogger, LogLevel } from "./log";
 export { get, post } from "./request";
+export { tagObject } from "./tag-object";

src/utils/tag-object.spec.ts

@@ -0,0 +1,60 @@
+import { tagObject } from "./tag-object";
+
+describe("tagObject", () => {
+  const setTag = jest.fn();
+
+  beforeEach(() => {
+    jest.clearAllMocks();
+  });
+  it("tags something simple", () => {
+    const span = {
+      setTag,
+    };
+    tagObject(span, "lambda_payload", { request: { myKey: "myValue" } });
+    expect(setTag).toBeCalledWith("lambda_payload.request.myKey", "myValue");
+  });
+  it("tags complex objects", () => {
+    const span = {
+      setTag,
+    };
+    tagObject(span, "lambda_payload", {
+      request: {
+        keyOne: "foobar",
+        myObject: {
+          anotherKey: ["array", "of", "values"],
+        },
+        val: null,
+        number: 1,
+      },
+    });
+    expect(setTag.mock.calls).toEqual([
+      ["lambda_payload.request.keyOne", "foobar"],
+      ["lambda_payload.request.myObject.anotherKey.0", "array"],
+      ["lambda_payload.request.myObject.anotherKey.1", "of"],
+      ["lambda_payload.request.myObject.anotherKey.2", "values"],
+      ["lambda_payload.request.val", null],
+      ["lambda_payload.request.number", 1],
+    ]);
+  });
+  it("tags arrays of objects", () => {
+    const span = {
+      setTag,
+    };
+    tagObject(span, "lambda_payload", {
+      request: {
+        vals: [{ thingOne: 1 }, { thingTwo: 2 }],
+      },
+    });
+    expect(setTag.mock.calls).toEqual([
+      ["lambda_payload.request.vals.0.thingOne", 1],
+      ["lambda_payload.request.vals.1.thingTwo", 2],
+    ]);
+  });
+  it("redacts common secret keys", () => {
+    const span = {
+      setTag,
+    };
+    tagObject(span, "lambda_payload", { request: { headers: { authorization: "myValue" } } });
+    expect(setTag).toBeCalledWith("lambda_payload.request.headers.authorization", "redacted");
+  });
+});

src/utils/tag-object.ts

@@ -0,0 +1,40 @@
+const redactableKeys = ["authorization", "x-authorization", "password", "token"];
+const maxDepth = 10;
+
+export function tagObject(currentSpan: any, key: string, obj: any, depth = 0): any {
+  if (depth >= maxDepth) {
+    return;
+  } else {
+    depth += 1;
+  }
+  if (obj === null) {
+    return currentSpan.setTag(key, obj);
+  }
+  if (typeof obj === "string") {
+    let parsed: string;
+    try {
+      parsed = JSON.parse(obj);
+    } catch (e) {
+      const redacted = redactVal(key, obj.substring(0, 5000));
+      return currentSpan.setTag(key, redacted);
+    }
+    return tagObject(currentSpan, key, parsed, depth);
+  }
+  if (typeof obj === "number") {
+    return currentSpan.setTag(key, obj);
+  }
+  if (typeof obj === "object") {
+    for (const [k, v] of Object.entries(obj)) {
+      tagObject(currentSpan, `${key}.${k}`, v, depth);
+    }
+    return;
+  }
+}
+
+function redactVal(k: string, v: string): string {
+  const splitKey = k.split(".").pop() || k;
+  if (redactableKeys.includes(splitKey)) {
+    return "redacted";
+  }
+  return v;
+}