fix: stop reading private aws-cdk-lib Function.environment field

README.md

@@ -409,6 +409,10 @@ When `addLambdaFunctions` is called, the Datadog CDK construct grants your Lambd
 
 The DatadogLambda construct takes in a list of lambda functions and installs the Datadog Lambda Library by attaching the Lambda Layers for [.NET][19], [Java][15], [Node.js][2], and [Python][1] to your functions. It redirects to a replacement handler that initializes the Lambda Library without any required code changes. Additional configurations added to the Datadog CDK construct will also translate into their respective environment variables under each lambda function (if applicable / required).
 
+**Configuring `DD_*` environment variables:** Set Datadog environment variables through `DatadogLambdaProps` (e.g. `enableDatadogTracing`, `logLevel`, `tags`). If you need to set a `DD_*` variable directly on a function, call `func.addEnvironment()` **after** `addLambdaFunctions()` -- values set before will be overridden by the construct.
+
+To set a value the construct should respect (rather than override), use `datadogLambda.setEnvironment(func, key, value)` **before** `addLambdaFunctions()`. This is useful for setting `DD_TAGS` per function: when source code integration is enabled, the construct appends git metadata (`git.commit.sha`, `git.repository_url`) to the value you set instead of overriding it.
+
 While Lambda function based log groups are handled by the `addLambdaFunctions` method automatically, the construct has an additional function `addForwarderToNonLambdaLogGroups` which subscribes the forwarder to any additional log groups of your choosing.
 
 ## Step Functions

src/datadog-lambda.ts

@@ -13,6 +13,7 @@ import * as logs from "aws-cdk-lib/aws-logs";
 import { ISecret, Secret } from "aws-cdk-lib/aws-secretsmanager";
 import { Construct } from "constructs";
 import log from "loglevel";
+import { getTrackedEnv, setTrackedEnv } from "./env-tracker";
 import {
   applyLayers,
   redirectHandlers,
@@ -181,6 +182,20 @@ export class DatadogLambda extends Construct {
     }
   }
 
+  /**
+   * Sets an environment variable on the given Lambda function and records it in the
+   * construct's internal tracker, so the construct treats it as a value it manages.
+   *
+   * Call this before `addLambdaFunctions()` to seed values the construct will respect.
+   * The main use case is setting `DD_TAGS` per function: when source code integration is
+   * enabled, the construct appends git metadata (`git.commit.sha`, `git.repository_url`)
+   * to the tracked `DD_TAGS` rather than overriding it.
+   */
+  public setEnvironment(lambdaFunction: LambdaFunction, key: string, value: string): void {
+    const [extractedLambdaFunction] = extractSingletonFunctions([lambdaFunction]);
+    setTrackedEnv(extractedLambdaFunction, key, value);
+  }
+
   public overrideGitMetadata(gitCommitSha: string, gitRepoUrl?: string): void {
     if (gitCommitSha) {
       this.gitCommitShaOverride = gitCommitSha;
@@ -191,23 +206,21 @@ export class DatadogLambda extends Construct {
 
     // If any lambdas have already been added, override the commit sha and url
     if (this.lambdas) {
-      this.lambdas.forEach((lambdaFunction: any) => {
-        const existingTags = lambdaFunction.environment.map.get(DD_TAGS);
-        if (existingTags === undefined) {
+      this.lambdas.forEach((lambdaFunction: LambdaFunction) => {
+        const existingTagValue = getTrackedEnv(lambdaFunction, DD_TAGS);
+        if (existingTagValue === undefined) {
           return;
         }
-        const tags = existingTags.value.split(",");
+        const tags = existingTagValue.split(",");
         if (gitCommitSha) {
-          const index = tags.findIndex((val: string) => val.split(":")[0] === "git.commit.sha");
-          tags[index] = `git.commit.sha:${gitCommitSha}`;
+          upsertTag(tags, "git.commit.sha", gitCommitSha);
         }
 
         if (gitRepoUrl) {
-          const index = tags.findIndex((val: string) => val.split(":")[0] === "git.repository_url");
-          tags[index] = `git.repository_url:${gitRepoUrl}`;
+          upsertTag(tags, "git.repository_url", gitRepoUrl);
         }
 
-        lambdaFunction.addEnvironment(DD_TAGS, tags.join(","));
+        setTrackedEnv(lambdaFunction, DD_TAGS, tags.join(","));
       });
     }
   }
@@ -305,6 +318,17 @@ function isSingletonFunction(fn: LambdaFunction): fn is lambda.SingletonFunction
   return fn.hasOwnProperty("lambdaFunction");
 }
 
+// Replaces the value of an existing `key:...` tag in place, or appends it if not present.
+function upsertTag(tags: string[], key: string, value: string): void {
+  const index = tags.findIndex((tag) => tag.split(":")[0] === key);
+  const entry = `${key}:${value}`;
+  if (index === -1) {
+    tags.push(entry);
+  } else {
+    tags[index] = entry;
+  }
+}
+
 export function validateProps(props: DatadogLambdaProps, apiKeyArnOverride = false): void {
   log.debug("Validating props...");
 

src/env-tracker.ts

@@ -0,0 +1,43 @@
+/*
+ * Unless explicitly stated otherwise all files in this repository are licensed
+ * under the Apache License Version 2.0.
+ *
+ * This product includes software developed at Datadog (https://www.datadoghq.com/).
+ * Copyright 2021 Datadog, Inc.
+ */
+
+import { LambdaFunction } from "./interfaces";
+
+// Tracks env vars written by this library per Function, so we can read them back without
+// touching aws-cdk-lib's private Function.environment field. aws-cdk-lib has no public
+// read accessor for env vars, so the WeakMap is the library's own bookkeeping only.
+//
+// This module is intentionally not re-exported from index.ts: the helpers are internal
+// and not part of the package's public API.
+//
+// WeakMap (vs Map) so Function objects can be garbage-collected when a stack goes out of
+// scope (e.g. between test cases) without this module holding a lingering reference.
+//
+// Note: env vars set on a Function via func.addEnvironment() outside of this library
+// are not tracked here and will be overridden if the library sets the same key.
+// Configure DD_* vars via DatadogLambdaProps, or call func.addEnvironment() after
+// calling addLambdaFunctions().
+const ddEnvTracker: WeakMap<LambdaFunction, Map<string, string>> = new WeakMap();
+
+export function setTrackedEnv(lam: LambdaFunction, key: string, value: string): void {
+  let envMap = ddEnvTracker.get(lam);
+  if (!envMap) {
+    envMap = new Map();
+    ddEnvTracker.set(lam, envMap);
+  }
+  envMap.set(key, value);
+  lam.addEnvironment(key, value);
+}
+
+export function getTrackedEnv(lam: LambdaFunction, key: string): string | undefined {
+  return ddEnvTracker.get(lam)?.get(key);
+}
+
+export function hasTrackedEnv(lam: LambdaFunction, key: string): boolean {
+  return ddEnvTracker.get(lam)?.has(key) ?? false;
+}

src/env.ts

@@ -6,10 +6,10 @@
  * Copyright 2021 Datadog, Inc.
  */
 
-import * as lambda from "aws-cdk-lib/aws-lambda";
 import log from "loglevel";
 import { runtimeLookup, RuntimeType } from "./constants";
-import { DatadogLambdaProps, DatadogLambdaStrictProps } from "./interfaces";
+import { setTrackedEnv, getTrackedEnv, hasTrackedEnv } from "./env-tracker";
+import { DatadogLambdaProps, DatadogLambdaStrictProps, LambdaFunction } from "./interfaces";
 
 export const AWS_LAMBDA_EXEC_WRAPPER_KEY = "AWS_LAMBDA_EXEC_WRAPPER";
 export const AWS_LAMBDA_EXEC_WRAPPER_VAL = "/opt/datadog_wrapper";
@@ -44,7 +44,7 @@ const execSync = require("child_process").execSync;
 const URL = require("url").URL;
 
 export function setGitEnvironmentVariables(
-  lambdas: any[],
+  lambdas: LambdaFunction[],
   gitCommitShaOverride?: string | undefined,
   gitRepoUrlOverride?: string | undefined,
 ): void {
@@ -61,12 +61,11 @@ export function setGitEnvironmentVariables(
 
   if (hash == "" || gitRepoUrl == "") return;
 
-  // We're using an any type here because AWS does not expose the `environment` field in their type
+  const tagsValue = `git.commit.sha:${hash},git.repository_url:${gitRepoUrl}`;
   lambdas.forEach((lam) => {
-    const tagsValue = `git.commit.sha:${hash},git.repository_url:${gitRepoUrl}`;
-    const existingTagValue = lam.environment.map.get(DD_TAGS)?.value;
+    const existingTagValue = getTrackedEnv(lam, DD_TAGS);
     const finalTagValue = existingTagValue ? `${existingTagValue},${tagsValue}` : tagsValue;
-    lam.addEnvironment(DD_TAGS, finalTagValue);
+    setTrackedEnv(lam, DD_TAGS, finalTagValue);
   });
 }
 
@@ -117,17 +116,16 @@ function filterSensitiveInfoFromRepository(repositoryUrl: string): string {
   }
 }
 
-export function applyEnvVariables(lam: lambda.Function, baseProps: DatadogLambdaStrictProps): void {
+export function applyEnvVariables(lam: LambdaFunction, baseProps: DatadogLambdaStrictProps): void {
   log.debug(`Setting environment variables...`);
-  const lam_with_env_vars: any = lam; //cast to any to access the private environment fields like in setGitEnvironmentVariables
-  const setEnvIfUndefined = (envVar: string, value: string | boolean) => {
-    if (lam_with_env_vars.environment.map.get(envVar) === undefined) {
-      lam.addEnvironment(envVar, value.toString().toLowerCase());
+
+  const setEnvIfUnset = (envVar: string, value: string | boolean) => {
+    if (!hasTrackedEnv(lam, envVar)) {
+      setTrackedEnv(lam, envVar, value.toString().toLowerCase());
     }
   };
 
-  //for each env variable, only set to default if it is NOT already set by user
-  setEnvIfUndefined(ENABLE_DD_TRACING_ENV_VAR, baseProps.enableDatadogTracing);
+  setEnvIfUnset(ENABLE_DD_TRACING_ENV_VAR, baseProps.enableDatadogTracing);
 
   const runtimeType = runtimeLookup[lam.runtime.name];
 
@@ -144,75 +142,74 @@ export function applyEnvVariables(lam: lambda.Function, baseProps: DatadogLambda
       baseProps.pythonLayerVersion >= 114) ||
     baseProps.datadogAppSecMode === "tracer"
   ) {
-    setEnvIfUndefined(ENABLE_DD_APPSEC_ENV_VAR, "true");
+    setEnvIfUnset(ENABLE_DD_APPSEC_ENV_VAR, "true");
   } else if (baseProps.datadogAppSecMode === "on" || baseProps.datadogAppSecMode === "extension") {
-    setEnvIfUndefined(AWS_LAMBDA_EXEC_WRAPPER_KEY, AWS_LAMBDA_EXEC_WRAPPER_VAL);
-    setEnvIfUndefined(ENABLE_DD_SERVERLESS_APPSEC_ENV_VAR, "true");
+    setEnvIfUnset(AWS_LAMBDA_EXEC_WRAPPER_KEY, AWS_LAMBDA_EXEC_WRAPPER_VAL);
+    setEnvIfUnset(ENABLE_DD_SERVERLESS_APPSEC_ENV_VAR, "true");
   }
 
-  setEnvIfUndefined(ENABLE_XRAY_TRACE_MERGING_ENV_VAR, baseProps.enableMergeXrayTraces);
+  setEnvIfUnset(ENABLE_XRAY_TRACE_MERGING_ENV_VAR, baseProps.enableMergeXrayTraces);
 
   if (baseProps.extensionLayerVersion || baseProps.extensionLayerArn) {
-    setEnvIfUndefined(INJECT_LOG_CONTEXT_ENV_VAR, "false");
+    setEnvIfUnset(INJECT_LOG_CONTEXT_ENV_VAR, "false");
   } else {
-    setEnvIfUndefined(INJECT_LOG_CONTEXT_ENV_VAR, baseProps.injectLogContext);
+    setEnvIfUnset(INJECT_LOG_CONTEXT_ENV_VAR, baseProps.injectLogContext);
   }
 
-  setEnvIfUndefined(ENABLE_DD_LOGS_ENV_VAR, baseProps.enableDatadogLogs);
-  setEnvIfUndefined(CAPTURE_LAMBDA_PAYLOAD_ENV_VAR, baseProps.captureLambdaPayload);
+  setEnvIfUnset(ENABLE_DD_LOGS_ENV_VAR, baseProps.enableDatadogLogs);
+  setEnvIfUnset(CAPTURE_LAMBDA_PAYLOAD_ENV_VAR, baseProps.captureLambdaPayload);
 
-  //Cloud Payload Tagging - set to "all" when enabled, empty string when disabled
   const CLOUD_PAYLOAD_TAGGING_VALUE = baseProps.captureCloudServicePayload ? "all" : "";
-  setEnvIfUndefined(DD_TRACE_CLOUD_REQUEST_PAYLOAD_TAGGING, CLOUD_PAYLOAD_TAGGING_VALUE);
-  setEnvIfUndefined(DD_TRACE_CLOUD_RESPONSE_PAYLOAD_TAGGING, CLOUD_PAYLOAD_TAGGING_VALUE);
+  setEnvIfUnset(DD_TRACE_CLOUD_REQUEST_PAYLOAD_TAGGING, CLOUD_PAYLOAD_TAGGING_VALUE);
+  setEnvIfUnset(DD_TRACE_CLOUD_RESPONSE_PAYLOAD_TAGGING, CLOUD_PAYLOAD_TAGGING_VALUE);
 
   if (baseProps.logLevel) {
-    setEnvIfUndefined(LOG_LEVEL_ENV_VAR, baseProps.logLevel);
+    setEnvIfUnset(LOG_LEVEL_ENV_VAR, baseProps.logLevel);
   }
 }
 
-export function setDDEnvVariables(lam: lambda.Function, props: DatadogLambdaProps): void {
+export function setDDEnvVariables(lam: LambdaFunction, props: DatadogLambdaProps): void {
   if (props.env) {
-    lam.addEnvironment(DD_ENV_ENV_VAR, props.env);
+    setTrackedEnv(lam, DD_ENV_ENV_VAR, props.env);
   }
   if (props.service) {
-    lam.addEnvironment(DD_SERVICE_ENV_VAR, props.service);
+    setTrackedEnv(lam, DD_SERVICE_ENV_VAR, props.service);
   }
   if (props.version) {
-    lam.addEnvironment(DD_VERSION_ENV_VAR, props.version);
+    setTrackedEnv(lam, DD_VERSION_ENV_VAR, props.version);
   }
   if (props.tags) {
-    lam.addEnvironment(DD_TAGS, props.tags);
+    setTrackedEnv(lam, DD_TAGS, props.tags);
   }
   if (props.enableColdStartTracing !== undefined) {
-    lam.addEnvironment(DD_COLD_START_TRACING, props.enableColdStartTracing.toString().toLowerCase());
+    setTrackedEnv(lam, DD_COLD_START_TRACING, props.enableColdStartTracing.toString().toLowerCase());
   }
   if (props.minColdStartTraceDuration !== undefined) {
-    lam.addEnvironment(DD_MIN_COLD_START_DURATION, props.minColdStartTraceDuration.toString().toLowerCase());
+    setTrackedEnv(lam, DD_MIN_COLD_START_DURATION, props.minColdStartTraceDuration.toString().toLowerCase());
   }
   if (props.coldStartTraceSkipLibs !== undefined) {
-    lam.addEnvironment(DD_COLD_START_TRACE_SKIP_LIB, props.coldStartTraceSkipLibs);
+    setTrackedEnv(lam, DD_COLD_START_TRACE_SKIP_LIB, props.coldStartTraceSkipLibs);
   }
   if (props.enableProfiling !== undefined) {
-    lam.addEnvironment(DD_PROFILING_ENABLED, props.enableProfiling.toString().toLowerCase());
+    setTrackedEnv(lam, DD_PROFILING_ENABLED, props.enableProfiling.toString().toLowerCase());
   }
   if (props.encodeAuthorizerContext !== undefined) {
-    lam.addEnvironment(DD_ENCODE_AUTHORIZER_CONTEXT, props.encodeAuthorizerContext.toString().toLowerCase());
+    setTrackedEnv(lam, DD_ENCODE_AUTHORIZER_CONTEXT, props.encodeAuthorizerContext.toString().toLowerCase());
   }
   if (props.decodeAuthorizerContext !== undefined) {
-    lam.addEnvironment(DD_DECODE_AUTHORIZER_CONTEXT, props.decodeAuthorizerContext.toString().toLowerCase());
+    setTrackedEnv(lam, DD_DECODE_AUTHORIZER_CONTEXT, props.decodeAuthorizerContext.toString().toLowerCase());
   }
   if (props.apmFlushDeadline !== undefined) {
-    lam.addEnvironment(DD_APM_FLUSH_DEADLINE_MILLISECONDS, props.apmFlushDeadline.toString().toLowerCase());
+    setTrackedEnv(lam, DD_APM_FLUSH_DEADLINE_MILLISECONDS, props.apmFlushDeadline.toString().toLowerCase());
   }
 
   if (props.llmObsEnabled !== undefined) {
-    lam.addEnvironment(DD_LLMOBS_ENABLED, props.llmObsEnabled.toString().toLowerCase());
+    setTrackedEnv(lam, DD_LLMOBS_ENABLED, props.llmObsEnabled.toString().toLowerCase());
   }
   if (props.llmObsMlApp !== undefined) {
-    lam.addEnvironment(DD_LLMOBS_ML_APP, props.llmObsMlApp);
+    setTrackedEnv(lam, DD_LLMOBS_ML_APP, props.llmObsMlApp);
   }
   if (props.llmObsAgentlessEnabled !== undefined) {
-    lam.addEnvironment(DD_LLMOBS_AGENTLESS_ENABLED, props.llmObsAgentlessEnabled.toString().toLowerCase());
+    setTrackedEnv(lam, DD_LLMOBS_AGENTLESS_ENABLED, props.llmObsAgentlessEnabled.toString().toLowerCase());
   }
 }