Fix template string not serialized with escaped { (#8161)

Escaping `{` only at the time of serializtion, which should help
avoid special case treatment of these string values. But if there
is a better approach I haven't thought of, let me know.

Fixes #8156

Signed-off-by: Anders Eknert <anders.eknert@apple.com>

v1/ast/term.go

@@ -25,7 +25,13 @@ import (
 	"github.com/open-policy-agent/opa/v1/util"
 )
 
-var errFindNotFound = errors.New("find: not found")
+var (
+	NullValue Value = Null{}
+
+	errFindNotFound = errors.New("find: not found")
+
+	varRegexp = regexp.MustCompile("^[[:alpha:]_][[:alpha:][:digit:]_]*$")
+)
 
 // Location records a position in source code.
 type Location = location.Location
@@ -544,8 +550,6 @@ func IsScalar(v Value) bool {
 // Null represents the null value defined by JSON.
 type Null struct{}
 
-var NullValue Value = Null{}
-
 // NullTerm creates a new Term with a Null value.
 func NullTerm() *Term {
 	return &Term{Value: NullValue}
@@ -927,14 +931,15 @@ func (ts *TemplateString) String() string {
 	for _, p := range ts.Parts {
 		switch x := p.(type) {
 		case *Expr:
-			str.WriteString("{")
+			str.WriteByte('{')
 			str.WriteString(p.String())
-			str.WriteString("}")
+			str.WriteByte('}')
 		case *Term:
 			s := p.String()
 			if _, ok := x.Value.(String); ok {
 				s = strings.TrimPrefix(s, "\"")
 				s = strings.TrimSuffix(s, "\"")
+				s = EscapeTemplateStringStringPart(s)
 			}
 			str.WriteString(s)
 		default:
@@ -942,7 +947,7 @@ func (ts *TemplateString) String() string {
 		}
 	}
 
-	str.WriteString("\"")
+	str.WriteByte('"')
 	return str.String()
 }
 
@@ -950,6 +955,53 @@ func TemplateStringTerm(multiLine bool, parts ...Node) *Term {
 	return &Term{Value: &TemplateString{MultiLine: multiLine, Parts: parts}}
 }
 
+// EscapeTemplateStringStringPart escapes unescaped left curly braces in s - i.e "{" becomes "\{".
+// The internal representation of string terms within a template string does **NOT**
+// treat '{' as special, but expects code dealing with template strings to escape them when
+// required, such as when serializing the complete template string. Code that programmatically
+// constructs template strings should not pre-escape left curly braces in string term parts.
+//
+// // TODO(anders): a future optimization would be to combine this with the other escaping done
+// // for strings (e.g. escaping quotes, backslashes, and JSON control characters) in a single operation
+// // to avoid multiple passes and allocations over the same string. That's currently done by
+// // strconv.Quote, so we would need to re-implement that logic in code of our own.
+// // NOTE(anders): I would love to come up with a better name for this component than
+// // "TemplateStringStringPart"..
+func EscapeTemplateStringStringPart(s string) string {
+	numUnescaped := countUnescapedLeftCurly(s)
+	if numUnescaped == 0 {
+		return s
+	}
+
+	l := len(s)
+	escaped := make([]byte, 0, l+numUnescaped)
+	if s[0] == '{' {
+		escaped = append(escaped, '\\', s[0])
+	} else {
+		escaped = append(escaped, s[0])
+	}
+
+	for i := 1; i < l; i++ {
+		if s[i] == '{' && s[i-1] != '\\' {
+			escaped = append(escaped, '\\', s[i])
+		} else {
+			escaped = append(escaped, s[i])
+		}
+	}
+
+	return util.ByteSliceToString(escaped)
+}
+
+func countUnescapedLeftCurly(s string) (n int) {
+	// Note(anders): while not the functions I'd intuitively reach for to solve this,
+	// they are hands down the fastest option here, as they're done in assembly, which
+	// performs about an order of magnitude better than a manual loop in Go.
+	if n = strings.Count(s, "{"); n > 0 {
+		n -= strings.Count(s, `\{`)
+	}
+	return n
+}
+
 // Var represents a variable as defined by the language.
 type Var string
 
@@ -1288,8 +1340,6 @@ func (ref Ref) Ptr() (string, error) {
 	return buf.String(), nil
 }
 
-var varRegexp = regexp.MustCompile("^[[:alpha:]_][[:alpha:][:digit:]_]*$")
-
 func IsVarCompatibleString(s string) bool {
 	return varRegexp.MatchString(s)
 }

v1/ast/term_bench_test.go

@@ -931,3 +931,49 @@ func BenchmarkPtr(b *testing.B) {
 		}
 	})
 }
+
+func BenchmarkEscapeTemplateStringStringPart(b *testing.B) {
+	inputs := []string{
+		"",
+		"{",
+		"\\{",
+		"}{",
+		"no curly!!!",
+		"{unes{caped",
+		"{{{{{{{{{{",
+	}
+	repeat := 100
+
+	for _, input := range inputs {
+		b.Run(fmt.Sprintf("%s * %d", input, repeat), func(b *testing.B) {
+			s := strings.Repeat(input, repeat)
+			for b.Loop() {
+				// expected output covered in test already
+				_ = EscapeTemplateStringStringPart(s)
+			}
+		})
+	}
+}
+
+func BenchmarkCountUnescapedLeftCurly(b *testing.B) {
+	inputs := []string{
+		"",
+		"{",
+		"\\{",
+		"}{",
+		"no curly!!!",
+		"{unes{caped",
+		"{{{{{{{{{{",
+	}
+	repeat := 100
+
+	for _, input := range inputs {
+		b.Run(fmt.Sprintf("%s * %d", input, repeat), func(b *testing.B) {
+			s := strings.Repeat(input, repeat)
+			for b.Loop() {
+				// expected output covered in test already
+				_ = countUnescapedLeftCurly(s)
+			}
+		})
+	}
+}

v1/ast/term_test.go

@@ -8,6 +8,7 @@ import (
 	"bytes"
 	"encoding/json"
 	"errors"
+	"fmt"
 	"math/rand"
 	"reflect"
 	"runtime"
@@ -568,6 +569,9 @@ func TestTermString(t *testing.T) {
 	// ensure that objects and sets have deterministic String() results
 	assertToString(t, SetTerm(VarTerm("y"), VarTerm("x")).Value, "{x, y}")
 	assertToString(t, ObjectTerm([2]*Term{VarTerm("y"), VarTerm("b")}, [2]*Term{VarTerm("x"), VarTerm("a")}).Value, "{x: a, y: b}")
+
+	assertToString(t, MustParseTerm(`$"foo {bar}"`).Value, `$"foo {bar}"`)
+	assertToString(t, MustParseTerm(`$"foo \{bar}"`).Value, `$"foo \{bar}"`)
 }
 
 func TestRefString_Escapes(t *testing.T) {
@@ -1689,6 +1693,65 @@ func TestLazyObjectCompare(t *testing.T) {
 	assertForced(t, x, true)
 }
 
+func TestEscapeTemplateStringStringPart(t *testing.T) {
+	t.Parallel()
+
+	cases := []struct {
+		inp string
+		exp string
+	}{
+		{inp: "", exp: ""},
+		{inp: "{", exp: "\\{"},
+		{inp: "\\{", exp: "\\{"},
+		{inp: "{\\{}", exp: "\\{\\{}"},
+		{inp: "\n{", exp: "\n\\{"},
+		{inp: "}{", exp: "}\\{"},
+		{inp: "no curly!!!", exp: "no curly!!!"},
+		{inp: "{unes{caped", exp: "\\{unes\\{caped"},
+		{inp: "{{{{{{{{{{", exp: "\\{\\{\\{\\{\\{\\{\\{\\{\\{\\{"},
+	}
+
+	for _, c := range cases {
+		t.Run(fmt.Sprintf("%q", c.inp), func(t *testing.T) {
+			t.Parallel()
+
+			if got := EscapeTemplateStringStringPart(c.inp); got != c.exp {
+				t.Fatalf("\nexp %q\ngot %q", c.exp, got)
+			}
+		})
+	}
+}
+
+func TestCountUnescapedLeftCurly(t *testing.T) {
+	t.Parallel()
+
+	cases := []struct {
+		inp string
+		exp int
+	}{
+		{inp: "", exp: 0},
+		{inp: "{", exp: 1},
+		{inp: "\\{", exp: 0},
+		{inp: "{\\{}", exp: 1},
+		{inp: "\n{", exp: 1},
+		{inp: "}{", exp: 1},
+		{inp: "no curly!!!", exp: 0},
+		{inp: "{unes{caped", exp: 2},
+		{inp: "{{{{{{{{{{", exp: 10},
+		{inp: "\\{{\\{{\\{{", exp: 3},
+	}
+
+	for _, c := range cases {
+		t.Run(fmt.Sprintf("%q", c.inp), func(t *testing.T) {
+			t.Parallel()
+
+			if got := countUnescapedLeftCurly(c.inp); got != c.exp {
+				t.Fatalf("\nexp %d\ngot %d", c.exp, got)
+			}
+		})
+	}
+}
+
 func TestTemplateStringEqual(t *testing.T) {
 	a := MustParseTerm(`$"hello {world}!"`).Value.(*TemplateString)
 	b := MustParseTerm(`$"hello {world}!"`).Value.(*TemplateString)

v1/format/format.go

@@ -1309,7 +1309,7 @@ func (w *writer) writeTemplateString(ts *ast.TemplateString, comments []*ast.Com
 	if ts.MultiLine {
 		w.write("`")
 	} else {
-		w.write("\"")
+		w.write(`"`)
 	}
 
 	for i, p := range ts.Parts {
@@ -1364,9 +1364,9 @@ func (w *writer) writeTemplateString(ts *ast.TemplateString, comments []*ast.Com
 		case *ast.Term:
 			if s, ok := x.Value.(ast.String); ok {
 				if ts.MultiLine {
-					w.write(string(s))
+					w.write(ast.EscapeTemplateStringStringPart(string(s)))
 				} else {
-					str := s.String()
+					str := ast.EscapeTemplateStringStringPart(s.String())
 					w.write(str[1 : len(str)-1])
 				}
 			} else {
@@ -1383,7 +1383,7 @@ func (w *writer) writeTemplateString(ts *ast.TemplateString, comments []*ast.Com
 	if ts.MultiLine {
 		w.write("`")
 	} else {
-		w.write("\"")
+		w.write(`"`)
 	}
 
 	return comments, nil

v1/format/format_test.go

@@ -723,6 +723,28 @@ a[_x[y][[z, w]]]`,
 			),
 			expected: "$`foo {\"bar\"} {`baz`}`",
 		},
+		{
+			note: "template-string, left curly in string part",
+			toFmt: ast.TemplateStringTerm(false,
+				ast.StringTerm(`allow if { `),
+				&ast.Expr{
+					Terms: ast.VarTerm("x"),
+				},
+				ast.StringTerm(" }"),
+			),
+			expected: `$"allow if \{ {x} }"`,
+		},
+		{
+			note: "template-string, multi-line, left curly in string part",
+			toFmt: ast.TemplateStringTerm(true,
+				ast.StringTerm("allow if {\n\t"),
+				&ast.Expr{
+					Terms: ast.VarTerm("x"),
+				},
+				ast.StringTerm("\n}"),
+			),
+			expected: "$`allow if \\{\n\t{x}\n}`",
+		},
 	}
 
 	for _, tc := range cases {
@@ -745,11 +767,7 @@ a[_x[y][[z, w]]]`,
 
 		// consistency check: disregarding source locations, it shouldn't panic
 		t.Run("no_loc/"+tc.note, func(t *testing.T) {
-			_, err := AstWithOpts(tc.toFmt, Opts{IgnoreLocations: true})
-			if err != nil {
-				t.Fatalf("Unexpected error: %s", err)
-			}
-			if err != nil {
+			if _, err := AstWithOpts(tc.toFmt, Opts{IgnoreLocations: true}); err != nil {
 				t.Fatalf("Unexpected error: %s", err)
 			}
 		})