fix(cloudfront,lambda): remove env vars when Lambda function is used …

…at edge

Automatically remove environment variables of a Lambda function when
it's used at edge. A warning is also shown.

Works with singleton functions.

Closes aws#9328
Closes aws#9453

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

@@ -52,6 +52,11 @@ export class CacheBehavior {
           if (edgeLambda.functionVersion.version === '$LATEST') {
             throw new Error('$LATEST function version cannot be used for Lambda@Edge');
           }
+
+          if (edgeLambda.functionVersion.lambda.removeEnvironment()) {
+            edgeLambda.functionVersion.node.addWarning(`Removed environment variables from function ${edgeLambda.functionVersion.node.path} because Lambda@Edge does not support environment variables`);
+          }
+
           return {
             lambdaFunctionArn: edgeLambda.functionVersion.functionArn,
             eventType: edgeLambda.eventType.toString(),

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

@@ -949,10 +949,15 @@ export class CloudFrontWebDistribution extends cdk.Resource implements IDistribu
     if (input.lambdaFunctionAssociations) {
       toReturn = Object.assign(toReturn, {
         lambdaFunctionAssociations: input.lambdaFunctionAssociations
-          .map(fna => ({
-            eventType: fna.eventType,
-            lambdaFunctionArn: fna.lambdaFunction && fna.lambdaFunction.functionArn,
-          })),
+          .map(fna => {
+            if (fna.lambdaFunction.lambda.removeEnvironment()) {
+              fna.lambdaFunction.lambda.node.addWarning(`Removed environment variables from function ${fna.lambdaFunction.node.path} because Lambda@Edge does not support environment variables`);
+            }
+            return {
+              eventType: fna.eventType,
+              lambdaFunctionArn: fna.lambdaFunction && fna.lambdaFunction.functionArn,
+            };
+          }),
       });
 
       // allow edgelambda.amazonaws.com to assume the functions' execution role.

packages/@aws-cdk/aws-cloudfront/package.json

@@ -64,6 +64,7 @@
   "license": "Apache-2.0",
   "devDependencies": {
     "@aws-cdk/assert": "0.0.0",
+    "@aws-cdk/cloud-assembly-schema": "0.0.0",
     "aws-sdk": "^2.715.0",
     "cdk-build-tools": "0.0.0",
     "cdk-integ-tools": "0.0.0",

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

@@ -1,6 +1,8 @@
 import '@aws-cdk/assert/jest';
+import { ABSENT } from '@aws-cdk/assert';
 import * as acm from '@aws-cdk/aws-certificatemanager';
 import * as lambda from '@aws-cdk/aws-lambda';
+import * as cxschema from '@aws-cdk/cloud-assembly-schema';
 import { App, Duration, Stack } from '@aws-cdk/core';
 import { CfnDistribution, Distribution, IOrigin, LambdaEdgeEventType, OriginBase, OriginProps, OriginProtocolPolicy, PriceClass } from '../lib';
 
@@ -366,6 +368,39 @@ describe('with Lambda@Edge functions', () => {
       });
     }).toThrow(/\$LATEST function version cannot be used for Lambda@Edge/);
   });
+
+  test('a warning is added when env vars are removed for edge functions', () => {
+    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(stack).toHaveResource('AWS::Lambda::Function', {
+      Environment: ABSENT,
+      Code: {
+        ZipFile: 'whateverwithenv',
+      },
+    });
+
+    expect(envLambdaFunction.currentVersion.node.metadata[0].type).toBe(cxschema.ArtifactMetadataEntryType.WARN);
+    expect(envLambdaFunction.currentVersion.node.metadata[0].data).toBe('Removed environment variables from function Stack/EnvFunction/CurrentVersion because Lambda@Edge does not support environment variables');
+  });
 });
 
 test('price class is included if provided', () => {

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

@@ -1,7 +1,8 @@
-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';
+import * as cxschema from '@aws-cdk/cloud-assembly-schema';
 import * as cdk from '@aws-cdk/core';
 import { nodeunitShim, Test } from 'nodeunit-shim';
 import {
@@ -476,6 +477,49 @@ nodeunitShim({
     test.done();
   },
 
+  'a warning is added when env vars of associated lambda are removed'(test: Test) {
+    const stack = new cdk.Stack();
+    const sourceBucket = new s3.Bucket(stack, 'Bucket');
+
+    const lambdaFunction = new lambda.SingletonFunction(stack, 'Lambda', {
+      uuid: 'xxxx-xxxx-xxxx-xxxx',
+      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.latestVersion,
+              }],
+            },
+          ],
+        },
+      ],
+    });
+
+    expect(stack).to(haveResourceLike('AWS::Lambda::Function', {
+      Environment: ABSENT,
+    }));
+
+    test.equal(lambdaFunction.node.metadata[0].type, cxschema.ArtifactMetadataEntryType.WARN);
+    test.equal(lambdaFunction.node.metadata[0].data, 'Removed environment variables from function Default/Lambda/$LATEST because Lambda@Edge does not support environment variables');
+
+    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/lib/function-base.ts

@@ -112,6 +112,13 @@ export interface IFunction extends IResource, ec2.IConnectable, iam.IGrantable {
    * Configures options for asynchronous invocation.
    */
   configureAsyncInvoke(options: EventInvokeConfigOptions): void;
+
+  /**
+   * Removes environment variables (useful when working with Lambda@Edge).
+   *
+   * This is a no-op for imported functions.
+   */
+  removeEnvironment(): boolean;
 }
 
 /**
@@ -318,6 +325,10 @@ export abstract class FunctionBase extends Resource implements IFunction {
     });
   }
 
+  public removeEnvironment(): boolean {
+    return false;
+  }
+
   /**
    * 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.

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

@@ -504,7 +504,7 @@ export class Function extends FunctionBase {
   /**
    * Environment variables for this function
    */
-  private readonly environment: { [key: string]: string };
+  private environment: { [key: string]: string };
 
   private readonly currentVersionOptions?: VersionOptions;
   private _currentVersion?: Version;
@@ -751,6 +751,14 @@ export class Function extends FunctionBase {
     return this._logGroup;
   }
 
+  public removeEnvironment(): boolean {
+    if (Object.keys(this.environment).length !== 0) {
+      this.environment = {};
+      return true;
+    }
+    return false;
+  }
+
   private renderEnvironment() {
     if (!this.environment || Object.keys(this.environment).length === 0) {
       return undefined;

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

@@ -81,6 +81,10 @@ export class SingletonFunction extends FunctionBase {
     down.node.addDependency(this.lambdaFunction);
   }
 
+  public removeEnvironment() {
+    return this.lambdaFunction.removeEnvironment();
+  }
+
   /**
    * Returns the construct tree node that corresponds to the lambda function.
    * @internal

packages/@aws-cdk/aws-lambda/test/test.lambda.ts

@@ -1,5 +1,5 @@
 import * as path from 'path';
-import { expect, haveResource, MatchStyle, ResourcePart } from '@aws-cdk/assert';
+import { ABSENT, expect, haveResource, MatchStyle, ResourcePart } from '@aws-cdk/assert';
 import * as ec2 from '@aws-cdk/aws-ec2';
 import * as iam from '@aws-cdk/aws-iam';
 import * as logs from '@aws-cdk/aws-logs';
@@ -1525,6 +1525,48 @@ export = {
 
     test.done();
   },
+
+  'can remove env vars'(test: Test) {
+    // GIVEN
+    const stack = new cdk.Stack();
+    const fn = new lambda.Function(stack, 'fn', {
+      code: new lambda.InlineCode('foo'),
+      handler: 'index.handler',
+      runtime: lambda.Runtime.NODEJS_12_X,
+      environment: {
+        KEY: 'value',
+      },
+    });
+
+    // WHEN
+    const removed = fn.removeEnvironment();
+
+    // THEN
+    expect(stack).to(haveResource('AWS::Lambda::Function', {
+      Environment: ABSENT,
+    }));
+    test.ok(removed);
+
+    test.done();
+  },
+
+  'removeEnvironment returns false without env vars'(test: Test) {
+    // GIVEN
+    const stack = new cdk.Stack();
+    const fn = new lambda.Function(stack, 'fn', {
+      code: new lambda.InlineCode('foo'),
+      handler: 'index.handler',
+      runtime: lambda.Runtime.NODEJS_12_X,
+    });
+
+    // WHEN
+    const removed = fn.removeEnvironment();
+
+    // THEN
+    test.equal(removed, false);
+
+    test.done();
+  },
 };
 
 function newTestLambda(scope: cdk.Construct) {