From 6a2f6fde4f97f254eb4ef3b79cab99f574abf72a Mon Sep 17 00:00:00 2001
From: Nikita Pivkin <nikita.pivkin@smartforce.io>
Date: Fri, 29 Mar 2024 00:50:36 +0300
Subject: [PATCH] fix(cloudformation): infer type after resolving a function
 (#6406)

---
 .../scanners/cloudformation/cftypes/types.go  | 24 +++++++++++++++++++
 .../parser/fn_find_in_map_test.go             | 23 ++++++++++++++++++
 .../cloudformation/parser/intrinsics.go       |  9 +++++--
 .../cloudformation/parser/property.go         |  8 +++++++
 .../cloudformation/parser/resource.go         |  7 ++----
 .../scanners/cloudformation/parser/util.go    |  4 +++-
 6 files changed, 67 insertions(+), 8 deletions(-)

diff --git a/pkg/iac/scanners/cloudformation/cftypes/types.go b/pkg/iac/scanners/cloudformation/cftypes/types.go
index 44d9c1fd2a93..0dc3b8b586a2 100644
--- a/pkg/iac/scanners/cloudformation/cftypes/types.go
+++ b/pkg/iac/scanners/cloudformation/cftypes/types.go
@@ -1,5 +1,7 @@
 package cftypes
 
+import "reflect"
+
 type CfType string
 
 const (
@@ -9,4 +11,26 @@ const (
 	Bool    CfType = "bool"
 	Map     CfType = "map"
 	List    CfType = "list"
+	Unknown CfType = "unknown"
 )
+
+func TypeFromGoValue(value interface{}) CfType {
+	switch reflect.TypeOf(value).Kind() {
+	case reflect.String:
+		return String
+	case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
+		return Int
+	case reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64:
+		return Int
+	case reflect.Float32, reflect.Float64:
+		return Float64
+	case reflect.Bool:
+		return Bool
+	case reflect.Map:
+		return Map
+	case reflect.Slice:
+		return List
+	default:
+		return Unknown
+	}
+}
diff --git a/pkg/iac/scanners/cloudformation/parser/fn_find_in_map_test.go b/pkg/iac/scanners/cloudformation/parser/fn_find_in_map_test.go
index bbfa372b7121..6063c39fc006 100644
--- a/pkg/iac/scanners/cloudformation/parser/fn_find_in_map_test.go
+++ b/pkg/iac/scanners/cloudformation/parser/fn_find_in_map_test.go
@@ -98,3 +98,26 @@ Resources:
 	nodeTypeProp := testRes.GetStringProperty("CacheNodeType", "")
 	assert.Equal(t, "cache.t2.micro", nodeTypeProp.Value())
 }
+
+func Test_InferType(t *testing.T) {
+	source := `---
+Mappings:
+  ApiDB:
+     MultiAZ:
+        development: False
+Resources:
+  ApiDB:
+    Type: AWS::RDS::DBInstance
+    Properties:
+      MultiAZ: !FindInMap [ApiDB, MultiAZ, development]
+`
+
+	ctx := createTestFileContext(t, source)
+	require.NotNil(t, ctx)
+
+	testRes := ctx.GetResourceByLogicalID("ApiDB")
+	require.NotNil(t, testRes)
+
+	nodeTypeProp := testRes.GetBoolProperty("MultiAZ")
+	assert.False(t, nodeTypeProp.Value())
+}
diff --git a/pkg/iac/scanners/cloudformation/parser/intrinsics.go b/pkg/iac/scanners/cloudformation/parser/intrinsics.go
index d455fd3d5c6e..1dadc4f6d6fd 100644
--- a/pkg/iac/scanners/cloudformation/parser/intrinsics.go
+++ b/pkg/iac/scanners/cloudformation/parser/intrinsics.go
@@ -78,8 +78,13 @@ func ResolveIntrinsicFunc(property *Property) (*Property, bool) {
 
 	for funcName := range property.AsMap() {
 		if fn := intrinsicFuncs[funcName]; fn != nil {
-			//
-			return fn(property)
+			prop, resolved := fn(property)
+			if prop == nil || !resolved {
+				return prop, false
+			}
+
+			prop.inferType()
+			return prop, true
 		}
 	}
 	return property, false
diff --git a/pkg/iac/scanners/cloudformation/parser/property.go b/pkg/iac/scanners/cloudformation/parser/property.go
index 3cdbbb36b58a..e667c7844f44 100644
--- a/pkg/iac/scanners/cloudformation/parser/property.go
+++ b/pkg/iac/scanners/cloudformation/parser/property.go
@@ -425,3 +425,11 @@ func convert(input interface{}) interface{} {
 	}
 	return input
 }
+
+func (p *Property) inferType() {
+	typ := cftypes.TypeFromGoValue(p.Inner.Value)
+	if typ == cftypes.Unknown {
+		return
+	}
+	p.Inner.Type = typ
+}
diff --git a/pkg/iac/scanners/cloudformation/parser/resource.go b/pkg/iac/scanners/cloudformation/parser/resource.go
index 69a864ad7cdf..bd1351f234df 100644
--- a/pkg/iac/scanners/cloudformation/parser/resource.go
+++ b/pkg/iac/scanners/cloudformation/parser/resource.go
@@ -100,11 +100,8 @@ func (r *Resource) GetProperty(path string) *Property {
 	first := pathParts[0]
 	property := &Property{}
 
-	for n, p := range r.properties() {
-		if n == first {
-			property = p
-			break
-		}
+	if p, exists := r.properties()[first]; exists {
+		property = p
 	}
 
 	if len(pathParts) == 1 || property.IsNil() {
diff --git a/pkg/iac/scanners/cloudformation/parser/util.go b/pkg/iac/scanners/cloudformation/parser/util.go
index a0792cf32865..03b9bf8da837 100644
--- a/pkg/iac/scanners/cloudformation/parser/util.go
+++ b/pkg/iac/scanners/cloudformation/parser/util.go
@@ -66,13 +66,15 @@ func setPropertyValueFromYaml(node *yaml.Node, propertyData *PropertyInner) erro
 	if node.Content == nil {
 
 		switch node.Tag {
-
 		case "!!int":
 			propertyData.Type = cftypes.Int
 			propertyData.Value, _ = strconv.Atoi(node.Value)
 		case "!!bool":
 			propertyData.Type = cftypes.Bool
 			propertyData.Value, _ = strconv.ParseBool(node.Value)
+		case "!!float":
+			propertyData.Type = cftypes.Float64
+			propertyData.Value, _ = strconv.ParseFloat(node.Value, 64)
 		case "!!str", "!!string":
 			propertyData.Type = cftypes.String
 			propertyData.Value = node.Value