fix(terraform): 'count' meta argument sourced from submodule output

Expand count blocks can depend on submodule returns

Do not expand unknown `count` blocks. Run `expandBlocks` in eval to
allow submodule returns to affect the `count` when using module
outputs.

(cherry picked from commit 494b512)

Author: Emyrk

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

@@ -162,14 +162,19 @@ func (e *evaluator) EvaluateAll(ctx context.Context) (terraform.Modules, map[str
 	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)
+	// A pointer for this module is needed up front to correctly set the module parent hierarchy.
+	// The actual instance is created at the end, when all terraform blocks
+	// are evaluated.
+	rootModule := new(terraform.Module)
 	submodules := e.evaluateSubmodules(ctx, rootModule, fsMap)
 
 	e.logger.Debug("Starting post-submodules evaluation...")
 	e.evaluateSteps()
 
 	e.logger.Debug("Module evaluation complete.")
+	// terraform.NewModule must be called at the end, as `e.blocks` can be
+	// changed up until the last moment.
+	*rootModule = *terraform.NewModule(e.projectRootPath, e.modulePath, e.blocks, e.ignores)
 	return append(terraform.Modules{rootModule}, submodules...), fsMap
 }
 
@@ -282,6 +287,10 @@ func (e *evaluator) evaluateSteps() {
 		e.logger.Debug("Starting iteration", log.Int("iteration", i))
 		e.evaluateStep()
 
+		// Always attempt to expand any blocks that might now be expandable
+		// due to new context being set.
+		e.blocks = e.expandBlocks(e.blocks)
+
 		// 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.logger.Debug("Context unchanged", log.Int("iteration", i))
@@ -330,8 +339,14 @@ func (e *evaluator) expandBlockForEaches(blocks terraform.Blocks) terraform.Bloc
 		}
 
 		forEachVal := forEachAttr.Value()
+		if !forEachVal.IsKnown() {
+			// Defer the expansion of the block if it is unknown. It might be known at a later
+			// execution step.
+			forEachFiltered = append(forEachFiltered, block)
+			continue
+		}
 
-		if forEachVal.IsNull() || !forEachVal.IsKnown() || !forEachAttr.IsIterable() {
+		if forEachVal.IsNull() || !forEachAttr.IsIterable() {
 			e.logger.Debug(`Failed to expand block. Invalid "for-each" argument. Must be known and iterable.`,
 				log.String("block", block.FullName()),
 				log.String("value", forEachVal.GoString()),
@@ -426,8 +441,15 @@ func (e *evaluator) expandBlockCounts(blocks terraform.Blocks) terraform.Blocks
 			countFiltered = append(countFiltered, block)
 			continue
 		}
-		count := 1
+
 		countAttrVal := countAttr.Value()
+		if !countAttrVal.IsKnown() {
+			// Defer to the next pass when the count might be known
+			countFiltered = append(countFiltered, block)
+			continue
+		}
+
+		count := 1
 		if !countAttrVal.IsNull() && countAttrVal.IsKnown() && countAttrVal.Type() == cty.Number {
 			count = int(countAttr.AsNumber())
 		}

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

@@ -0,0 +1,156 @@
+package parser
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestBlockExpandWithSubmoduleOutput(t *testing.T) {
+	// `count` meta attributes are incorrectly handled when referencing
+	// a module output.
+	files := map[string]string{
+		"main.tf": `
+module "foo" {
+	source = "./modules/foo"
+}
+data "this_resource" "this" {
+	count = module.foo.staticZero
+}
+data "that_resource" "this" {
+	count = module.foo.staticFive
+}
+
+data "for_each_resource_empty" "this" {
+	for_each = module.foo.empty_list
+}
+data "for_each_resource_abc" "this" {
+	for_each = module.foo.list_abc
+}
+
+data "dynamic_block" "that" {
+	dynamic "element" {
+		for_each = module.foo.list_abc
+		content {
+			foo = element.value
+		}
+	}
+}
+`,
+		"modules/foo/main.tf": `
+output "staticZero" {	
+	value = 0
+}
+output "staticFive" {	
+	value = 5
+}
+
+output "empty_list" {
+	value = []
+}
+output "list_abc" {
+	value = ["a", "b", "c"]
+}
+`,
+	}
+
+	modules := parse(t, files)
+	require.Len(t, modules, 2)
+
+	datas := modules.GetDatasByType("this_resource")
+	require.Empty(t, datas)
+
+	datas = modules.GetDatasByType("that_resource")
+	require.Len(t, datas, 5)
+
+	datas = modules.GetDatasByType("for_each_resource_empty")
+	require.Empty(t, datas)
+
+	datas = modules.GetDatasByType("for_each_resource_abc")
+	require.Len(t, datas, 3)
+
+	dyn := modules.GetDatasByType("dynamic_block")
+	require.Len(t, dyn, 1)
+	require.Len(t, dyn[0].GetBlocks("element"), 3, "dynamic expand")
+}
+
+func TestBlockExpandWithSubmoduleOutputNested(t *testing.T) {
+	files := map[string]string{
+		"main.tf": `
+module "alpha" {
+  source = "./nestedcount"
+  set_count = 2
+}
+module "beta" {
+  source = "./nestedcount"
+  set_count = module.alpha.set_count
+}
+module "charlie" {
+  count = module.beta.set_count - 1
+  source = "./nestedcount"
+  set_count = module.beta.set_count
+}
+data "repeatable" "foo" {
+  count = module.charlie[0].set_count
+  value = "foo"
+}
+`,
+		"setcount/main.tf": `
+variable "set_count" {
+    type = number
+}
+output "set_count" {
+  value = var.set_count
+}
+`,
+		"nestedcount/main.tf": `
+variable "set_count" {
+  type = number
+}
+module "nested_mod" {
+  source = "../setcount"
+  set_count = var.set_count
+}
+output "set_count" {
+  value = module.nested_mod.set_count
+}
+`,
+	}
+
+	modules := parse(t, files)
+	require.Len(t, modules, 7)
+
+	datas := modules.GetDatasByType("repeatable")
+	assert.Len(t, datas, 2)
+}
+
+func TestBlockCountModules(t *testing.T) {
+	t.Skip(
+		"This test is currently failing. " +
+			"The count passed to `module bar` is not being set correctly. " +
+			"The count value is sourced from the output of `module foo`. " +
+			"Submodules cannot be dependent on the output of other submodules right now. ",
+	)
+	// `count` meta attributes are incorrectly handled when referencing
+	// a module output.
+	files := map[string]string{
+		"main.tf": `
+module "foo" {
+	source = "./modules/foo"
+}
+module "bar" {
+	source = "./modules/foo"
+    count = module.foo.staticZero
+}
+`,
+		"modules/foo/main.tf": `
+output "staticZero" {	
+	value = 0
+}
+`,
+	}
+
+	modules := parse(t, files)
+	require.Len(t, modules, 2)
+}

pkg/iac/terraform/block.go

@@ -620,7 +620,7 @@ func (b *Block) expandDynamic() ([]*Block, error) {
 		return nil, fmt.Errorf("invalid for-each in %s block: %w", b.FullLocalName(), err)
 	}
 
-	if !forEachVal.IsKnown() {
+	if !forEachVal.IsWhollyKnown() {
 		return nil, errors.New("for-each must be known")
 	}