fix(terraform): hcl object expressions to return references (#8271)

Signed-off-by: nikpivkin <nikita.pivkin@smartforce.io>
Co-authored-by: nikpivkin <nikita.pivkin@smartforce.io>
Co-authored-by: simar7 <1254783+simar7@users.noreply.github.com>
Co-authored-by: Simar <simar@linux.com>

pkg/iac/adapters/terraform/aws/lambda/adapt.go

@@ -86,8 +86,8 @@ func (a *adapter) adaptPermission(permission *terraform.Block) lambda.Permission
 	sourceARNAttr := permission.GetAttribute("source_arn")
 	sourceARN := sourceARNAttr.AsStringValueOrDefault("", permission)
 
-	if len(sourceARNAttr.AllReferences()) > 0 {
-		sourceARN = iacTypes.String(sourceARNAttr.AllReferences()[0].NameLabel(), sourceARNAttr.GetMetadata())
+	if refs := sourceARNAttr.AllReferences(); len(refs) > 0 {
+		sourceARN = iacTypes.String(refs[0].NameLabel(), sourceARNAttr.GetMetadata())
 	}
 
 	return lambda.Permission{

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

@@ -2559,6 +2559,38 @@ resource "aws_s3_bucket" "example" {
 	assert.Nil(t, val)
 }
 
+func Test_AttrIsRefToOtherBlock(t *testing.T) {
+	fsys := testutil.CreateFS(t, map[string]string{
+		"main.tf": `locals {
+  baz_idx = 0
+}
+
+resource "foo" "bar" {
+  attr = baz.qux[local.baz_idx].attr
+}
+
+resource "baz" "qux" {
+  count = 1
+  attr  = "test"
+}
+`,
+	})
+
+	parser := New(fsys, "", OptionStopOnHCLError(true))
+	require.NoError(t, parser.ParseFS(t.Context(), "."))
+
+	modules, err := parser.EvaluateAll(t.Context())
+	require.NoError(t, err)
+
+	require.Len(t, modules, 1)
+	root := modules[0]
+	foo := root.GetResourcesByType("foo")[0]
+	fooAttr := foo.GetAttribute("attr")
+	b, err := root.GetReferencedBlock(fooAttr, foo)
+	require.NoError(t, err)
+	require.NotNil(t, b)
+}
+
 func TestConfigWithEphemeralBlock(t *testing.T) {
 	fsys := fstest.MapFS{
 		"main.tf": &fstest.MapFile{Data: []byte(`ephemeral "random_password" "password" {

pkg/iac/terraform/attribute.go

@@ -796,7 +796,7 @@ func (a *Attribute) AllReferences(blocks ...*Block) []*Reference {
 	if a == nil {
 		return nil
 	}
-	refs := a.extractReferences()
+	refs := a.referencesFromExpression(a.hclAttribute.Expr)
 	for _, block := range blocks {
 		for _, ref := range refs {
 			if ref.TypeLabel() == "each" {
@@ -809,68 +809,31 @@ func (a *Attribute) AllReferences(blocks ...*Block) []*Reference {
 	return refs
 }
 
-// nolint
-func (a *Attribute) referencesFromExpression(expression hcl.Expression) []*Reference {
-	var refs []*Reference
-	switch t := expression.(type) {
-	case *hclsyntax.ConditionalExpr:
-		if ref, err := createDotReferenceFromTraversal(a.module, t.TrueResult.Variables()...); err == nil {
-			refs = append(refs, ref)
+func (a *Attribute) referencesFromExpression(expr hcl.Expression) []*Reference {
+	if reflect.TypeOf(expr).String() == "*json.expression" {
+		if ref, err := createDotReferenceFromTraversal(a.module, expr.Variables()...); err == nil {
+			return []*Reference{ref}
 		}
-		if ref, err := createDotReferenceFromTraversal(a.module, t.FalseResult.Variables()...); err == nil {
-			refs = append(refs, ref)
-		}
-		if ref, err := createDotReferenceFromTraversal(a.module, t.Condition.Variables()...); err == nil {
-			refs = append(refs, ref)
-		}
-	case *hclsyntax.ScopeTraversalExpr:
-		if ref, err := createDotReferenceFromTraversal(a.module, t.Variables()...); err == nil {
-			refs = append(refs, ref)
-		}
-	case *hclsyntax.TemplateWrapExpr:
-		refs = a.referencesFromExpression(t.Wrapped)
-	case *hclsyntax.TemplateExpr:
-		for _, part := range t.Parts {
-			ref, err := createDotReferenceFromTraversal(a.module, part.Variables()...)
-			if err != nil {
-				continue
-			}
-			refs = append(refs, ref)
-		}
-	case *hclsyntax.TupleConsExpr:
-		for _, v := range t.Variables() {
-			if ref, err := createDotReferenceFromTraversal(a.module, v); err == nil {
-				refs = append(refs, ref)
-			}
-		}
-	case *hclsyntax.RelativeTraversalExpr:
-		switch s := t.Source.(type) {
-		case *hclsyntax.IndexExpr:
-			if collectionRef, err := createDotReferenceFromTraversal(a.module, s.Collection.Variables()...); err == nil {
-				key, _ := s.Key.Value(a.ctx.Inner())
-				collectionRef.SetKey(key)
-				refs = append(refs, collectionRef)
-			}
-		default:
-			if ref, err := createDotReferenceFromTraversal(a.module, t.Source.Variables()...); err == nil {
-				refs = append(refs, ref)
-			}
-		}
-	default:
-		if reflect.TypeOf(expression).String() == "*json.expression" {
-			if ref, err := createDotReferenceFromTraversal(a.module, expression.Variables()...); err == nil {
-				refs = append(refs, ref)
-			}
-		}
-	}
-	return refs
-}
-
-func (a *Attribute) extractReferences() []*Reference {
-	if a == nil {
 		return nil
 	}
-	return a.referencesFromExpression(a.hclAttribute.Expr)
+
+	vars := expr.Variables()
+	refs := make([]*Reference, 0, len(vars))
+	for _, v := range vars {
+		ref, err := createDotReferenceFromTraversal(a.module, v)
+		if err != nil {
+			continue
+		}
+
+		if relExpr, ok := expr.(*hclsyntax.RelativeTraversalExpr); ok {
+			if idxExpr, ok := relExpr.Source.(*hclsyntax.IndexExpr); ok {
+				key, _ := idxExpr.Key.Value(a.ctx.Inner())
+				ref.SetKey(key)
+			}
+		}
+		refs = append(refs, ref)
+	}
+	return refs
 }
 
 func (a *Attribute) IsResourceBlockReference(resourceType string) bool {

pkg/iac/terraform/attribute_test.go

@@ -0,0 +1,132 @@
+package terraform
+
+import (
+	"testing"
+
+	"github.com/hashicorp/hcl/v2"
+	"github.com/hashicorp/hcl/v2/hclsyntax"
+	"github.com/hashicorp/hcl/v2/json"
+	"github.com/stretchr/testify/require"
+
+	"github.com/aquasecurity/trivy/pkg/iac/terraform/context"
+	"github.com/aquasecurity/trivy/pkg/iac/types"
+)
+
+func Test_AllReferences(t *testing.T) {
+	cases := []struct {
+		input string
+		refs  []string
+	}{
+		{
+			input: "42", // literal
+			refs:  []string{},
+		},
+		{
+			input: "5 == 5", // comparison
+			refs:  []string{},
+		},
+		{
+			input: "var.foo",
+			refs:  []string{"variable.foo"},
+		},
+		{
+			input: "resource.foo.bar[local.idx].name",
+			refs:  []string{"foo.bar", "locals.idx"},
+		},
+		{
+			input: "resource.foo.bar[0].name",
+			refs:  []string{"foo.bar[0].name"},
+		},
+		{
+			input: "resource.aws_instance.id",
+			refs:  []string{"aws_instance.id"},
+		},
+		{
+			input: "data.aws_ami.ubuntu.most_recent",
+			refs:  []string{"data.aws_ami.ubuntu.most_recent"},
+		},
+		{
+			input: "5 == 5 ? var.foo : data.aws_ami.ubuntu.most_recent", // conditional
+			refs:  []string{"variable.foo", "data.aws_ami.ubuntu.most_recent"},
+		},
+		{
+			input: `{x = 1, y = data.aws_ami.ubuntu.most_recent}`,
+			refs:  []string{"data.aws_ami.ubuntu.most_recent"},
+		},
+		{
+			input: `{foo = 1 == 1 ? var.bar : data.aws_ami.ubuntu.most_recent}`,
+			refs:  []string{"variable.bar", "data.aws_ami.ubuntu.most_recent"},
+		},
+		{
+			input: `[var.foo, var.bar]`,
+			refs:  []string{"variable.foo", "variable.bar"},
+		},
+		{
+			// Expression in the key
+			input: `{(local.foo): local.bar}`,
+			refs:  []string{"locals.foo", "locals.bar"},
+		},
+	}
+
+	for _, test := range cases {
+		t.Run(test.input, func(t *testing.T) {
+			ctx := context.NewContext(&hcl.EvalContext{}, nil)
+
+			exp, diag := hclsyntax.ParseExpression([]byte(test.input), "", hcl.Pos{Line: 1, Column: 1})
+			if diag.HasErrors() {
+				require.NoError(t, diag)
+			}
+
+			a := NewAttribute(&hcl.Attribute{
+				Name:      "test",
+				Expr:      exp,
+				Range:     hcl.Range{},
+				NameRange: hcl.Range{},
+			}, ctx, "", types.Metadata{}, Reference{}, "", nil)
+
+			refs := a.AllReferences()
+			humanRefs := make([]string, 0, len(refs))
+			for _, ref := range refs {
+				humanRefs = append(humanRefs, ref.HumanReadable())
+			}
+
+			require.ElementsMatch(t, test.refs, humanRefs)
+		})
+	}
+}
+
+func Test_AllReferences_JSON(t *testing.T) {
+	tests := []struct {
+		src      string
+		expected []string
+	}{
+		{
+			src:      `"hello ${noun}"`,
+			expected: []string{"noun"},
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.src, func(t *testing.T) {
+			expr, diag := json.ParseExpression([]byte(tt.src), "")
+			if diag.HasErrors() {
+				require.NoError(t, diag)
+			}
+
+			attr := NewAttribute(&hcl.Attribute{
+				Name:      "test",
+				Expr:      expr,
+				Range:     hcl.Range{},
+				NameRange: hcl.Range{},
+			}, context.NewContext(&hcl.EvalContext{}, nil), "", types.Metadata{}, Reference{}, "", nil)
+
+			refs := attr.AllReferences()
+			humanRefs := make([]string, 0, len(refs))
+			for _, ref := range refs {
+				humanRefs = append(humanRefs, ref.HumanReadable())
+			}
+
+			require.ElementsMatch(t, tt.expected, humanRefs)
+		})
+	}
+}