From 6e7a1a925b5b009d60ae493ac59278e0e9f79e0d Mon Sep 17 00:00:00 2001
From: Elad Ben-Israel <benisrae@amazon.com>
Date: Mon, 13 May 2019 09:40:11 +0300
Subject: [PATCH] fix(cdk): CfnMapping.findInMap with tokens

Do not validate existence of second key if the first
key is a token.

Fixes #1363
---
 packages/@aws-cdk/cdk/lib/cfn-mapping.ts    |  4 +-
 packages/@aws-cdk/cdk/test/test.mappings.ts | 43 ++++++++++++++++++++-
 2 files changed, 44 insertions(+), 3 deletions(-)

diff --git a/packages/@aws-cdk/cdk/lib/cfn-mapping.ts b/packages/@aws-cdk/cdk/lib/cfn-mapping.ts
index ac2fd0cfe7827..7ec630bbf312c 100644
--- a/packages/@aws-cdk/cdk/lib/cfn-mapping.ts
+++ b/packages/@aws-cdk/cdk/lib/cfn-mapping.ts
@@ -38,8 +38,8 @@ export class CfnMapping extends CfnRefElement {
       throw new Error(`Mapping doesn't contain top-level key '${key1}'`);
     }
 
-    // opportunistically check that the key exists (if the key does not contain tokens)
-    if (!Token.isToken(key2) && !(key2 in this.mapping[key1])) {
+    // opportunistically check that the second key exists (if the key does not contain tokens)
+    if (!Token.isToken(key1) && !Token.isToken(key2) && !(key2 in this.mapping[key1])) {
       throw new Error(`Mapping doesn't contain second-level key '${key2}'`);
     }
 
diff --git a/packages/@aws-cdk/cdk/test/test.mappings.ts b/packages/@aws-cdk/cdk/test/test.mappings.ts
index f9407e1c1b86d..4f524ae1f1e07 100644
--- a/packages/@aws-cdk/cdk/test/test.mappings.ts
+++ b/packages/@aws-cdk/cdk/test/test.mappings.ts
@@ -62,5 +62,46 @@ export = {
     test.deepEqual(stack.node.resolve(v1), expected);
     test.deepEqual(stack.node.resolve(v2), expected);
     test.done();
-  }
+  },
+
+  'no validation if first key is token and second is a static string'(test: Test) {
+    // GIVEN
+    const stack = new Stack();
+    const mapping = new CfnMapping(stack, 'mapping', {
+      mapping: {
+        'us-east-1': {
+          size: 12
+        }
+      }
+    });
+
+    // WHEN
+    const v = mapping.findInMap(Aws.region, 'size');
+
+    // THEN
+    test.deepEqual(stack.node.resolve(v), {
+      "Fn::FindInMap": [ 'mapping', { Ref: "AWS::Region" }, "size" ]
+    });
+    test.done();
+  },
+
+  'validate first key if it is a string and second is a token'(test: Test) {
+    // GIVEN
+    const stack = new Stack();
+    const mapping = new CfnMapping(stack, 'mapping', {
+      mapping: {
+        size: {
+          'us-east-1': 12
+        }
+      }
+    });
+
+    // WHEN
+    const v = mapping.findInMap('size', Aws.region);
+
+    // THEN
+    test.throws(() => mapping.findInMap('not-found', Aws.region), /Mapping doesn't contain top-level key 'not-found'/);
+    test.deepEqual(stack.node.resolve(v), { 'Fn::FindInMap': [ 'mapping', 'size', { Ref: 'AWS::Region' } ] });
+    test.done();
+  },
 };