From 64cfcb5b1741b6d538fc05042bd901791e7590a1 Mon Sep 17 00:00:00 2001
From: Ayush Goyal <goyalayu@amazon.com>
Date: Mon, 5 Oct 2020 20:02:18 +0530
Subject: [PATCH] feat(core): Add environment prop to CustomResourceProvider

---
 .../custom-resource-provider.ts               | 26 ++++++++++++++++++
 .../custom-resource-provider.test.ts          | 27 +++++++++++++++++++
 2 files changed, 53 insertions(+)

diff --git a/packages/@aws-cdk/core/lib/custom-resource-provider/custom-resource-provider.ts b/packages/@aws-cdk/core/lib/custom-resource-provider/custom-resource-provider.ts
index 821307e466a77..82160f3a63ad6 100644
--- a/packages/@aws-cdk/core/lib/custom-resource-provider/custom-resource-provider.ts
+++ b/packages/@aws-cdk/core/lib/custom-resource-provider/custom-resource-provider.ts
@@ -60,6 +60,13 @@ export interface CustomResourceProviderProps {
    * @default Size.mebibytes(128)
    */
   readonly memorySize?: Size;
+
+  /**
+   * Key-value pairs that are passed to Lambda as Environment
+   *
+   * @default - No environment variables.
+   */
+  readonly environment?: { [key: string]: string };
 }
 
 /**
@@ -177,6 +184,7 @@ export class CustomResourceProvider extends CoreConstruct {
         Handler: `${ENTRYPOINT_FILENAME}.handler`,
         Role: role.getAtt('Arn'),
         Runtime: 'nodejs12.x',
+        Environment: this.renderEnvironmentVariables(props.environment),
       },
     });
 
@@ -184,4 +192,22 @@ export class CustomResourceProvider extends CoreConstruct {
 
     this.serviceToken = Token.asString(handler.getAtt('Arn'));
   }
+
+  private renderEnvironmentVariables(env?: { [key: string]: string }) {
+    if (!env || Object.keys(env).length === 0) {
+      return undefined;
+    }
+
+    // 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 } = {};
+    const keys = Object.keys(env).sort();
+
+    for (const key of keys) {
+      variables[key] = env[key];
+    }
+
+    return { Variables: variables };
+  }
 }
diff --git a/packages/@aws-cdk/core/test/custom-resource-provider/custom-resource-provider.test.ts b/packages/@aws-cdk/core/test/custom-resource-provider/custom-resource-provider.test.ts
index f3dbff18ed651..3a08468d9dc32 100644
--- a/packages/@aws-cdk/core/test/custom-resource-provider/custom-resource-provider.test.ts
+++ b/packages/@aws-cdk/core/test/custom-resource-provider/custom-resource-provider.test.ts
@@ -169,4 +169,31 @@ nodeunitShim({
     test.deepEqual(lambda.Properties.Timeout, 300);
     test.done();
   },
+
+  'environment variables'(test: Test) {
+    // GIVEN
+    const stack = new Stack();
+
+    // WHEN
+    CustomResourceProvider.getOrCreate(stack, 'Custom:MyResourceType', {
+      codeDirectory: TEST_HANDLER,
+      runtime: CustomResourceProviderRuntime.NODEJS_12,
+      environment: {
+        B: 'b',
+        A: 'a',
+      },
+    });
+
+    // THEN
+    const template = toCloudFormation(stack);
+    const lambda = template.Resources.CustomMyResourceTypeCustomResourceProviderHandler29FBDD2A;
+    test.deepEqual(lambda.Properties.Environment, {
+      Variables: {
+        A: 'a',
+        B: 'b',
+      },
+    });
+    test.done();
+  },
 });
+