feat(cloudformation): support default values and list results in Fn::FindInMap (#9515)

Signed-off-by: nikpivkin <[email protected]>

Author: nikpivkin

pkg/iac/scanners/cloudformation/parser/fn_find_in_map.go

@@ -1,45 +1,80 @@
 package parser
 
 import (
+	"fmt"
+
 	"github.com/aquasecurity/trivy/pkg/iac/scanners/cloudformation/cftypes"
 )
 
-func ResolveFindInMap(property *Property) (resolved *Property, success bool) {
+func ResolveFindInMap(property *Property) (*Property, bool) {
 	if !property.isFunction() {
 		return property, true
 	}
 
 	refValue := property.AsMap()["Fn::FindInMap"].AsList()
 
-	if len(refValue) != 3 {
-		return abortIntrinsic(property, "Fn::FindInMap should have exactly 3 values, returning original Property")
+	if len(refValue) < 3 || len(refValue) > 4 {
+		return abortIntrinsic(property, "Fn::FindInMap expects 3 or 4 arguments")
 	}
 
-	mapName := refValue[0].AsString()
-	topLevelKey := refValue[1].AsString()
-	secondaryLevelKey := refValue[2].AsString()
-
 	if property.ctx == nil {
-		return abortIntrinsic(property, "the property does not have an attached context, returning original Property")
+		return abortIntrinsic(property, "property context is missing")
 	}
 
-	m, ok := property.ctx.Mappings[mapName]
-	if !ok {
-		return abortIntrinsic(property, "could not find map %s, returning original Property")
+	var defaultValue any
+	if len(refValue) == 4 {
+		if m := refValue[3].AsMap(); m != nil {
+			if defProp, exists := m["DefaultValue"]; exists && defProp != nil {
+				defaultValue = defProp.RawValue()
+			}
+		}
+	}
+
+	mapName := refValue[0].AsString()
+	topKey := refValue[1].AsString()
+	secKey := refValue[2].AsString()
+
+	value, err := resolveMapping(property.ctx, mapName, topKey, secKey)
+	if err != nil {
+		if defaultValue == nil {
+			return abortIntrinsic(property, err.Error())
+		}
+		value = defaultValue
 	}
 
-	mapContents := m.(map[string]any)
+	switch v := value.(type) {
+	case string:
+		return property.deriveResolved(cftypes.String, v), true
+	case []any:
+		elems := make([]*Property, len(v))
+		for i, el := range v {
+			elems[i] = property.deriveResolved(cftypes.String, el)
+		}
+		return property.deriveResolved(cftypes.List, elems), true
+	default:
+		return abortIntrinsic(property, fmt.Sprintf("unsupported type in mapping: %T", v))
+	}
+}
 
-	k, ok := mapContents[topLevelKey]
+func resolveMapping(ctx *FileContext, mapName, topKey, secKey string) (any, error) {
+	m, ok := ctx.Mappings[mapName]
 	if !ok {
-		return abortIntrinsic(property, "could not find %s in the %s map, returning original Property", topLevelKey, mapName)
+		return nil, fmt.Errorf("map %s not found", mapName)
+	}
+	mapContents, ok := m.(map[string]any)
+	if !ok {
+		return nil, fmt.Errorf("map %s has invalid type", mapName)
 	}
 
+	k, ok := mapContents[topKey]
+	if !ok {
+		return nil, fmt.Errorf("key %s not found in map %s", topKey, mapName)
+	}
 	mapValues := k.(map[string]any)
 
-	prop, ok := mapValues[secondaryLevelKey]
+	prop, ok := mapValues[secKey]
 	if !ok {
-		return abortIntrinsic(property, "could not find a value for %s in %s, returning original Property", secondaryLevelKey, topLevelKey)
+		return nil, fmt.Errorf("key %s not found in %s", secKey, topKey)
 	}
-	return property.deriveResolved(cftypes.String, prop), true
+	return prop, nil
 }

pkg/iac/scanners/cloudformation/parser/fn_find_in_map_test.go

@@ -3,100 +3,71 @@ package parser
 import (
 	"testing"
 
+	"github.com/samber/lo"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
 
-func Test_resolve_find_in_map_value(t *testing.T) {
-
+func Test_FindInMap(t *testing.T) {
 	source := `---
 Parameters:
   Environment: 
     Type: String
-    Default: production
+    Default: dev
 Mappings:
   CacheNodeTypes:
     production:
       NodeType: cache.t2.large
+      CacheSecurityGroupNames: [ "sg-1", "sg-2" ]
     test:
       NodeType: cache.t2.small
+      CacheSecurityGroupNames: [ "sg-3" ]
     dev:
       NodeType: cache.t2.micro
+      CacheSecurityGroupNames: [ "sg-4" ]
 Resources:
-    ElasticacheSecurityGroup:
-      Type: 'AWS::EC2::SecurityGroup'
-      Properties:
-        GroupDescription: Elasticache Security Group
-        SecurityGroupIngress:
-          - IpProtocol: tcp
-            FromPort: 11211
-            ToPort: 11211
-            SourceSecurityGroupName: !Ref InstanceSecurityGroup
-    ElasticacheCluster:
-      Type: 'AWS::ElastiCache::CacheCluster'
-      Properties:    
-        Engine: memcached
-        CacheNodeType: !FindInMap [ CacheNodeTypes, production, NodeType ]
-        NumCacheNodes: '1'
-        VpcSecurityGroupIds:
-          - !GetAtt 
-            - ElasticacheSecurityGroup
-            - GroupId
-`
-	ctx := createTestFileContext(t, source)
-	require.NotNil(t, ctx)
-
-	testRes := ctx.GetResourceByLogicalID("ElasticacheCluster")
-	assert.NotNil(t, testRes)
-
-	nodeTypeProp := testRes.GetStringProperty("CacheNodeType", "")
-	assert.Equal(t, "cache.t2.large", nodeTypeProp.Value())
-}
+  ElasticacheCluster:
+    Type: 'AWS::ElastiCache::CacheCluster'
+    Properties:    
+      Engine: memcached
+      CacheNodeType: !FindInMap [ CacheNodeTypes, !Ref Environment, NodeType ]
+      NumCacheNodes: '1'
 
-func Test_resolve_find_in_map_with_nested_intrinsic_value(t *testing.T) {
+  ElasticacheClusterWithDefault:
+    Type: 'AWS::ElastiCache::CacheCluster'
+    Properties:
+      Engine: memcached
+      CacheNodeType: !FindInMap [ CacheNodeTypes, staging, NodeType, DefaultValue: cache.t2.medium ]
+      NumCacheNodes: '1'
 
-	source := `---
-Parameters:
-  Environment: 
-    Type: String
-    Default: dev
-Mappings:
-  CacheNodeTypes:
-    production:
-      NodeType: cache.t2.large
-    test:
-      NodeType: cache.t2.small
-    dev:
-      NodeType: cache.t2.micro
-Resources:
-    ElasticacheSecurityGroup:
-      Type: 'AWS::EC2::SecurityGroup'
-      Properties:
-        GroupDescription: Elasticache Security Group
-        SecurityGroupIngress:
-          - IpProtocol: tcp
-            FromPort: 11211
-            ToPort: 11211
-            SourceSecurityGroupName: !Ref InstanceSecurityGroup
-    ElasticacheCluster:
-      Type: 'AWS::ElastiCache::CacheCluster'
-      Properties:    
-        Engine: memcached
-        CacheNodeType: !FindInMap [ CacheNodeTypes, !Ref Environment, NodeType ]
-        NumCacheNodes: '1'
-        VpcSecurityGroupIds:
-          - !GetAtt 
-            - ElasticacheSecurityGroup
-            - GroupId
+  ElasticacheClusterList:
+    Type: 'AWS::ElastiCache::CacheCluster'
+    Properties:
+      Engine: memcached
+      CacheSecurityGroupNames: !FindInMap [ CacheNodeTypes, production, CacheSecurityGroupNames ]
+      NumCacheNodes: '1'
 `
+
 	ctx := createTestFileContext(t, source)
 	require.NotNil(t, ctx)
 
-	testRes := ctx.GetResourceByLogicalID("ElasticacheCluster")
-	assert.NotNil(t, testRes)
-
-	nodeTypeProp := testRes.GetStringProperty("CacheNodeType", "")
+	cluster := ctx.GetResourceByLogicalID("ElasticacheCluster")
+	require.NotNil(t, cluster)
+	nodeTypeProp := cluster.GetStringProperty("CacheNodeType", "")
 	assert.Equal(t, "cache.t2.micro", nodeTypeProp.Value())
+
+	clusterDefault := ctx.GetResourceByLogicalID("ElasticacheClusterWithDefault")
+	require.NotNil(t, clusterDefault)
+	nodeTypePropDefault := clusterDefault.GetStringProperty("CacheNodeType", "")
+	assert.Equal(t, "cache.t2.medium", nodeTypePropDefault.Value())
+
+	clusterList := ctx.GetResourceByLogicalID("ElasticacheClusterList")
+	require.NotNil(t, clusterList)
+	sgNamesProp := clusterList.GetProperty("CacheSecurityGroupNames").AsList()
+	groupNames := lo.Map(sgNamesProp, func(prop *Property, _ int) any {
+		return prop.AsString()
+	})
+	assert.ElementsMatch(t, []any{"sg-1", "sg-2"}, groupNames)
 }
 
 func Test_InferType(t *testing.T) {