From 623675d2f8fb2786f23beb87994e687e8a7c6612 Mon Sep 17 00:00:00 2001
From: Adam Ruka <adamruka@amazon.com>
Date: Fri, 5 Mar 2021 17:20:53 -0800
Subject: [PATCH] fix(cfn-include): allow dynamic mappings to be used in
 Fn::FindInMap (#13428)

The template parsing logic in cloudformation-include always searched for the Mapping
in the template based on the first argument passed to Fn::FindInMap.
However, that doesn't work if that first argument is a dynamic expression,
like `{ Ref: Param }`.

Check for that case explicitly, and don't search for the Mapping
if the first argument to Fn::FindInMap is a dynamic expression.


----

*By submitting this pull request, I confirm that my contribution is made under the terms of the Apache-2.0 license*
---
 .../find-in-map-with-dynamic-mapping.json     | 30 +++++++++++++++++++
 .../test/valid-templates.test.ts              |  8 +++++
 packages/@aws-cdk/core/lib/cfn-parse.ts       | 16 +++++++---
 3 files changed, 50 insertions(+), 4 deletions(-)
 create mode 100644 packages/@aws-cdk/cloudformation-include/test/test-templates/find-in-map-with-dynamic-mapping.json

diff --git a/packages/@aws-cdk/cloudformation-include/test/test-templates/find-in-map-with-dynamic-mapping.json b/packages/@aws-cdk/cloudformation-include/test/test-templates/find-in-map-with-dynamic-mapping.json
new file mode 100644
index 0000000000000..aedf3250272fc
--- /dev/null
+++ b/packages/@aws-cdk/cloudformation-include/test/test-templates/find-in-map-with-dynamic-mapping.json
@@ -0,0 +1,30 @@
+{
+  "Parameters": {
+    "Stage": {
+      "Type": "String",
+      "AllowedValues": ["beta"],
+      "Default": "beta"
+    }
+  },
+  "Mappings": {
+    "beta": {
+      "region": {
+        "key1": "name"
+      }
+    }
+  },
+  "Resources": {
+    "Bucket": {
+      "Type": "AWS::S3::Bucket",
+      "Properties": {
+        "BucketName": {
+          "Fn::FindInMap": [
+            { "Ref": "Stage" },
+            "region",
+            "key1"
+          ]
+        }
+      }
+    }
+  }
+}
diff --git a/packages/@aws-cdk/cloudformation-include/test/valid-templates.test.ts b/packages/@aws-cdk/cloudformation-include/test/valid-templates.test.ts
index 873add618b9d0..34b403a2e7a99 100644
--- a/packages/@aws-cdk/cloudformation-include/test/valid-templates.test.ts
+++ b/packages/@aws-cdk/cloudformation-include/test/valid-templates.test.ts
@@ -746,6 +746,14 @@ describe('CDK Include', () => {
     }).toThrow(/Mapping with name 'NonExistentMapping' was not found in the template/);
   });
 
+  test('can ingest a template that uses Fn::FindInMap with the first argument being a dynamic reference', () => {
+    includeTestTemplate(stack, 'find-in-map-with-dynamic-mapping.json');
+
+    expect(stack).toMatchTemplate(
+      loadTestFileToJsObject('find-in-map-with-dynamic-mapping.json'),
+    );
+  });
+
   test('handles renaming Mapping references', () => {
     const cfnTemplate = includeTestTemplate(stack, 'only-mapping-and-bucket.json');
     const someMapping = cfnTemplate.getMapping('SomeMapping');
diff --git a/packages/@aws-cdk/core/lib/cfn-parse.ts b/packages/@aws-cdk/core/lib/cfn-parse.ts
index 1a5c245b61c4e..93376ae19d365 100644
--- a/packages/@aws-cdk/core/lib/cfn-parse.ts
+++ b/packages/@aws-cdk/core/lib/cfn-parse.ts
@@ -566,11 +566,19 @@ export class CfnParser {
       case 'Fn::FindInMap': {
         const value = this.parseValue(object[key]);
         // the first argument to FindInMap is the mapping name
-        const mapping = this.finder.findMapping(value[0]);
-        if (!mapping) {
-          throw new Error(`Mapping used in FindInMap expression with name '${value[0]}' was not found in the template`);
+        let mappingName: string;
+        if (Token.isUnresolved(value[0])) {
+          // the first argument can be a dynamic expression like Ref: Param;
+          // if it is, we can't find the mapping in advance
+          mappingName = value[0];
+        } else {
+          const mapping = this.finder.findMapping(value[0]);
+          if (!mapping) {
+            throw new Error(`Mapping used in FindInMap expression with name '${value[0]}' was not found in the template`);
+          }
+          mappingName = mapping.logicalId;
         }
-        return Fn._findInMap(mapping.logicalId, value[1], value[2]);
+        return Fn._findInMap(mappingName, value[1], value[2]);
       }
       case 'Fn::Select': {
         const value = this.parseValue(object[key]);