From a69190045acf91c48e9070ab647c2877d81d4803 Mon Sep 17 00:00:00 2001
From: Jonathan Goldwasser <jogold@users.noreply.github.com>
Date: Thu, 13 Jun 2019 20:25:19 +0200
Subject: [PATCH] feat(cloudformation): add option to restrict data returned
 AwsCustomResource (#2859)

Specifying `outputPath` allows to restrict the data returned by the custom resource to a specific
path in the API response. This can be used to limit the data returned by the custom resource if
working with API calls that could potentially result in custom respone objects exceeding the hard
limit of 4096 bytes.

Closes #2825
---
 .../lib/aws-custom-resource-provider/index.ts | 27 +++++++---
 .../lib/aws-custom-resource.ts                | 13 +++++
 .../test/test.aws-custom-resource-provider.ts | 49 ++++++++++++++++++-
 3 files changed, 82 insertions(+), 7 deletions(-)

diff --git a/packages/@aws-cdk/aws-cloudformation/lib/aws-custom-resource-provider/index.ts b/packages/@aws-cdk/aws-cloudformation/lib/aws-custom-resource-provider/index.ts
index 5dcf0d9c2e14f..fa9882991392a 100644
--- a/packages/@aws-cdk/aws-cloudformation/lib/aws-custom-resource-provider/index.ts
+++ b/packages/@aws-cdk/aws-cloudformation/lib/aws-custom-resource-provider/index.ts
@@ -33,12 +33,26 @@ function fixBooleans(object: object) {
       : v);
 }
 
+/**
+ * Filters the keys of an object.
+ */
+function filterKeys(object: object, pred: (key: string) => boolean) {
+  return Object.entries(object)
+    .reduce(
+      (acc, [k, v]) => pred(k)
+        ? { ...acc, [k]: v }
+        : acc,
+        {}
+    );
+}
+
 export async function handler(event: AWSLambda.CloudFormationCustomResourceEvent, context: AWSLambda.Context) {
   try {
     console.log(JSON.stringify(event));
     console.log('AWS SDK VERSION: ' + (AWS as any).VERSION);
 
     let physicalResourceId = (event as any).PhysicalResourceId;
+    let flatData: { [key: string]: string } = {};
     let data: { [key: string]: string } = {};
     const call: AwsSdkCall | undefined = event.ResourceProperties[event.RequestType];
 
@@ -47,18 +61,19 @@ export async function handler(event: AWSLambda.CloudFormationCustomResourceEvent
 
       try {
         const response = await awsService[call.action](call.parameters && fixBooleans(call.parameters)).promise();
-        data = flatten(response);
+        flatData = flatten(response);
+        data = call.outputPath
+          ? filterKeys(flatData, k => k.startsWith(call.outputPath!))
+          : flatData;
       } catch (e) {
         if (!call.catchErrorPattern || !new RegExp(call.catchErrorPattern).test(e.code)) {
           throw e;
         }
       }
 
-      if (call.physicalResourceIdPath) {
-        physicalResourceId = data[call.physicalResourceIdPath];
-      } else {
-        physicalResourceId = call.physicalResourceId!;
-      }
+      physicalResourceId = call.physicalResourceIdPath
+        ? flatData[call.physicalResourceIdPath]
+        : call.physicalResourceId;
     }
 
     await respond('SUCCESS', 'OK', physicalResourceId, data);
diff --git a/packages/@aws-cdk/aws-cloudformation/lib/aws-custom-resource.ts b/packages/@aws-cdk/aws-cloudformation/lib/aws-custom-resource.ts
index cf4d738db2fbf..6affcb325a3b4 100644
--- a/packages/@aws-cdk/aws-cloudformation/lib/aws-custom-resource.ts
+++ b/packages/@aws-cdk/aws-cloudformation/lib/aws-custom-resource.ts
@@ -71,6 +71,18 @@ export interface AwsSdkCall {
    * @default use latest available API version
    */
   readonly apiVersion?: string;
+
+  /**
+   * Restrict the data returned by the custom resource to a specific path in
+   * the API response. Use this to limit the data returned by the custom
+   * resource if working with API calls that could potentially result in custom
+   * response objects exceeding the hard limit of 4096 bytes.
+   *
+   * Example for ECS / updateService: 'service.deploymentConfiguration.maximumPercent'
+   *
+   * @default return all data
+   */
+  readonly outputPath?: string;
 }
 
 export interface AwsCustomResourceProps {
@@ -158,6 +170,7 @@ export class AwsCustomResource extends cdk.Construct {
 
   /**
    * Returns response data for the AWS SDK call.
+   *
    * Example for S3 / listBucket : 'Buckets.0.Name'
    *
    * @param dataPath the path to the data
diff --git a/packages/@aws-cdk/aws-cloudformation/test/test.aws-custom-resource-provider.ts b/packages/@aws-cdk/aws-cloudformation/test/test.aws-custom-resource-provider.ts
index a28f7035acad3..f1d659a65be75 100644
--- a/packages/@aws-cdk/aws-cloudformation/test/test.aws-custom-resource-provider.ts
+++ b/packages/@aws-cdk/aws-cloudformation/test/test.aws-custom-resource-provider.ts
@@ -223,5 +223,52 @@ export = {
     test.equal(request.isDone(), true);
 
     test.done();
-  }
+  },
+
+  async 'restrict output path'(test: Test) {
+    const listObjectsFake = sinon.fake.resolves({
+      Contents: [
+        {
+          Key: 'first-key',
+          ETag: 'first-key-etag'
+        },
+        {
+          Key: 'second-key',
+          ETag: 'second-key-etag',
+        }
+      ]
+    } as SDK.S3.ListObjectsOutput);
+
+    AWS.mock('S3', 'listObjects', listObjectsFake);
+
+    const event: AWSLambda.CloudFormationCustomResourceCreateEvent = {
+      ...eventCommon,
+      RequestType: 'Create',
+      ResourceProperties: {
+        ServiceToken: 'token',
+        Create: {
+          service: 'S3',
+          action: 'listObjects',
+          parameters: {
+            Bucket: 'my-bucket'
+          },
+          physicalResourceId: 'id',
+          outputPath: 'Contents.0'
+        } as AwsSdkCall
+      }
+    };
+
+    const request = createRequest(body =>
+      body.Status === 'SUCCESS' &&
+      body.PhysicalResourceId === 'id' &&
+      body.Data!['Contents.0.Key'] === 'first-key' &&
+      body.Data!['Contents.1.Key'] === undefined
+    );
+
+    await handler(event, {} as AWSLambda.Context);
+
+    test.equal(request.isDone(), true);
+
+    test.done();
+  },
 };