functions: Improve marks support for lookup

Several changes to lookup to improve how we handle marked values:

- If the entire collection is marked, preserve the marks on any result
  (whether successful or fallback)
- If a returned value from the collection is marked, preserve the marks
  from only that value, combined with any overall collection marks
- Retain marks on the fallback value when it is returned, combined with
  any overall collection marks
- Include marks on the key in the result, as otherwise the result it
  ends up selecting could imply what the sensitive value was
- Retain collection marks when returning an unknown value for a not
  wholly-known collection

See also zclconf/go-cty#98

lang/funcs/collection.go

@@ -214,12 +214,14 @@ var IndexFunc = function.New(&function.Spec{
 var LookupFunc = function.New(&function.Spec{
 	Params: []function.Parameter{
 		{
-			Name: "inputMap",
-			Type: cty.DynamicPseudoType,
+			Name:        "inputMap",
+			Type:        cty.DynamicPseudoType,
+			AllowMarked: true,
 		},
 		{
-			Name: "key",
-			Type: cty.String,
+			Name:        "key",
+			Type:        cty.String,
+			AllowMarked: true,
 		},
 	},
 	VarParam: &function.Parameter{
@@ -228,6 +230,7 @@ var LookupFunc = function.New(&function.Spec{
 		AllowUnknown:     true,
 		AllowDynamicType: true,
 		AllowNull:        true,
+		AllowMarked:      true,
 	},
 	Type: func(args []cty.Value) (ret cty.Type, err error) {
 		if len(args) < 1 || len(args) > 3 {
@@ -242,7 +245,8 @@ var LookupFunc = function.New(&function.Spec{
 				return cty.DynamicPseudoType, nil
 			}
 
-			key := args[1].AsString()
+			keyVal, _ := args[1].Unmark()
+			key := keyVal.AsString()
 			if ty.HasAttribute(key) {
 				return args[0].GetAttr(key).Type(), nil
 			} else if len(args) == 3 {
@@ -268,34 +272,46 @@ var LookupFunc = function.New(&function.Spec{
 		defaultValueSet := false
 
 		if len(args) == 3 {
+			// intentionally leave default value marked
 			defaultVal = args[2]
 			defaultValueSet = true
 		}
 
-		mapVar := args[0]
-		lookupKey := args[1].AsString()
+		// keep track of marks from the collection and key
+		var markses []cty.ValueMarks
+
+		// unmark collection, retain marks to reapply later
+		mapVar, mapMarks := args[0].Unmark()
+		markses = append(markses, mapMarks)
+
+		// include marks on the key in the result
+		keyVal, keyMarks := args[1].Unmark()
+		if len(keyMarks) > 0 {
+			markses = append(markses, keyMarks)
+		}
+		lookupKey := keyVal.AsString()
 
 		if !mapVar.IsKnown() {
-			return cty.UnknownVal(retType), nil
+			return cty.UnknownVal(retType).WithMarks(markses...), nil
 		}
 
 		if mapVar.Type().IsObjectType() {
 			if mapVar.Type().HasAttribute(lookupKey) {
-				return mapVar.GetAttr(lookupKey), nil
+				return mapVar.GetAttr(lookupKey).WithMarks(markses...), nil
 			}
 		} else if mapVar.HasIndex(cty.StringVal(lookupKey)) == cty.True {
-			return mapVar.Index(cty.StringVal(lookupKey)), nil
+			return mapVar.Index(cty.StringVal(lookupKey)).WithMarks(markses...), nil
 		}
 
 		if defaultValueSet {
 			defaultVal, err = convert.Convert(defaultVal, retType)
 			if err != nil {
 				return cty.NilVal, err
 			}
-			return defaultVal, nil
+			return defaultVal.WithMarks(markses...), nil
 		}
 
-		return cty.UnknownVal(cty.DynamicPseudoType), fmt.Errorf(
+		return cty.UnknownVal(cty.DynamicPseudoType).WithMarks(markses...), fmt.Errorf(
 			"lookup failed to find '%s'", lookupKey)
 	},
 })

lang/funcs/collection_test.go

@@ -795,6 +795,88 @@ func TestLookup(t *testing.T) {
 			cty.DynamicVal, // if the key is unknown then we don't know which object attribute and thus can't know the type
 			false,
 		},
+		{ // successful marked collection lookup returns marked value
+			[]cty.Value{
+				cty.MapVal(map[string]cty.Value{
+					"boop": cty.StringVal("beep"),
+				}).Mark("a"),
+				cty.StringVal("boop"),
+				cty.StringVal("nope"),
+			},
+			cty.StringVal("beep").Mark("a"),
+			false,
+		},
+		{ // apply collection marks to unknown return vaue
+			[]cty.Value{
+				cty.MapVal(map[string]cty.Value{
+					"boop": cty.StringVal("beep"),
+					"frob": cty.UnknownVal(cty.String),
+				}).Mark("a"),
+				cty.StringVal("frob"),
+				cty.StringVal("nope"),
+			},
+			cty.UnknownVal(cty.String).Mark("a"),
+			false,
+		},
+		{ // propagate collection marks to default when returning
+			[]cty.Value{
+				cty.MapVal(map[string]cty.Value{
+					"boop": cty.StringVal("beep"),
+				}).Mark("a"),
+				cty.StringVal("frob"),
+				cty.StringVal("nope").Mark("b"),
+			},
+			cty.StringVal("nope").WithMarks(cty.NewValueMarks("a", "b")),
+			false,
+		},
+		{ // on unmarked collection, return only marks from found value
+			[]cty.Value{
+				cty.MapVal(map[string]cty.Value{
+					"boop": cty.StringVal("beep").Mark("a"),
+					"frob": cty.StringVal("honk").Mark("b"),
+				}),
+				cty.StringVal("frob"),
+				cty.StringVal("nope").Mark("c"),
+			},
+			cty.StringVal("honk").Mark("b"),
+			false,
+		},
+		{ // on unmarked collection, return default exactly on missing
+			[]cty.Value{
+				cty.MapVal(map[string]cty.Value{
+					"boop": cty.StringVal("beep").Mark("a"),
+					"frob": cty.StringVal("honk").Mark("b"),
+				}),
+				cty.StringVal("squish"),
+				cty.StringVal("nope").Mark("c"),
+			},
+			cty.StringVal("nope").Mark("c"),
+			false,
+		},
+		{ // retain marks on default if converted
+			[]cty.Value{
+				cty.MapVal(map[string]cty.Value{
+					"boop": cty.StringVal("beep").Mark("a"),
+					"frob": cty.StringVal("honk").Mark("b"),
+				}),
+				cty.StringVal("squish"),
+				cty.NumberIntVal(5).Mark("c"),
+			},
+			cty.StringVal("5").Mark("c"),
+			false,
+		},
+		{ // propagate marks from key
+			[]cty.Value{
+				cty.MapVal(map[string]cty.Value{
+					"boop": cty.StringVal("beep"),
+					"frob": cty.StringVal("honk"),
+				}),
+				cty.StringVal("boop").Mark("a"),
+				cty.StringVal("nope"),
+			},
+			cty.StringVal("beep").Mark("a"),
+			false,
+		},
 	}
 
 	for _, test := range tests {