fix: return the raw strings when formatting (#7525)

prevent `\u0000` from being changed to `\x00` 

Signed-off-by: sspaink <[email protected]>

Author: sspaink

v1/format/format.go

@@ -1207,10 +1207,29 @@ func (w *writer) writeTermParens(parens bool, term *ast.Term, comments []*ast.Co
 	case ast.String:
 		if term.Location.Text[0] == '`' {
 			// To preserve raw strings, we need to output the original text,
-			// not what x.String() would give us.
 			w.write(string(term.Location.Text))
 		} else {
-			w.write(x.String())
+			// x.String() cannot be used by default because it can change the input string "\u0000" to "\x00"
+			var after, quote string
+			var found bool
+			// term.Location.Text could contain the prefix `else :=`, remove it
+			switch term.Location.Text[len(term.Location.Text)-1] {
+			case '"':
+				quote = "\""
+				_, after, found = strings.Cut(string(term.Location.Text), quote)
+			case '`':
+				quote = "`"
+				_, after, found = strings.Cut(string(term.Location.Text), quote)
+			}
+
+			if !found {
+				// If no quoted string was found, that means it is a key being formatted to a string
+				// e.g. partial_set.y to partial_set["y"]
+				w.write(x.String())
+			} else {
+				w.write(quote + after)
+			}
+
 		}
 	case ast.Var:
 		w.write(w.formatVar(x))

v1/format/format_test.go

@@ -69,7 +69,7 @@ func TestFormatNilLocationEmptyBody(t *testing.T) {
 func TestFormatNilLocationFunctionArgs(t *testing.T) {
 	b := ast.NewBody()
 	s := ast.StringTerm(" ")
-	s.SetLocation(location.NewLocation([]byte("foo"), "p.rego", 2, 2))
+	s.SetLocation(location.NewLocation([]byte("\" \""), "p.rego", 2, 2))
 	b.Append(ast.Split.Expr(ast.NewTerm(ast.Var("__local1__")), s, ast.NewTerm(ast.Var("__local2__"))))
 	exp := "split(__local1__, \" \", __local2__)\n"
 	bs, err := Ast(b)

v1/format/testfiles/v0/test_unicode.rego

@@ -0,0 +1,15 @@
+package test
+
+x := "\u0000"
+x := "\u0000 \""
+
+_fg := {
+    "black":    "\u001b[30m",
+    "red":      "\u001b[31m",
+    "green":    "\u001b[32m",
+    "yellow":   "\u001b[33m",
+    "blue":     "\u001b[34m",
+    "magenta":  "\u001b[35m",
+    "cyan":     "\u001b[36m",
+    "white":    "\u001b[37m",
+}

v1/format/testfiles/v0/test_unicode.rego.formatted

@@ -0,0 +1,16 @@
+package test
+
+x := "\u0000"
+
+x := "\u0000 \""
+
+_fg := {
+	"black": "\u001b[30m",
+	"red": "\u001b[31m",
+	"green": "\u001b[32m",
+	"yellow": "\u001b[33m",
+	"blue": "\u001b[34m",
+	"magenta": "\u001b[35m",
+	"cyan": "\u001b[36m",
+	"white": "\u001b[37m",
+}

v1/format/testfiles/v1/test_else_strings.rego

@@ -0,0 +1,14 @@
+package test
+
+a := "foo" if {
+	false
+} else := `{"foo":"bar"}`
+
+
+a := "foo" if {
+	false
+} else := "`{\"foo\":\"bar\"}`"
+
+a := "foo" if {
+	false
+} else := "foo"

v1/format/testfiles/v1/test_else_strings.rego.formatted

@@ -0,0 +1,13 @@
+package test
+
+a := "foo" if {
+	false
+} else := `{"foo":"bar"}`
+
+a := "foo" if {
+	false
+} else := "`{\"foo\":\"bar\"}`"
+
+a := "foo" if {
+	false
+} else := "foo"

v1/format/testfiles/v1/test_unicode.rego

@@ -0,0 +1,22 @@
+package test
+
+x := "\u0000"
+x := "\u0000 \""
+
+authorize = "\u0000" if {
+    input.user == "\u0000"
+} else = "\u0000" if {
+    input.path[0] == "\u0000"
+    input.source_network == "\u0000"
+}
+
+_fg := {
+    "black":    "\u001b[30m",
+    "red":      "\u001b[31m",
+    "green":    "\u001b[32m",
+    "yellow":   "\u001b[33m",
+    "blue":     "\u001b[34m",
+    "magenta":  "\u001b[35m",
+    "cyan":     "\u001b[36m",
+    "white":    "\u001b[37m",
+}

v1/format/testfiles/v1/test_unicode.rego.formatted

@@ -0,0 +1,22 @@
+package test
+
+x := "\u0000"
+x := "\u0000 \""
+
+authorize := "\u0000" if {
+	input.user == "\u0000"
+} else := "\u0000" if {
+	input.path[0] == "\u0000"
+	input.source_network == "\u0000"
+}
+
+_fg := {
+	"black": "\u001b[30m",
+	"red": "\u001b[31m",
+	"green": "\u001b[32m",
+	"yellow": "\u001b[33m",
+	"blue": "\u001b[34m",
+	"magenta": "\u001b[35m",
+	"cyan": "\u001b[36m",
+	"white": "\u001b[37m",
+}