fix(lambda-nodejs): NodejsFunction construct incompatible with lambda@edge

packages/@aws-cdk/aws-cloudfront/lib/private/cache-behavior.ts

@@ -48,15 +48,10 @@ export class CacheBehavior {
       smoothStreaming: this.props.smoothStreaming,
       viewerProtocolPolicy: this.props.viewerProtocolPolicy ?? ViewerProtocolPolicy.ALLOW_ALL,
       lambdaFunctionAssociations: this.props.edgeLambdas
-        ? this.props.edgeLambdas.map(edgeLambda => {
-          if (edgeLambda.functionVersion.version === '$LATEST') {
-            throw new Error('$LATEST function version cannot be used for Lambda@Edge');
-          }
-          return {
-            lambdaFunctionArn: edgeLambda.functionVersion.functionArn,
-            eventType: edgeLambda.eventType.toString(),
-          };
-        })
+        ? this.props.edgeLambdas.map(edgeLambda => ({
+          lambdaFunctionArn: edgeLambda.functionVersion.edgeArn,
+          eventType: edgeLambda.eventType.toString(),
+        }))
         : undefined,
     };
   }

packages/@aws-cdk/aws-cloudfront/lib/web_distribution.ts

@@ -896,7 +896,7 @@ export class CloudFrontWebDistribution extends cdk.Resource implements IDistribu
         lambdaFunctionAssociations: input.lambdaFunctionAssociations
           .map(fna => ({
             eventType: fna.eventType,
-            lambdaFunctionArn: fna.lambdaFunction && fna.lambdaFunction.functionArn,
+            lambdaFunctionArn: fna.lambdaFunction && fna.lambdaFunction.edgeArn,
           })),
       });
 

packages/@aws-cdk/aws-cloudfront/test/distribution.test.ts

@@ -538,6 +538,59 @@ describe('with Lambda@Edge functions', () => {
       });
     }).toThrow(/\$LATEST function version cannot be used for Lambda@Edge/);
   });
+
+  test('with removable env vars', () => {
+    const envLambdaFunction = new lambda.Function(stack, 'EnvFunction', {
+      runtime: lambda.Runtime.NODEJS,
+      code: lambda.Code.fromInline('whateverwithenv'),
+      handler: 'index.handler',
+    });
+    envLambdaFunction.addEnvironment('KEY', 'value', { removeInEdge: true });
+
+    new Distribution(stack, 'MyDist', {
+      defaultBehavior: {
+        origin,
+        edgeLambdas: [
+          {
+            functionVersion: envLambdaFunction.currentVersion,
+            eventType: LambdaEdgeEventType.ORIGIN_REQUEST,
+          },
+        ],
+      },
+    });
+
+    expect(stack).toHaveResource('AWS::Lambda::Function', {
+      Environment: ABSENT,
+      Code: {
+        ZipFile: 'whateverwithenv',
+      },
+    });
+  });
+
+  test('with incompatible env vars', () => {
+    const envLambdaFunction = new lambda.Function(stack, 'EnvFunction', {
+      runtime: lambda.Runtime.NODEJS,
+      code: lambda.Code.fromInline('whateverwithenv'),
+      handler: 'index.handler',
+      environment: {
+        KEY: 'value',
+      },
+    });
+
+    new Distribution(stack, 'MyDist', {
+      defaultBehavior: {
+        origin,
+        edgeLambdas: [
+          {
+            functionVersion: envLambdaFunction.currentVersion,
+            eventType: LambdaEdgeEventType.ORIGIN_REQUEST,
+          },
+        ],
+      },
+    });
+
+    expect(() => app.synth()).toThrow(/KEY/);
+  });
 });
 
 test('price class is included if provided', () => {

packages/@aws-cdk/aws-cloudfront/test/web_distribution.test.ts

@@ -1,4 +1,4 @@
-import { expect, haveResource, haveResourceLike } from '@aws-cdk/assert';
+import { ABSENT, expect, haveResource, haveResourceLike } from '@aws-cdk/assert';
 import * as certificatemanager from '@aws-cdk/aws-certificatemanager';
 import * as lambda from '@aws-cdk/aws-lambda';
 import * as s3 from '@aws-cdk/aws-s3';
@@ -426,8 +426,7 @@ nodeunitShim({
     const stack = new cdk.Stack();
     const sourceBucket = new s3.Bucket(stack, 'Bucket');
 
-    const lambdaFunction = new lambda.SingletonFunction(stack, 'Lambda', {
-      uuid: 'xxxx-xxxx-xxxx-xxxx',
+    const lambdaFunction = new lambda.Function(stack, 'Lambda', {
       code: lambda.Code.inline('foo'),
       handler: 'index.handler',
       runtime: lambda.Runtime.NODEJS_10_X,
@@ -444,7 +443,7 @@ nodeunitShim({
               isDefaultBehavior: true,
               lambdaFunctionAssociations: [{
                 eventType: LambdaEdgeEventType.ORIGIN_REQUEST,
-                lambdaFunction: lambdaFunction.latestVersion,
+                lambdaFunction: lambdaFunction.addVersion('1'),
               }],
             },
           ],
@@ -459,13 +458,7 @@ nodeunitShim({
             {
               'EventType': 'origin-request',
               'LambdaFunctionARN': {
-                'Fn::Join': [
-                  '',
-                  [
-                    { 'Fn::GetAtt': ['SingletonLambdaxxxxxxxxxxxxxxxx69D4268A', 'Arn'] },
-                    ':$LATEST',
-                  ],
-                ],
+                'Ref': 'LambdaVersion1BB7548E1',
               },
             },
           ],
@@ -476,6 +469,82 @@ nodeunitShim({
     test.done();
   },
 
+  'associate a lambda with removable env vars'(test: Test) {
+    const app = new cdk.App();
+    const stack = new cdk.Stack(app, 'Stack');
+    const sourceBucket = new s3.Bucket(stack, 'Bucket');
+
+    const lambdaFunction = new lambda.Function(stack, 'Lambda', {
+      code: lambda.Code.inline('foo'),
+      handler: 'index.handler',
+      runtime: lambda.Runtime.NODEJS_10_X,
+    });
+    lambdaFunction.addEnvironment('KEY', 'value', { removeInEdge: true });
+
+    new CloudFrontWebDistribution(stack, 'AnAmazingWebsiteProbably', {
+      originConfigs: [
+        {
+          s3OriginSource: {
+            s3BucketSource: sourceBucket,
+          },
+          behaviors: [
+            {
+              isDefaultBehavior: true,
+              lambdaFunctionAssociations: [{
+                eventType: LambdaEdgeEventType.ORIGIN_REQUEST,
+                lambdaFunction: lambdaFunction.addVersion('1'),
+              }],
+            },
+          ],
+        },
+      ],
+    });
+
+    expect(stack).to(haveResource('AWS::Lambda::Function', {
+      Environment: ABSENT,
+    }));
+
+    test.done();
+  },
+
+  'throws when associating a lambda with incompatible env vars'(test: Test) {
+    const app = new cdk.App();
+    const stack = new cdk.Stack(app, 'Stack');
+    const sourceBucket = new s3.Bucket(stack, 'Bucket');
+
+    const lambdaFunction = new lambda.Function(stack, 'Lambda', {
+      code: lambda.Code.inline('foo'),
+      handler: 'index.handler',
+      runtime: lambda.Runtime.NODEJS_10_X,
+      environment: {
+        KEY: 'value',
+      },
+    });
+
+    new CloudFrontWebDistribution(stack, 'AnAmazingWebsiteProbably', {
+      originConfigs: [
+        {
+          s3OriginSource: {
+            s3BucketSource: sourceBucket,
+          },
+          behaviors: [
+            {
+              isDefaultBehavior: true,
+              lambdaFunctionAssociations: [{
+                eventType: LambdaEdgeEventType.ORIGIN_REQUEST,
+                lambdaFunction: lambdaFunction.addVersion('1'),
+              }],
+            },
+          ],
+        },
+      ],
+    });
+
+    test.throws(() => app.synth(), /KEY/);
+
+    test.done();
+  },
+
   'distribution has a defaultChild'(test: Test) {
     const stack = new cdk.Stack();
     const sourceBucket = new s3.Bucket(stack, 'Bucket');

packages/@aws-cdk/aws-lambda-nodejs/lib/function.ts

@@ -85,7 +85,7 @@ export class NodejsFunction extends lambda.Function {
 
       // Enable connection reuse for aws-sdk
       if (props.awsSdkConnectionReuse ?? true) {
-        this.addEnvironment('AWS_NODEJS_CONNECTION_REUSE_ENABLED', '1');
+        this.addEnvironment('AWS_NODEJS_CONNECTION_REUSE_ENABLED', '1', { removeInEdge: true });
       }
     } finally {
       // We can only restore after the code has been bound to the function

packages/@aws-cdk/aws-lambda/lib/function-base.ts

@@ -318,6 +318,15 @@ export abstract class FunctionBase extends Resource implements IFunction {
     });
   }
 
+  /**
+   * Checks whether this function is compatible for Lambda@Edge.
+   *
+   * @internal
+   */
+  public _checkEdgeCompatibility(): void {
+    return;
+  }
+
   /**
    * Returns the construct tree node that corresponds to the lambda function.
    * For use internally for constructs, when the tree is set up in non-standard ways. Ex: SingletonFunction.
@@ -417,4 +426,8 @@ class LatestVersion extends FunctionBase implements IVersion {
   public addAlias(aliasName: string, options: AliasOptions = {}) {
     return addAlias(this, this, aliasName, options);
   }
+
+  public get edgeArn(): never {
+    throw new Error('$LATEST function version cannot be used for Lambda@Edge');
+  }
 }

packages/@aws-cdk/aws-lambda/lib/function.ts

@@ -515,7 +515,7 @@ export class Function extends FunctionBase {
   /**
    * Environment variables for this function
    */
-  private readonly environment: { [key: string]: string };
+  private environment: { [key: string]: EnvironmentConfig } = {};
 
   private readonly currentVersionOptions?: VersionOptions;
   private _currentVersion?: Version;
@@ -558,7 +558,7 @@ export class Function extends FunctionBase {
     const code = props.code.bind(this);
     verifyCodeConfig(code, props.runtime);
 
-    let profilingGroupEnvironmentVariables = {};
+    let profilingGroupEnvironmentVariables: { [key: string]: string } = {};
     if (props.profilingGroup && props.profiling !== false) {
       this.validateProfilingEnvironmentVariables(props);
       props.profilingGroup.grantPublish(this.role);
@@ -582,7 +582,10 @@ export class Function extends FunctionBase {
       };
     }
 
-    this.environment = { ...profilingGroupEnvironmentVariables, ...(props.environment || {}) };
+    const env = { ...profilingGroupEnvironmentVariables, ...props.environment };
+    for (const [key, value] of Object.entries(env)) {
+      this.addEnvironment(key, value);
+    }
 
     this.deadLetterQueue = this.buildDeadLetterQueue(props);
 
@@ -675,9 +678,10 @@ export class Function extends FunctionBase {
    * If this is a ref to a Lambda function, this operation results in a no-op.
    * @param key The environment variable key.
    * @param value The environment variable's value.
+   * @param options Environment variable options.
    */
-  public addEnvironment(key: string, value: string): this {
-    this.environment[key] = value;
+  public addEnvironment(key: string, value: string, options?: EnvironmentOptions): this {
+    this.environment[key] = { value, ...options };
     return this;
   }
 
@@ -764,27 +768,43 @@ export class Function extends FunctionBase {
     return this._logGroup;
   }
 
+  /** @internal */
+  public _checkEdgeCompatibility(): void {
+    // Check env vars
+    const envEntries = Object.entries(this.environment);
+    for (const [key, config] of envEntries) {
+      if (config.removeInEdge) {
+        delete this.environment[key];
+        this.node.addInfo(`Removed ${key} environment variable for Lambda@Edge compatibility`);
+      }
+    }
+    const envKeys = Object.keys(this.environment);
+    if (envKeys.length !== 0) {
+      throw new Error(`The function ${this.node.path} contains environment variables [${envKeys}] and is not compatible with Lambda@Edge. \
+Environment variables can be marked for removal when used in Lambda@Edge by setting the \'removeInEdge\' property in the \'addEnvironment()\' API.`);
+    }
+
+    return;
+  }
+
   private renderEnvironment() {
     if (!this.environment || Object.keys(this.environment).length === 0) {
       return undefined;
     }
 
-    // for backwards compatibility we do not sort environment variables in case
-    // _currentVersion is not defined. otherwise, this would have invalidated
+    const variables: { [key: string]: string } = {};
+    // Sort environment so the hash of the function used to create
+    // `currentVersion` is not affected by key order (this is how lambda does
+    // it). For backwards compatibility we do not sort environment variables in case
+    // _currentVersion is not defined. Otherwise, this would have invalidated
     // the template, and for example, may cause unneeded updates for nested
     // stacks.
-    if (!this._currentVersion) {
-      return {
-        variables: this.environment,
-      };
-    }
+    const keys = this._currentVersion
+      ? Object.keys(this.environment).sort()
+      : Object.keys(this.environment);
 
-    // sort environment so the hash of the function used to create
-    // `currentVersion` is not affected by key order (this is how lambda does
-    // it).
-    const variables: { [key: string]: string } = {};
-    for (const key of Object.keys(this.environment).sort()) {
-      variables[key] = this.environment[key];
+    for (const key of keys) {
+      variables[key] = this.environment[key].value;
     }
 
     return { variables };
@@ -905,6 +925,27 @@ export class Function extends FunctionBase {
   }
 }
 
+/**
+ * Environment variables options
+ */
+export interface EnvironmentOptions {
+  /**
+   * When used in Lambda@Edge via edgeArn() API, these environment
+   * variables will be removed. If not set, an error will be thrown.
+   * @see https://docs.aws.amazon.com/AmazonCloudFront/latest/DeveloperGuide/lambda-requirements-limits.html#lambda-requirements-lambda-function-configuration
+   *
+   * @default false - using the function in Lambda@Edge will throw
+   */
+  readonly removeInEdge?: boolean
+}
+
+/**
+ * Configuration for an environment variable
+ */
+interface EnvironmentConfig extends EnvironmentOptions {
+  readonly value: string;
+}
+
 /**
  * Given an opaque (token) ARN, returns a CloudFormation expression that extracts the function
  * name from the ARN.