fix(terraform): for_each on a map returns a resource for every key (#9156)

Author: sontek

pkg/iac/adapters/terraform/aws/iam/roles_test.go

@@ -2,6 +2,7 @@ package iam
 
 import (
 	"sort"
+	"strings"
 	"testing"
 
 	"github.com/aquasecurity/iamgo"
@@ -377,3 +378,171 @@ data "aws_partition" "current" {}
 		})
 	}
 }
+
+// validateLambdaEcsKeys validates that attachment references contain both lambda and ecs-tasks keys
+func validateLambdaEcsKeys(t *testing.T, attachmentRefs []string) {
+	hasLambda := false
+	hasEcs := false
+	for _, ref := range attachmentRefs {
+		if strings.Contains(ref, "lambda") {
+			hasLambda = true
+		}
+		if strings.Contains(ref, "ecs-tasks") {
+			hasEcs = true
+		}
+	}
+	if !hasLambda || !hasEcs {
+		t.Errorf("expected attachment refs to include both lambda and ecs-tasks keys, got %v", attachmentRefs)
+	}
+}
+
+func Test_forEachReferences(t *testing.T) {
+	tests := []struct {
+		name          string
+		terraform     string
+		expectedCount int
+	}{
+		{
+			name: "computed local with for_each map",
+			terraform: `
+locals {
+  platform_role_principals = {
+    lambda    = "lambda.amazonaws.com"
+    ecs-tasks = "ecs-tasks.amazonaws.com"
+  }
+}
+
+data "aws_iam_policy_document" "platform_role_assume_policy" {
+  for_each = local.platform_role_principals
+
+  statement {
+    actions = ["sts:AssumeRole"]
+    principals {
+      type        = "Service"
+      identifiers = [each.value]
+    }
+  }
+}
+
+resource "aws_iam_role" "platform_role" {
+  for_each           = local.platform_role_principals
+  name               = "platform-${each.key}"
+  assume_role_policy = data.aws_iam_policy_document.platform_role_assume_policy[each.key].json
+}
+
+locals {
+  platform_roles = {
+    for role_key, role_res in aws_iam_role.platform_role :
+    role_key => {
+      role = role_res.name
+    }
+  }
+}
+
+data "aws_iam_policy_document" "administrative_policy_doc" {
+  statement {
+    resources = ["*"]
+    actions   = ["Tag:GetResources", "Tag:TagResources", "Tag:UntagResources"]
+  }
+}
+
+resource "aws_iam_policy" "administrative_policy" {
+  name   = "administrative-policy"
+  policy = data.aws_iam_policy_document.administrative_policy_doc.json
+}
+
+resource "aws_iam_role_policy_attachment" "administrative_policy_attachment" {
+  for_each   = local.platform_roles
+  role       = each.value.role
+  policy_arn = aws_iam_policy.administrative_policy.arn
+}`,
+			expectedCount: 2,
+		},
+		{
+			name: "direct for_each reference",
+			terraform: `
+locals {
+  roles = {
+    lambda    = "lambda.amazonaws.com"
+    ecs-tasks = "ecs-tasks.amazonaws.com"
+  }
+}
+
+resource "aws_iam_role" "platform_role" {
+  for_each = local.roles
+  name     = "platform-${each.key}"
+}
+
+resource "aws_iam_role_policy_attachment" "test" {
+  for_each = aws_iam_role.platform_role
+  role     = each.value.name
+  policy_arn = "arn:aws:iam::aws:policy/AmazonS3FullAccess"
+}`,
+			expectedCount: 2,
+		},
+		{
+			name: "for_each with computed local reference",
+			terraform: `
+locals {
+  role_principals = {
+    lambda    = "lambda.amazonaws.com"
+    ecs-tasks = "ecs-tasks.amazonaws.com"
+  }
+}
+
+resource "aws_iam_role" "platform_role" {
+  for_each = local.role_principals
+  name     = "platform-${each.key}"
+}
+
+locals {
+  platform_roles = {
+    for role_key, role_res in aws_iam_role.platform_role :
+    role_key => {
+      role = role_res.name
+    }
+  }
+}
+
+resource "aws_iam_role_policy_attachment" "test" {
+  for_each = local.platform_roles
+  role     = each.value.role
+  policy_arn = "arn:aws:iam::aws:policy/AmazonS3FullAccess"
+}`,
+			expectedCount: 2,
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			modules := tftestutil.CreateModulesFromSource(t, test.terraform, ".tf")
+			attachments := modules.GetResourcesByType("aws_iam_role_policy_attachment")
+
+			// Debug output for troubleshooting
+			t.Logf("Total resources found: %d", len(attachments))
+			for i, attachment := range attachments {
+				t.Logf("Attachment %d: %s", i, attachment.Reference().String())
+				t.Logf("  - FullName: %s", attachment.FullName())
+				t.Logf("  - TypeLabel: %s", attachment.TypeLabel())
+				t.Logf("  - NameLabel: %s", attachment.NameLabel())
+				if key := attachment.Reference().RawKey(); !key.IsNull() && key.IsKnown() {
+					t.Logf("  - Key: %s (%s)", key.GoString(), key.Type().GoString())
+				}
+			}
+
+			var attachmentRefs []string
+			for _, a := range attachments {
+				attachmentRefs = append(attachmentRefs, a.Reference().String())
+			}
+
+			sort.Strings(attachmentRefs)
+
+			if len(attachments) != test.expectedCount {
+				t.Fatalf("expected %d policy attachments, got %d: %v", test.expectedCount, len(attachments), attachmentRefs)
+			}
+
+			// Validate that both lambda and ecs-tasks keys are present
+			validateLambdaEcsKeys(t, attachmentRefs)
+		})
+	}
+}

pkg/iac/scanners/terraform/parser/evaluator.go

@@ -142,6 +142,14 @@ func (e *evaluator) EvaluateAll(ctx context.Context) (terraform.Modules, map[str
 	e.blocks = e.expandBlocks(e.blocks)
 	e.blocks = e.expandBlocks(e.blocks)
 
+	// Re-evaluate locals after all expansions to ensure computed locals
+	// that depend on expanded resources get the correct values
+	e.ctx.Replace(e.getValuesByBlockType("locals"), "local")
+
+	// Final expansion round to handle any previously deferred for_each blocks
+	// that now have correct local values
+	e.blocks = e.expandBlockForEaches(e.blocks)
+
 	// rootModule is initialized here, but not fully evaluated until all submodules are evaluated.
 	// Initializing it up front to keep the module hierarchy of parents correct.
 	rootModule := terraform.NewModule(e.projectRootPath, e.modulePath, e.blocks, e.ignores)
@@ -309,7 +317,7 @@ func (e *evaluator) expandBlockForEaches(blocks terraform.Blocks) terraform.Bloc
 
 		forEachAttr := block.GetAttribute("for_each")
 
-		if forEachAttr.IsNil() || block.IsExpanded() || !isBlockSupportsForEachMetaArgument(block) {
+		if forEachAttr.IsNil() || block.IsExpanded() || !isBlockSupportsForEachMetaArgument(block) || e.shouldDeferForEachExpansion(forEachAttr) {
 			forEachFiltered = append(forEachFiltered, block)
 			continue
 		}
@@ -545,7 +553,8 @@ func (e *evaluator) getValuesByBlockType(blockType string) cty.Value {
 			}
 			values[b.Label()] = val
 		case "locals", "moved", "import":
-			maps.Copy(values, b.Values().AsValueMap())
+			localValues := b.Values().AsValueMap()
+			maps.Copy(values, localValues)
 		case "provider", "module", "check":
 			if b.Label() == "" {
 				continue
@@ -597,12 +606,16 @@ func (e *evaluator) getResources() map[string]cty.Value {
 			typeValues = make(map[string]cty.Value)
 			values[ref.TypeLabel()] = typeValues
 		}
-		typeValues[ref.NameLabel()] = blockInstanceValues(b, typeValues, b.Values())
+
+		instanceVal := blockInstanceValues(b, typeValues, b.Values())
+		typeValues[ref.NameLabel()] = instanceVal
 	}
 
-	return lo.MapValues(values, func(v map[string]cty.Value, _ string) cty.Value {
+	result := lo.MapValues(values, func(v map[string]cty.Value, _ string) cty.Value {
 		return cty.ObjectVal(v)
 	})
+
+	return result
 }
 
 // blockInstanceValues returns a cty.Value containing the values of the block instances.
@@ -640,3 +653,52 @@ func blockInstanceValues(b *terraform.Block, typeValues map[string]cty.Value, va
 func isForEachKey(key cty.Value) bool {
 	return key.Type().Equals(cty.Number) || key.Type().Equals(cty.String)
 }
+
+func (e *evaluator) shouldDeferForEachExpansion(forEachAttr *terraform.Attribute) bool {
+	// Check if the for_each references a local value
+	for _, ref := range forEachAttr.AllReferences() {
+		if ref.BlockType().Name() == "locals" {
+			// Get the local value from context
+			localName := ref.NameLabel()
+			if localVal := e.ctx.Get("local", localName); localVal != cty.NilVal && localVal.IsKnown() {
+				// Check if this local contains unresolved dynamic values or looks like
+				// it was computed from a single resource object instead of an expanded collection
+				if e.localHasDynamicValues(localVal) {
+					return true
+				}
+			}
+		}
+	}
+	return false
+}
+
+func (e *evaluator) localHasDynamicValues(localVal cty.Value) bool {
+	if !localVal.Type().IsObjectType() {
+		return false
+	}
+
+	valueMap := localVal.AsValueMap()
+	if valueMap == nil {
+		return false
+	}
+
+	// Check if the local contains DynamicVal which indicates unresolved references
+	for _, val := range valueMap {
+		if val.Type() == cty.DynamicPseudoType {
+			return true
+		}
+
+		// If it's an object, check recursively for dynamic values
+		if val.Type().IsObjectType() {
+			if subMap := val.AsValueMap(); subMap != nil {
+				for _, subVal := range subMap {
+					if subVal.Type() == cty.DynamicPseudoType {
+						return true
+					}
+				}
+			}
+		}
+	}
+
+	return false
+}