From b5a4860a36fa5007acfb89e00c99af4f5488cf36 Mon Sep 17 00:00:00 2001
From: Martin Atkins <mart@degeneration.co.uk>
Date: Tue, 16 Apr 2024 13:24:34 -0700
Subject: [PATCH] lang/marks: MarkPaths helper function

This helper is a slightly more convenient wrapper around cty's own
"MarkWithPaths" function that applies a single given mark at zero or more
paths in a given value.

In several parts of Terraform we preserve sensitive marks in particular
for serialization and cannot preserve anything else, and so this helper
is useful for reapplying the saved sensitive paths back to the stored
value during decoding.
---
 internal/lang/marks/paths.go      | 24 ++++++++++++++
 internal/lang/marks/paths_test.go | 52 +++++++++++++++++++++++++++++++
 2 files changed, 76 insertions(+)

diff --git a/internal/lang/marks/paths.go b/internal/lang/marks/paths.go
index affd23e23b..461ec1bb42 100644
--- a/internal/lang/marks/paths.go
+++ b/internal/lang/marks/paths.go
@@ -36,3 +36,27 @@ func PathsWithMark(pvms []cty.PathValueMarks, wantMark any) (withWanted []cty.Pa
 
 	return withWanted, withOthers
 }
+
+// MarkPaths transforms the given value by marking each of the given paths
+// with the given mark value.
+func MarkPaths(val cty.Value, mark any, paths []cty.Path) cty.Value {
+	if len(paths) == 0 {
+		// No-allocations path for the common case where there are no marked paths at all.
+		return val
+	}
+
+	// For now we'll use cty's slightly lower-level function to achieve this
+	// result. This is a little inefficient due to an additional dynamic
+	// allocation for the intermediate data structure, so if that becomes
+	// a problem in practice then we may wish to write a more direct
+	// implementation here.
+	markses := make([]cty.PathValueMarks, len(paths))
+	marks := cty.NewValueMarks(mark)
+	for i, path := range paths {
+		markses[i] = cty.PathValueMarks{
+			Path:  path,
+			Marks: marks,
+		}
+	}
+	return val.MarkWithPaths(markses)
+}
diff --git a/internal/lang/marks/paths_test.go b/internal/lang/marks/paths_test.go
index c74a6fce5f..f5efd2829f 100644
--- a/internal/lang/marks/paths_test.go
+++ b/internal/lang/marks/paths_test.go
@@ -56,3 +56,55 @@ func TestPathsWithMark(t *testing.T) {
 		t.Errorf("wrong set of entries with other marks\n%s", diff)
 	}
 }
+
+func TestMarkPaths(t *testing.T) {
+	value := cty.ObjectVal(map[string]cty.Value{
+		"s": cty.StringVal(".s"),
+		"l": cty.ListVal([]cty.Value{
+			cty.StringVal(".l[0]"),
+			cty.StringVal(".l[1]"),
+		}),
+		"m": cty.MapVal(map[string]cty.Value{
+			"a": cty.StringVal(`.m["a"]`),
+			"b": cty.StringVal(`.m["b"]`),
+		}),
+		"o": cty.ObjectVal(map[string]cty.Value{
+			"a": cty.StringVal(".o.a"),
+			"b": cty.StringVal(".o.b"),
+		}),
+		"t": cty.TupleVal([]cty.Value{
+			cty.StringVal(`.t[0]`),
+			cty.StringVal(`.t[1]`),
+		}),
+	})
+	sensitivePaths := []cty.Path{
+		cty.GetAttrPath("s"),
+		cty.GetAttrPath("l").IndexInt(1),
+		cty.GetAttrPath("m").IndexString("a"),
+		cty.GetAttrPath("o").GetAttr("b"),
+		cty.GetAttrPath("t").IndexInt(0),
+	}
+	got := MarkPaths(value, Sensitive, sensitivePaths)
+	want := cty.ObjectVal(map[string]cty.Value{
+		"s": cty.StringVal(".s").Mark(Sensitive),
+		"l": cty.ListVal([]cty.Value{
+			cty.StringVal(".l[0]"),
+			cty.StringVal(".l[1]").Mark(Sensitive),
+		}),
+		"m": cty.MapVal(map[string]cty.Value{
+			"a": cty.StringVal(`.m["a"]`).Mark(Sensitive),
+			"b": cty.StringVal(`.m["b"]`),
+		}),
+		"o": cty.ObjectVal(map[string]cty.Value{
+			"a": cty.StringVal(".o.a"),
+			"b": cty.StringVal(".o.b").Mark(Sensitive),
+		}),
+		"t": cty.TupleVal([]cty.Value{
+			cty.StringVal(`.t[0]`).Mark(Sensitive),
+			cty.StringVal(`.t[1]`),
+		}),
+	})
+	if diff := cmp.Diff(want, got, ctydebug.CmpOptions); diff != "" {
+		t.Errorf("wrong result\n%s", diff)
+	}
+}