perf(misconf): optimize work with context (#6968)

nikpivkin · web-flow · commit 2b6d8d9227fb · 2024-08-08T22:09:36.000Z
Signed-off-by: nikpivkin &lt;nikita.pivkin@smartforce.io&gt;
diff --git a/pkg/iac/scanners/terraform/parser/evaluator.go b/pkg/iac/scanners/terraform/parser/evaluator.go
@@ -96,9 +96,8 @@ func (e *evaluator) evaluateStep() {
 	e.ctx.Set(e.getValuesByBlockType("locals"), "local")
 	e.ctx.Set(e.getValuesByBlockType("provider"), "provider")
 
-	resources := e.getValuesByBlockType("resource")
-	for key, resource := range resources.AsValueMap() {
-		e.ctx.Set(resource, key)
+	for typ, resource := range e.getResources() {
+		e.ctx.Set(resource, typ)
 	}
 
 	e.ctx.Set(e.getValuesByBlockType("data"), "data")
@@ -224,10 +223,12 @@ func (e *evaluator) evaluateSteps() {
 	var lastContext hcl.EvalContext
 	for i := 0; i < maxContextIterations; i++ {
 
+		e.debug.Log("Starting iteration %d", i)
 		e.evaluateStep()
 
 		// if ctx matches the last evaluation, we can bail, nothing left to resolve
 		if i > 0 && reflect.DeepEqual(lastContext.Variables, e.ctx.Inner().Variables) {
+			e.debug.Log("Context unchanged at i=%d", i)
 			break
 		}
 		if len(e.ctx.Inner().Variables) != len(lastContext.Variables) {
@@ -330,15 +331,16 @@ func (e *evaluator) expandBlockForEaches(blocks terraform.Blocks, isDynamic bool
 			}
 
 			clone := block.Clone(idx)
-
 			ctx := clone.Context()
-
 			e.copyVariables(block, clone)
 
-			ctx.SetByDot(idx, "each.key")
-			ctx.SetByDot(val, "each.value")
-			ctx.Set(idx, block.TypeLabel(), "key")
-			ctx.Set(val, block.TypeLabel(), "value")
+			eachObj := cty.ObjectVal(map[string]cty.Value{
+				"key":   idx,
+				"value": val,
+			})
+
+			ctx.Set(eachObj, "each")
+			ctx.Set(eachObj, block.TypeLabel())
 
 			if isDynamic {
 				if iterAttr := block.GetAttribute("iterator"); iterAttr.IsNotNil() {
@@ -354,9 +356,7 @@ func (e *evaluator) expandBlockForEaches(blocks terraform.Blocks, isDynamic bool
 
 			forEachFiltered = append(forEachFiltered, clone)
 
-			values := clone.Values()
-			clones[idx.AsString()] = values
-			e.ctx.SetByDot(values, clone.GetMetadata().Reference())
+			clones[idx.AsString()] = clone.Values()
 		})
 
 		metadata := block.GetMetadata()
@@ -434,11 +434,12 @@ func (e *evaluator) copyVariables(from, to *terraform.Block) {
 		return
 	}
 
-	srcValue := e.ctx.Root().Get(fromBase, fromRel)
+	rootCtx := e.ctx.Root()
+	srcValue := rootCtx.Get(fromBase, fromRel)
 	if srcValue == cty.NilVal {
 		return
 	}
-	e.ctx.Root().Set(srcValue, fromBase, toRel)
+	rootCtx.Set(srcValue, fromBase, toRel)
 }
 
 func (e *evaluator) evaluateVariable(b *terraform.Block) (cty.Value, error) {
@@ -530,7 +531,7 @@ func (e *evaluator) getValuesByBlockType(blockType string) cty.Value {
 				continue
 			}
 			values[b.Label()] = b.Values()
-		case "resource", "data":
+		case "data":
 			if len(b.Labels()) < 2 {
 				continue
 			}
@@ -553,3 +554,28 @@ func (e *evaluator) getValuesByBlockType(blockType string) cty.Value {
 
 	return cty.ObjectVal(values)
 }
+
+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 {
+			continue
+		}
+
+		val, exists := values[b.Labels()[0]]
+		if !exists {
+			val = make(map[string]cty.Value)
+			values[b.Labels()[0]] = val
+		}
+		val[b.Labels()[1]] = b.Values()
+	}
+
+	return lo.MapValues(values, func(v map[string]cty.Value, _ string) cty.Value {
+		return cty.ObjectVal(v)
+	})
+}
diff --git a/pkg/iac/terraform/block.go b/pkg/iac/terraform/block.go
@@ -85,7 +85,7 @@ func NewBlock(hclBlock *hcl.Block, ctx *context.Context, moduleBlock *Block, par
 	}
 
 	b := Block{
-		id:           uuid.New().String(),
+		id:           uuid.NewString(),
 		context:      ctx,
 		hclBlock:     hclBlock,
 		moduleBlock:  moduleBlock,
@@ -446,6 +446,9 @@ func (b *Block) Attributes() map[string]*Attribute {
 func (b *Block) Values() cty.Value {
 	values := createPresetValues(b)
 	for _, attribute := range b.GetAttributes() {
+		if attribute.Name() == "for_each" {
+			continue
+		}
 		values[attribute.Name()] = attribute.Value()
 	}
 	return cty.ObjectVal(postProcessValues(b, values))
diff --git a/pkg/iac/terraform/context/context.go b/pkg/iac/terraform/context/context.go
@@ -46,17 +46,29 @@ func (c *Context) Get(parts ...string) cty.Value {
 	if len(parts) == 0 {
 		return cty.NilVal
 	}
-	src := c.ctx.Variables
-	for i, part := range parts {
-		if i == len(parts)-1 {
-			return src[part]
+
+	curr := c.ctx.Variables[parts[0]]
+	if len(parts) == 1 {
+		return curr
+	}
+
+	for i, part := range parts[1:] {
+		if !curr.Type().HasAttribute(part) {
+			return cty.NilVal
+		}
+
+		attr := curr.GetAttr(part)
+
+		if i == len(parts)-2 { // iteration from the first element
+			return attr
 		}
-		nextPart := src[part]
-		if nextPart == cty.NilVal {
+
+		if !(attr.IsKnown() && attr.Type().IsObjectType()) {
 			return cty.NilVal
 		}
-		src = nextPart.AsValueMap()
+		curr = attr
 	}
+
 	return cty.NilVal
 }
 
@@ -97,13 +109,12 @@ func mergeVars(src cty.Value, parts []string, value cty.Value) cty.Value {
 	}
 
 	data := make(map[string]cty.Value)
-	if src.Type().IsObjectType() && !src.IsNull() && src.LengthInt() > 0 {
+	if isNotEmptyObject(src) {
 		data = src.AsValueMap()
-		tmp, ok := src.AsValueMap()[parts[0]]
-		if !ok {
-			src = cty.ObjectVal(make(map[string]cty.Value))
+		if attr, ok := data[parts[0]]; ok {
+			src = attr
 		} else {
-			src = tmp
+			src = cty.EmptyObjectVal
 		}
 	}
 
@@ -118,14 +129,16 @@ func mergeObjects(a, b cty.Value) cty.Value {
 	for key, val := range a.AsValueMap() {
 		output[key] = val
 	}
-	for key, val := range b.AsValueMap() {
-		old, exists := output[key]
-		if exists && isNotEmptyObject(old) && isNotEmptyObject(val) {
-			output[key] = mergeObjects(old, val)
+	b.ForEachElement(func(key, val cty.Value) (stop bool) {
+		k := key.AsString()
+		old := output[k]
+		if old.IsKnown() && isNotEmptyObject(old) && isNotEmptyObject(val) {
+			output[k] = mergeObjects(old, val)
 		} else {
-			output[key] = val
+			output[k] = val
 		}
-	}
+		return false
+	})
 	return cty.ObjectVal(output)
 }
 
diff --git a/pkg/iac/terraform/context/context_test.go b/pkg/iac/terraform/context/context_test.go
@@ -52,6 +52,43 @@ func Test_ContextVariablesPreservation(t *testing.T) {
 
 }
 
+func Test_SetWithMerge(t *testing.T) {
+	hctx := hcl.EvalContext{
+		Variables: map[string]cty.Value{
+			"my": cty.ObjectVal(map[string]cty.Value{
+				"someValue": cty.ObjectVal(map[string]cty.Value{
+					"foo": cty.StringVal("test"),
+					"bar": cty.ObjectVal(map[string]cty.Value{
+						"foo": cty.StringVal("test"),
+					}),
+				}),
+			}),
+		},
+	}
+
+	ctx := NewContext(&hctx, nil)
+
+	val := cty.ObjectVal(map[string]cty.Value{
+		"foo2": cty.StringVal("test2"),
+		"bar": cty.ObjectVal(map[string]cty.Value{
+			"foo2": cty.StringVal("test2"),
+		}),
+	})
+
+	ctx.Set(val, "my", "someValue")
+	got := ctx.Get("my", "someValue")
+	expected := cty.ObjectVal(map[string]cty.Value{
+		"foo":  cty.StringVal("test"),
+		"foo2": cty.StringVal("test2"),
+		"bar": cty.ObjectVal(map[string]cty.Value{
+			"foo":  cty.StringVal("test"),
+			"foo2": cty.StringVal("test2"),
+		}),
+	})
+
+	assert.Equal(t, expected, got)
+}
+
 func Test_ContextVariablesPreservationByDot(t *testing.T) {
 
 	underlying := &hcl.EvalContext{}
