fix: make operations with reference key safe (#1425)

aquasecurity · Aug 30, 2023 · 6fa220e · 6fa220e
1 parent 16e71c7
commit 6fa220e
Show file tree

Hide file tree

Showing 4 changed files with 144 additions and 19 deletions.
diff --git a/pkg/terraform/attribute.go b/pkg/terraform/attribute.go
@@ -884,9 +884,7 @@ func createDotReferenceFromTraversal(parentRef string, traversals ...hcl.Travers
 	if err != nil {
 		return nil, err
 	}
-	if !key.IsNull() {
-		ref.SetKey(key)
-	}
+	ref.SetKey(key)
 	return ref, nil
 }
 

diff --git a/pkg/terraform/block.go b/pkg/terraform/block.go
@@ -74,9 +74,7 @@ func NewBlock(hclBlock *hcl.Block, ctx *context.Context, moduleBlock *Block, par
 	ref, _ := newReference(parts, parent)
 	if len(index) > 0 {
 		key := index[0]
-		if !key.IsNull() {
-			ref.SetKey(key)
-		}
+		ref.SetKey(key)
 	}
 
 	metadata := defsecTypes.NewMetadata(rng, ref.String())

diff --git a/pkg/terraform/reference.go b/pkg/terraform/reference.go
@@ -114,8 +114,7 @@ func (r Reference) String() string {
 	return base
 }
 
-func (r Reference) RefersTo(a Reference) bool {
-	other := a
+func (r Reference) RefersTo(other Reference) bool {
 
 	if r.BlockType() != other.BlockType() {
 		return false
@@ -133,33 +132,43 @@ func (r Reference) RefersTo(a Reference) bool {
 }
 
 func (r *Reference) SetKey(key cty.Value) {
+	if key.IsNull() || !key.IsKnown() {
+		return
+	}
 	r.key = key
 }
+
 func (r Reference) KeyBracketed() string {
-	if r.key.IsNull() || !r.key.IsKnown() {
-		return ""
-	}
-	switch r.key.Type() {
-	case cty.Number:
-		f := r.key.AsBigFloat()
-		f64, _ := f.Float64()
-		return fmt.Sprintf("[%d]", int(f64))
-	case cty.String:
-		return fmt.Sprintf("[%q]", r.key.AsString())
+	switch v := key(r).(type) {
+	case int:
+		return fmt.Sprintf("[%d]", v)
+	case string:
+		if v == "" {
+			return ""
+		}
+		return fmt.Sprintf("[%q]", v)
 	default:
 		return ""
 	}
 }
+
 func (r Reference) RawKey() cty.Value {
 	return r.key
 }
 
 func (r Reference) Key() string {
+	return fmt.Sprintf("%v", key(r))
+}
+
+func key(r Reference) interface{} {
+	if r.key.IsNull() || !r.key.IsKnown() {
+		return ""
+	}
 	switch r.key.Type() {
 	case cty.Number:
 		f := r.key.AsBigFloat()
 		f64, _ := f.Float64()
-		return fmt.Sprintf("%d", int(f64))
+		return int(f64)
 	case cty.String:
 		return r.key.AsString()
 	default:

diff --git a/pkg/terraform/reference_test.go b/pkg/terraform/reference_test.go
@@ -4,6 +4,8 @@ import (
 	"testing"
 
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"github.com/zclconf/go-cty/cty"
 )
 
 func Test_ReferenceParsing(t *testing.T) {
@@ -49,3 +51,121 @@ func Test_ReferenceParsing(t *testing.T) {
 		})
 	}
 }
+
+func Test_SetKey(t *testing.T) {
+	tests := []struct {
+		name string
+		key  cty.Value
+		want cty.Value
+	}{
+		{
+			name: "happy",
+			key:  cty.StringVal("str"),
+			want: cty.StringVal("str"),
+		},
+		{
+			name: "null key",
+			key:  cty.NullVal(cty.String),
+			want: cty.Value{},
+		},
+		{
+			name: "unknown key",
+			key:  cty.UnknownVal(cty.String),
+			want: cty.Value{},
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			t.Run(tt.name, func(t *testing.T) {
+				ref, err := newReference([]string{"resource", "test"}, "")
+				require.NoError(t, err)
+
+				ref.SetKey(tt.key)
+
+				assert.Equal(t, tt.want, ref.RawKey())
+			})
+		})
+	}
+}
+
+func Test_Key(t *testing.T) {
+
+	tests := []struct {
+		name string
+		key  cty.Value
+		want string
+	}{
+		{
+			name: "empty key",
+			want: "",
+		},
+		{
+			name: "str key",
+			key:  cty.StringVal("some_value"),
+			want: "some_value",
+		},
+		{
+			name: "number key",
+			key:  cty.NumberIntVal(122),
+			want: "122",
+		},
+		{
+			name: "bool key",
+			key:  cty.BoolVal(true),
+			want: "",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			t.Run(tt.name, func(t *testing.T) {
+				ref, err := newReference([]string{"resource", "test"}, "")
+				require.NoError(t, err)
+
+				ref.SetKey(tt.key)
+
+				assert.Equal(t, tt.want, ref.Key())
+			})
+		})
+	}
+}
+
+func Test_KeyBracketed(t *testing.T) {
+	tests := []struct {
+		name string
+		key  cty.Value
+		want string
+	}{
+		{
+			name: "empty key",
+			want: "",
+		},
+		{
+			name: "str key",
+			key:  cty.StringVal("some_value"),
+			want: "[\"some_value\"]",
+		},
+		{
+			name: "number key",
+			key:  cty.NumberIntVal(122),
+			want: "[122]",
+		},
+		{
+			name: "bool key",
+			key:  cty.BoolVal(true),
+			want: "",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			ref, err := newReference([]string{"resource", "test"}, "")
+			require.NoError(t, err)
+
+			ref.SetKey(tt.key)
+
+			assert.Equal(t, tt.want, ref.KeyBracketed())
+		})
+	}
+}