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

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

@@ -1,3 +1,4 @@
+import { Lazy } from '@aws-cdk/core';
 import { CfnDistribution } from '../cloudfront.generated';
 import { AddBehaviorOptions, ViewerProtocolPolicy } from '../distribution';
 
@@ -48,15 +49,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: Lazy.stringValue({ produce: () => edgeLambda.functionVersion.edgeArn }),
+          eventType: edgeLambda.eventType.toString(),
+        }))
         : undefined,
     };
   }

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

@@ -895,7 +895,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 && cdk.Lazy.stringValue({ produce: () => fna.lambdaFunction.edgeArn }),
           })),
       });
 

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

@@ -523,20 +523,72 @@ describe('with Lambda@Edge functions', () => {
     });
   });
 
-  test('fails creation when attempting to add the $LATEST function version as an edge Lambda to the default behavior', () => {
-    expect(() => {
-      new Distribution(stack, 'MyDist', {
-        defaultBehavior: {
-          origin,
-          edgeLambdas: [
-            {
-              functionVersion: lambdaFunction.latestVersion,
-              eventType: LambdaEdgeEventType.ORIGIN_RESPONSE,
-            },
-          ],
-        },
-      });
-    }).toThrow(/\$LATEST function version cannot be used for Lambda@Edge/);
+  test('fails synthesis when attempting to add the $LATEST function version as an edge Lambda to the default behavior', () => {
+    new Distribution(stack, 'MyDist', {
+      defaultBehavior: {
+        origin,
+        edgeLambdas: [
+          {
+            functionVersion: lambdaFunction.latestVersion,
+            eventType: LambdaEdgeEventType.ORIGIN_RESPONSE,
+          },
+        ],
+      },
+    });
+    expect(() => app.synth()).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/);
   });
 });
 

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');
+  }
 }