chore(terraform): hook into evaluateStep behavior with custom hooks

evaluateStep handles many terraform semantics such as default
values for 'variables'. A hook into these steps allows defining
additional semantics, likely to mirror those of the actual provider
implementation.

Author: Emyrk

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

@@ -39,6 +39,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(
@@ -56,6 +59,7 @@ func newEvaluator(
 	logger *log.Logger,
 	allowDownloads bool,
 	skipCachedModules bool,
+	stepHooks []EvaluateStepHook,
 ) *evaluator {
 
 	// create a context to store variables and make functions available
@@ -88,9 +92,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")
@@ -104,6 +111,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

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

@@ -10,6 +10,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

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

@@ -55,6 +55,7 @@ type Parser struct {
 	// cwd is optional, if left to empty string, 'os.Getwd'
 	// will be used for populating 'path.cwd' in terraform.
 	cwd string
+	stepHooks         []EvaluateStepHook
 }
 
 // New creates a new Parser
@@ -70,6 +71,7 @@ func New(moduleFS fs.FS, moduleSource string, opts ...Option) *Parser {
 		configsFS:      moduleFS,
 		logger:         slog.Default(),
 		tfvars:         make(map[string]cty.Value),
+		stepHooks:      make([]EvaluateStepHook, 0),
 	}
 
 	for _, option := range opts {
@@ -322,6 +324,7 @@ func (p *Parser) Load(_ context.Context) (*evaluator, error) {
 		p.logger.With(log.Prefix("terraform evaluator")),
 		p.allowDownloads,
 		p.skipCachedModules,
+		p.stepHooks,
 	), nil
 }
 

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"
 )
@@ -2291,6 +2292,80 @@ func TestTFVarsFileDoesNotExist(t *testing.T) {
 	assert.ErrorContains(t, err, "file does not exist")
 }
 
+func Test_OptionsWithEvalHook(t *testing.T) {
+	fs := testutil.CreateFS(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(t.Context(), "."))
+
+	modules, err := parser.EvaluateAll(t.Context())
+	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(map[string]string{
 		"main.tf": `resource "test" "this" {