From d1d0ede069396f69e36c0d7da25e134f0bbae21f Mon Sep 17 00:00:00 2001
From: Kristin Laemmert <mildwonkey@users.noreply.github.com>
Date: Tue, 16 Oct 2018 12:09:04 -0700
Subject: [PATCH] lang/funcs: return default value if provided when object
 lookup fails to find attr

---
 lang/funcs/collection.go      | 26 ++++++++++++++------------
 lang/funcs/collection_test.go | 17 +++++++++++++++++
 2 files changed, 31 insertions(+), 12 deletions(-)

diff --git a/lang/funcs/collection.go b/lang/funcs/collection.go
index 8920cef595..499392f823 100644
--- a/lang/funcs/collection.go
+++ b/lang/funcs/collection.go
@@ -4,7 +4,6 @@ import (
 	"fmt"
 	"sort"
 
-	"github.com/hashicorp/hcl2/hcl"
 	"github.com/zclconf/go-cty/cty"
 	"github.com/zclconf/go-cty/cty/convert"
 	"github.com/zclconf/go-cty/cty/function"
@@ -497,12 +496,21 @@ var LookupFunc = function.New(&function.Spec{
 		if len(args) < 1 || len(args) > 3 {
 			return cty.NilType, fmt.Errorf("lookup() takes two or three arguments, got %d", len(args))
 		}
+
 		ty := args[0].Type()
+		key := args[1].AsString()
 		switch {
 		case ty.IsObjectType():
+			if ty.HasAttribute(key) {
+				return args[0].GetAttr(key).Type(), nil
+			} else if len(args) == 3 {
+				// if the key isn't found but a default is provided,
+				// return the default type
+				return args[2].Type(), nil
+			}
 			return cty.DynamicPseudoType, nil
 		default:
-			return args[0].Type().ElementType(), nil
+			return ty.ElementType(), nil
 		}
 	},
 	Impl: func(args []cty.Value, retType cty.Type) (ret cty.Value, err error) {
@@ -521,17 +529,11 @@ var LookupFunc = function.New(&function.Spec{
 			return cty.UnknownVal(retType), nil
 		}
 
-		if retType == cty.DynamicPseudoType {
-			// we're dealing with object-typed values
-			traversal := hcl.TraverseAttr{Name: args[1].AsString()}
-			ret, diags := traversal.TraversalStep(args[0])
-			if diags.HasErrors() {
-				return cty.UnknownVal(retType), diags
+		if mapVar.Type().IsObjectType() {
+			if mapVar.Type().HasAttribute(lookupKey) {
+				return mapVar.GetAttr(lookupKey), nil
 			}
-			return ret, nil
-		}
-
-		if mapVar.HasIndex(cty.StringVal(lookupKey)) == cty.True {
+		} else if mapVar.HasIndex(cty.StringVal(lookupKey)) == cty.True {
 			v := mapVar.Index(cty.StringVal(lookupKey))
 			if ty := v.Type(); !ty.Equals(cty.NilType) {
 				switch {
diff --git a/lang/funcs/collection_test.go b/lang/funcs/collection_test.go
index 4589c4c2f6..26a50af2f6 100644
--- a/lang/funcs/collection_test.go
+++ b/lang/funcs/collection_test.go
@@ -1165,6 +1165,14 @@ func TestLookup(t *testing.T) {
 			cty.NilVal,
 			true,
 		},
+		{ // Invalid key
+			[]cty.Value{
+				mapWithObjects,
+				cty.StringVal("bar"),
+			},
+			cty.NilVal,
+			true,
+		},
 		{ // Supplied default with valid key
 			[]cty.Value{
 				simpleMap,
@@ -1183,6 +1191,15 @@ func TestLookup(t *testing.T) {
 			cty.StringVal("bar"),
 			false,
 		},
+		{ // Supplied default with valid key
+			[]cty.Value{
+				mapWithObjects,
+				cty.StringVal("foobar"),
+				cty.StringVal(""),
+			},
+			cty.StringVal(""),
+			false,
+		},
 		{ // Supplied default with invalid key
 			[]cty.Value{
 				simpleMap,