From 5f27b011a112783e665c9edf0337d1e4d6666289 Mon Sep 17 00:00:00 2001
From: Marcin Rabenda <xrn.design@gmail.com>
Date: Fri, 9 Jun 2023 14:17:51 +0200
Subject: [PATCH] feat(intrinsics): add support for FindInMap default (#546)

---
 cloudformation/intrinsics.go | 38 ++++++++++++++++++++++++++++++------
 goformation_test.go          |  7 +++++++
 2 files changed, 39 insertions(+), 6 deletions(-)

diff --git a/cloudformation/intrinsics.go b/cloudformation/intrinsics.go
index b6045a0df5..3280ac6278 100644
--- a/cloudformation/intrinsics.go
+++ b/cloudformation/intrinsics.go
@@ -72,6 +72,24 @@ func str3Wrap(fn func(interface{}, interface{}, interface{}) string) intrinsics.
 	}
 }
 
+func findInMapWrap(fn func(interface{}, interface{}, interface{}, ...interface{}) string) intrinsics.IntrinsicHandler {
+	return func(name string, input interface{}, template interface{}) interface{} {
+		if arr, ok := input.([]interface{}); ok {
+			if len(arr) < 3 {
+				return nil
+			}
+
+			mapName := arr[0]
+			topLevelKey := arr[1]
+			secondLevelKey := arr[2]
+			additional := arr[3:]
+
+			return fn(mapName, topLevelKey, secondLevelKey, additional...)
+		}
+		return nil
+	}
+}
+
 func str2AWrap(fn func(interface{}, []string) string) intrinsics.IntrinsicHandler {
 	return func(name string, input interface{}, template interface{}) interface{} {
 		if arr, ok := input.([]interface{}); ok {
@@ -107,7 +125,7 @@ var EncoderIntrinsics = map[string]intrinsics.IntrinsicHandler{
 	"Fn::If":          str3Wrap(If),
 	"Fn::Not":         strAWrap(Not),
 	"Fn::Or":          strAWrap(Or),
-	"Fn::FindInMap":   str3Wrap(FindInMap),
+	"Fn::FindInMap":   findInMapWrap(FindInMap),
 	"Fn::GetAtt":      strSplit2Wrap(GetAtt),
 	"Fn::GetAZs":      strWrap(GetAZs),
 	"Fn::ImportValue": strWrap(ImportValue),
@@ -225,13 +243,21 @@ func CIDRPtr(ipBlock, count, cidrBits interface{}) *string {
 	return String(CIDR(ipBlock, count, cidrBits))
 }
 
-// FindInMap returns the value corresponding to keys in a two-level map that is declared in the Mappings section.
-func FindInMap(mapName, topLevelKey, secondLevelKey interface{}) string {
-	return encode(fmt.Sprintf(`{ "Fn::FindInMap" : [ %q, %q, %q ] }`, mapName, topLevelKey, secondLevelKey))
+func FindInMap(mapName, topLevelKey, secondLevelKey interface{}, additional ...interface{}) string {
+	var defaultValue interface{}
+	if len(additional) > 0 {
+		defaultValue = additional[0]
+	}
+
+	if defaultValue == nil {
+		return encode(fmt.Sprintf(`{ "Fn::FindInMap" : [ %q, %q, %q ] }`, mapName, topLevelKey, secondLevelKey))
+	} else {
+		return encode(fmt.Sprintf(`{ "Fn::FindInMap" : [ %q, %q, %q, { "DefaultValue": %q }] }`, mapName, topLevelKey, secondLevelKey, defaultValue))
+	}
 }
 
-func FindInMapPtr(mapName, topLevelKey, secondLevelKey interface{}) *string {
-	return String(FindInMap(mapName, topLevelKey, secondLevelKey))
+func FindInMapPtr(mapName, topLevelKey, secondLevelKey interface{}, additional ...interface{}) *string {
+	return String(FindInMap(mapName, topLevelKey, secondLevelKey, additional...))
 }
 
 // If returns one value if the specified condition evaluates to true and another value if the specified condition evaluates to false. Currently, AWS CloudFormation supports the Fn::If intrinsic function in the metadata attribute, update policy attribute, and property values in the Resources section and Outputs sections of a template. You can use the AWS::NoValue pseudo parameter as a return value to remove the corresponding property.
diff --git a/goformation_test.go b/goformation_test.go
index d8366a31ab..b636c48c3b 100644
--- a/goformation_test.go
+++ b/goformation_test.go
@@ -1052,6 +1052,13 @@ var _ = Describe("Goformation", func() {
 					"Fn::FindInMap": []interface{}{"test-map", "test-top-level-key", "test-second-level-key"},
 				},
 			},
+			{
+				Name:  "Fn::FindInMap with DefaultValue",
+				Input: cloudformation.FindInMap("test-map", "test-top-level-key", "test-second-level-key", "test-default-value"),
+				Expected: map[string]interface{}{
+					"Fn::FindInMap": []interface{}{"test-map", "test-top-level-key", "test-second-level-key", map[string]interface{}{"DefaultValue": "test-default-value"}},
+				},
+			},
 			{
 				Name:  "Fn::GetAZs",
 				Input: cloudformation.GetAZs("test-region"),