fix(terraform): evaluateStep to correctly set EvalContext for multiple instances of blocks (#8555)

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

Author: Emyrk

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

@@ -0,0 +1,34 @@
+package parser
+
+import "github.com/zclconf/go-cty/cty"
+
+// insertTupleElement inserts a value into a tuple at the specified index.
+// If the idx is outside the bounds of the list, it grows the tuple to
+// the new size, and fills in `cty.NilVal` for the missing elements.
+//
+// This function will not panic. If the list value is not a list, it will
+// be replaced with an empty list.
+func insertTupleElement(list cty.Value, idx int, val cty.Value) cty.Value {
+	if list.IsNull() || !list.Type().IsTupleType() {
+		// better than a panic
+		list = cty.EmptyTupleVal
+	}
+
+	if idx < 0 {
+		// Nothing to do?
+		return list
+	}
+
+	// Create a new list of the correct length, copying in the old list
+	// values for matching indices.
+	newList := make([]cty.Value, max(idx+1, list.LengthInt()))
+	for it := list.ElementIterator(); it.Next(); {
+		key, elem := it.Element()
+		elemIdx, _ := key.AsBigFloat().Int64()
+		newList[elemIdx] = elem
+	}
+	// Insert the new value.
+	newList[idx] = val
+
+	return cty.TupleVal(newList)
+}

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

@@ -0,0 +1,103 @@
+package parser
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+	"github.com/zclconf/go-cty/cty"
+)
+
+func Test_insertTupleElement(t *testing.T) {
+	t.Parallel()
+
+	tests := []struct {
+		name  string
+		start cty.Value
+		index int
+		value cty.Value
+		want  cty.Value
+	}{
+		{
+			name:  "empty",
+			start: cty.Value{},
+			index: 0,
+			value: cty.NilVal,
+			want:  cty.TupleVal([]cty.Value{cty.NilVal}),
+		},
+		{
+			name:  "empty to length",
+			start: cty.Value{},
+			index: 2,
+			value: cty.NilVal,
+			want:  cty.TupleVal([]cty.Value{cty.NilVal, cty.NilVal, cty.NilVal}),
+		},
+		{
+			name:  "insert to empty",
+			start: cty.EmptyTupleVal,
+			index: 1,
+			value: cty.NumberIntVal(5),
+			want:  cty.TupleVal([]cty.Value{cty.NilVal, cty.NumberIntVal(5)}),
+		},
+		{
+			name:  "insert to existing",
+			start: cty.TupleVal([]cty.Value{cty.NumberIntVal(1), cty.NumberIntVal(2), cty.NumberIntVal(3)}),
+			index: 1,
+			value: cty.NumberIntVal(5),
+			want:  cty.TupleVal([]cty.Value{cty.NumberIntVal(1), cty.NumberIntVal(5), cty.NumberIntVal(3)}),
+		},
+		{
+			name:  "insert to existing, extends",
+			start: cty.TupleVal([]cty.Value{cty.NumberIntVal(1), cty.NumberIntVal(2), cty.NumberIntVal(3)}),
+			index: 4,
+			value: cty.NumberIntVal(5),
+			want: cty.TupleVal([]cty.Value{
+				cty.NumberIntVal(1), cty.NumberIntVal(2),
+				cty.NumberIntVal(3), cty.NilVal,
+				cty.NumberIntVal(5),
+			}),
+		},
+		{
+			name:  "mixed list",
+			start: cty.TupleVal([]cty.Value{cty.StringVal("a"), cty.NumberIntVal(2), cty.NumberIntVal(3)}),
+			index: 1,
+			value: cty.BoolVal(true),
+			want: cty.TupleVal([]cty.Value{
+				cty.StringVal("a"), cty.BoolVal(true), cty.NumberIntVal(3),
+			}),
+		},
+		{
+			name:  "replace end",
+			start: cty.TupleVal([]cty.Value{cty.StringVal("a"), cty.NumberIntVal(2), cty.NumberIntVal(3)}),
+			index: 2,
+			value: cty.StringVal("end"),
+			want: cty.TupleVal([]cty.Value{
+				cty.StringVal("a"), cty.NumberIntVal(2), cty.StringVal("end"),
+			}),
+		},
+
+		// Some bad arguments
+		{
+			name:  "negative index",
+			start: cty.TupleVal([]cty.Value{cty.StringVal("a"), cty.NumberIntVal(2), cty.NumberIntVal(3)}),
+			index: -1,
+			value: cty.BoolVal(true),
+			want:  cty.TupleVal([]cty.Value{cty.StringVal("a"), cty.NumberIntVal(2), cty.NumberIntVal(3)}),
+		},
+		{
+			name:  "non-list",
+			start: cty.BoolVal(true),
+			index: 1,
+			value: cty.BoolVal(true),
+			want:  cty.TupleVal([]cty.Value{cty.NilVal, cty.BoolVal(true)}),
+		},
+	}
+
+	for _, tt := range tests {
+		tt := tt
+		t.Run(tt.name, func(t *testing.T) {
+			t.Parallel()
+
+			require.Equal(t, tt.want, insertTupleElement(tt.start, tt.index, tt.value))
+		})
+	}
+}

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

@@ -4,6 +4,7 @@ import (
 	"context"
 	"errors"
 	"io/fs"
+	"maps"
 	"reflect"
 	"slices"
 
@@ -521,7 +522,6 @@ func (e *evaluator) getValuesByBlockType(blockType string) cty.Value {
 	values := make(map[string]cty.Value)
 
 	for _, b := range blocksOfType {
-
 		switch b.Type() {
 		case "variable": // variables are special in that their value comes from the "default" attribute
 			val, err := e.evaluateVariable(b)
@@ -536,9 +536,7 @@ func (e *evaluator) getValuesByBlockType(blockType string) cty.Value {
 			}
 			values[b.Label()] = val
 		case "locals", "moved", "import":
-			for key, val := range b.Values().AsValueMap() {
-				values[key] = val
-			}
+			maps.Copy(values, b.Values().AsValueMap())
 		case "provider", "module", "check":
 			if b.Label() == "" {
 				continue
@@ -549,19 +547,27 @@ func (e *evaluator) getValuesByBlockType(blockType string) cty.Value {
 				continue
 			}
 
-			blockMap, ok := values[b.Labels()[0]]
+			// Data blocks should all be loaded into the top level 'values'
+			// object. The hierarchy of the map is:
+			//  values = map[<type>]map[<name>] =
+			//              Block -> Block's attributes as a cty.Object
+			//              Tuple(Block) -> Instances of the block
+			//              Object(Block) -> Field values are instances of the block
+			ref := b.Reference()
+			typeValues, ok := values[ref.TypeLabel()]
 			if !ok {
-				values[b.Labels()[0]] = cty.ObjectVal(make(map[string]cty.Value))
-				blockMap = values[b.Labels()[0]]
+				typeValues = cty.ObjectVal(make(map[string]cty.Value))
+				values[ref.TypeLabel()] = typeValues
 			}
 
-			valueMap := blockMap.AsValueMap()
+			valueMap := typeValues.AsValueMap()
 			if valueMap == nil {
 				valueMap = make(map[string]cty.Value)
 			}
+			valueMap[ref.NameLabel()] = blockInstanceValues(b, valueMap)
 
-			valueMap[b.Labels()[1]] = b.Values()
-			values[b.Labels()[0]] = cty.ObjectVal(valueMap)
+			// Update the map of all blocks with the same type.
+			values[ref.TypeLabel()] = cty.ObjectVal(valueMap)
 		}
 	}
 
@@ -572,23 +578,57 @@ func (e *evaluator) getResources() map[string]cty.Value {
 	values := make(map[string]map[string]cty.Value)
 
 	for _, b := range e.blocks {
-		if b.Type() != "resource" {
-			continue
-		}
-
-		if len(b.Labels()) < 2 {
+		if b.Type() != "resource" || len(b.Labels()) < 2 {
 			continue
 		}
 
-		val, exists := values[b.Labels()[0]]
+		ref := b.Reference()
+		typeValues, exists := values[ref.TypeLabel()]
 		if !exists {
-			val = make(map[string]cty.Value)
-			values[b.Labels()[0]] = val
+			typeValues = make(map[string]cty.Value)
+			values[ref.TypeLabel()] = typeValues
 		}
-		val[b.Labels()[1]] = b.Values()
+		typeValues[ref.NameLabel()] = blockInstanceValues(b, typeValues)
 	}
 
 	return lo.MapValues(values, func(v map[string]cty.Value, _ string) cty.Value {
 		return cty.ObjectVal(v)
 	})
 }
+
+// blockInstanceValues returns a cty.Value containing the values of the block instances.
+// If the count argument is used, a tuple is returned where the index corresponds to the argument index.
+// If the for_each argument is used, an object is returned where the key corresponds to the argument key.
+// In other cases, the values of the block itself are returned.
+func blockInstanceValues(b *terraform.Block, typeValues map[string]cty.Value) cty.Value {
+	ref := b.Reference()
+	key := ref.RawKey()
+
+	switch {
+	case key.Type().Equals(cty.Number) && b.GetAttribute("count") != nil:
+		idx, _ := key.AsBigFloat().Int64()
+		return insertTupleElement(typeValues[ref.NameLabel()], int(idx), b.Values())
+	case isForEachKey(key) && b.GetAttribute("for_each") != nil:
+		keyStr := ref.Key()
+
+		instancesVal, exists := typeValues[ref.NameLabel()]
+		if !exists || !instancesVal.CanIterateElements() {
+			instancesVal = cty.EmptyObjectVal
+		}
+
+		instances := instancesVal.AsValueMap()
+		if instances == nil {
+			instances = make(map[string]cty.Value)
+		}
+
+		instances[keyStr] = b.Values()
+		return cty.ObjectVal(instances)
+
+	default:
+		return b.Values()
+	}
+}
+
+func isForEachKey(key cty.Value) bool {
+	return key.Type().Equals(cty.Number) || key.Type().Equals(cty.String)
+}

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

@@ -1704,6 +1704,102 @@ resource "test_resource" "this" {
 	assert.Equal(t, "test_value", attr.GetRawValue())
 }
 
+func TestPopulateContextWithBlockInstances(t *testing.T) {
+
+	tests := []struct {
+		name  string
+		files map[string]string
+	}{
+		{
+			name: "data blocks with count",
+			files: map[string]string{
+				"main.tf": `data "d" "foo" {
+  count = 1
+  value = "Index ${count.index}"
+}
+
+data "b" "foo" {
+  count = 1
+  value = data.d.foo[0].value
+}
+
+data "c" "foo" {
+  count = 1
+  value = data.b.foo[0].value
+}`,
+			},
+		},
+		{
+			name: "resource blocks with count",
+			files: map[string]string{
+				"main.tf": `resource "d" "foo" {
+  count = 1
+  value = "Index ${count.index}"
+}
+
+resource "b" "foo" {
+  count = 1
+  value = d.foo[0].value
+}
+
+resource "c" "foo" {
+  count = 1
+  value = b.foo[0].value
+}`,
+			},
+		},
+		{
+			name: "data blocks with for_each",
+			files: map[string]string{
+				"main.tf": `data "d" "foo" {
+  for_each = toset([0])
+  value = "Index ${each.key}"
+}
+
+data "b" "foo" {
+  for_each = data.d.foo 
+  value = each.value.value
+}
+
+data "c" "foo" {
+  for_each = data.b.foo 
+  value = each.value.value
+}`,
+			},
+		},
+		{
+			name: "resource blocks with for_each",
+			files: map[string]string{
+				"main.tf": `resource "d" "foo" {
+  for_each = toset([0])
+  value = "Index ${each.key}"
+}
+
+resource "b" "foo" {
+  for_each = d.foo 
+  value = each.value.value
+}
+
+resource "c" "foo" {
+  for_each = b.foo 
+  value = each.value.value
+}`,
+			},
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			modules := parse(t, tt.files)
+			require.Len(t, modules, 1)
+			for _, b := range modules.GetBlocks() {
+				attr := b.GetAttribute("value")
+				assert.Equal(t, "Index 0", attr.Value().AsString())
+			}
+		})
+	}
+}
+
 // TestNestedModulesOptions ensures parser options are carried to the nested
 // submodule evaluators.
 // The test will include an invalid module that will fail to download