coder · Emyrk · Mar 19, 2025 · Mar 19, 2025 · Mar 19, 2025 · Mar 19, 2025
diff --git a/pkg/iac/scanners/terraform/parser/evaluator.go b/pkg/iac/scanners/terraform/parser/evaluator.go
@@ -38,6 +38,9 @@ type evaluator struct {
 	parentParser      *Parser
 	allowDownloads    bool
 	skipCachedModules bool
+	// stepHooks are functions that are called after each evaluation step.
+	// They can be used to provide additional semantics to other terraform blocks.
+	stepHooks []EvaluateStepHook
 }
 
 func newEvaluator(
@@ -55,6 +58,7 @@ func newEvaluator(
 	logger *log.Logger,
 	allowDownloads bool,
 	skipCachedModules bool,
+	stepHooks []EvaluateStepHook,
 ) *evaluator {
 
 	// create a context to store variables and make functions available
@@ -87,9 +91,12 @@ func newEvaluator(
 		logger:            logger,
 		allowDownloads:    allowDownloads,
 		skipCachedModules: skipCachedModules,
+		stepHooks:         stepHooks,
 	}
 }
 
+type EvaluateStepHook func(ctx *tfcontext.Context, blocks terraform.Blocks, inputVars map[string]cty.Value)
+
 func (e *evaluator) evaluateStep() {
 
 	e.ctx.Set(e.getValuesByBlockType("variable"), "var")
@@ -103,6 +110,10 @@ func (e *evaluator) evaluateStep() {
 	e.ctx.Set(e.getValuesByBlockType("data"), "data")
 	e.ctx.Set(e.getValuesByBlockType("output"), "output")
 	e.ctx.Set(e.getValuesByBlockType("module"), "module")
+
+	for _, hook := range e.stepHooks {
+		hook(e.ctx, e.blocks, e.inputVars)
+	}
 }
 
 // exportOutputs is used to export module outputs to the parent module
@@ -140,14 +151,20 @@ func (e *evaluator) EvaluateAll(ctx context.Context) (terraform.Modules, map[str
 	e.blocks = e.expandBlocks(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
 }
 
@@ -254,6 +271,9 @@ 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) {
@@ -401,8 +421,15 @@ func (e *evaluator) expandBlockCounts(blocks terraform.Blocks) terraform.Blocks
 			countFiltered = append(countFiltered, block)
 			continue
 		}
-		count := 1
+
 		countAttrVal := countAttr.Value()
+		if countAttrVal.IsNull() {
+			// 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())
 		}

diff --git a/pkg/iac/scanners/terraform/parser/option.go b/pkg/iac/scanners/terraform/parser/option.go
@@ -8,6 +8,12 @@ import (
 
 type Option func(p *Parser)
 
+func OptionWithEvalHook(hooks EvaluateStepHook) Option {
+	return func(p *Parser) {
+		p.stepHooks = append(p.stepHooks, hooks)
+	}
+}
+
 func OptionWithTFVarsPaths(paths ...string) Option {
 	return func(p *Parser) {
 		p.tfvarsPaths = paths

diff --git a/pkg/iac/scanners/terraform/parser/parser.go b/pkg/iac/scanners/terraform/parser/parser.go
@@ -51,6 +51,7 @@ type Parser struct {
 	fsMap             map[string]fs.FS
 	configsFS         fs.FS
 	skipPaths         []string
+	stepHooks         []EvaluateStepHook
 }
 
 // New creates a new Parser
@@ -66,6 +67,7 @@ func New(moduleFS fs.FS, moduleSource string, opts ...Option) *Parser {
 		configsFS:      moduleFS,
 		logger:         log.WithPrefix("terraform parser").With("module", "root"),
 		tfvars:         make(map[string]cty.Value),
+		stepHooks:      make([]EvaluateStepHook, 0),
 	}
 
 	for _, option := range opts {
@@ -311,6 +313,7 @@ func (p *Parser) Load(ctx context.Context) (*evaluator, error) {
 		log.WithPrefix("terraform evaluator"),
 		p.allowDownloads,
 		p.skipCachedModules,
+		p.stepHooks,
 	), nil
 }
 

diff --git a/pkg/iac/scanners/terraform/parser/parser_test.go b/pkg/iac/scanners/terraform/parser/parser_test.go
@@ -17,6 +17,7 @@ import (
 
 	"github.com/aquasecurity/trivy/internal/testutil"
 	"github.com/aquasecurity/trivy/pkg/iac/terraform"
+	tfcontext "github.com/aquasecurity/trivy/pkg/iac/terraform/context"
 	"github.com/aquasecurity/trivy/pkg/log"
 	"github.com/aquasecurity/trivy/pkg/set"
 )
@@ -1704,6 +1705,128 @@ resource "test_resource" "this" {
 	assert.Equal(t, "test_value", attr.GetRawValue())
 }
 
+func TestBlockCount(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
+}
+`,
+		"modules/foo/main.tf": `
+output "staticZero" {	
+	value = 0
+}
+output "staticFive" {	
+	value = 5
+}
+`,
+	}
+
+	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)
+}
+
+func TestBlockCountNested(t *testing.T) {
+	// `count` meta attributes are incorrectly handled when referencing
+	// a module output.
+	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 = 0' module 'bar' is still loaded")
+	// `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)
+}
+
 // TestNestedModulesOptions ensures parser options are carried to the nested
 // submodule evaluators.
 // The test will include an invalid module that will fail to download
@@ -2127,6 +2250,80 @@ func TestTFVarsFileDoesNotExist(t *testing.T) {
 	assert.ErrorContains(t, err, "file does not exist")
 }
 
+func Test_OptionsWithEvalHook(t *testing.T) {
+	fs := testutil.CreateFS(t, map[string]string{
+		"main.tf": `
+data "your_custom_data" "this" {
+  default = ["foo", "foh", "fum"]
+  unaffected = "bar"
+}
+
+// Testing the hook affects some value, which is used in another evaluateStep
+// action (expanding blocks)
+data "random_thing" "that" {
+  dynamic "repeated" {
+    for_each = data.your_custom_data.this.value
+	content {
+      value = repeated.value
+	}
+  }
+}
+
+locals {
+	referenced = data.your_custom_data.this.value
+	static_ref = data.your_custom_data.this.unaffected
+}
+`})
+
+	parser := New(fs, "", OptionWithEvalHook(
+		// A basic example of how to have a 'default' value for a data block.
+		// To see a more practical example, see how 'evaluateVariable' handles
+		// the 'default' value of a variable.
+		func(ctx *tfcontext.Context, blocks terraform.Blocks, inputVars map[string]cty.Value) {
+			dataBlocks := blocks.OfType("data")
+			for _, block := range dataBlocks {
+				if len(block.Labels()) >= 1 && block.Labels()[0] == "your_custom_data" {
+					def := block.GetAttribute("default")
+					ctx.Set(cty.ObjectVal(map[string]cty.Value{
+						"value": def.Value(),
+					}), "data", "your_custom_data", "this")
+				}
+			}
+
+		},
+	))
+
+	require.NoError(t, parser.ParseFS(context.TODO(), "."))
+
+	modules, _, err := parser.EvaluateAll(context.TODO())
+	require.NoError(t, err)
+	assert.Len(t, modules, 1)
+
+	rootModule := modules[0]
+
+	// Check the default value of the data block
+	blocks := rootModule.GetDatasByType("your_custom_data")
+	assert.Len(t, blocks, 1)
+	expList := cty.TupleVal([]cty.Value{cty.StringVal("foo"), cty.StringVal("foh"), cty.StringVal("fum")})
+	assert.True(t, expList.Equals(blocks[0].GetAttribute("default").Value()).True(), "default value matched list")
+	assert.Equal(t, "bar", blocks[0].GetAttribute("unaffected").Value().AsString())
+
+	// Check the referenced 'data.your_custom_data.this.value' exists in the eval
+	// context, and it is the default value of the data block.
+	locals := rootModule.GetBlocks().OfType("locals")
+	assert.Len(t, locals, 1)
+	assert.True(t, expList.Equals(locals[0].GetAttribute("referenced").Value()).True(), "referenced value matched list")
+	assert.Equal(t, "bar", locals[0].GetAttribute("static_ref").Value().AsString())
+
+	// Check the dynamic block is expanded correctly
+	dynamicBlocks := rootModule.GetDatasByType("random_thing")
+	assert.Len(t, dynamicBlocks, 1)
+	assert.Len(t, dynamicBlocks[0].GetBlocks("repeated"), 3)
+	for i, repeat := range dynamicBlocks[0].GetBlocks("repeated") {
+		assert.Equal(t, expList.Index(cty.NumberIntVal(int64(i))), repeat.GetAttribute("value").Value())
+	}
+}
+
 func Test_OptionsWithTfVars(t *testing.T) {
 	fs := testutil.CreateFS(t, map[string]string{
 		"main.tf": `resource "test" "this" {