diff --git a/v1/ast/term.go b/v1/ast/term.go index 8d13fef05e..59f81ed6d0 100644 --- a/v1/ast/term.go +++ b/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) } diff --git a/v1/ast/term_bench_test.go b/v1/ast/term_bench_test.go index acb5b14264..08994c9702 100644 --- a/v1/ast/term_bench_test.go +++ b/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) + } + }) + } +} diff --git a/v1/ast/term_test.go b/v1/ast/term_test.go index 6f7e781dcb..ba22605bae 100644 --- a/v1/ast/term_test.go +++ b/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) diff --git a/v1/format/format.go b/v1/format/format.go index 251f45c3a6..ba4bf73181 100644 --- a/v1/format/format.go +++ b/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 diff --git a/v1/format/format_test.go b/v1/format/format_test.go index 69dc44c7ce..ccaa51cf03 100644 --- a/v1/format/format_test.go +++ b/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) } })