Skip to content

Commit 058c4ca

Browse files
committed
funcr: Escape strings when they are not known-safe
1 parent 653e8e8 commit 058c4ca

File tree

2 files changed

+119
-39
lines changed

2 files changed

+119
-39
lines changed

funcr/funcr.go

Lines changed: 53 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -236,7 +236,7 @@ func (f Formatter) render(builtins, args []interface{}) string {
236236
if hook := f.opts.RenderBuiltinsHook; hook != nil {
237237
vals = hook(f.sanitize(vals))
238238
}
239-
f.flatten(buf, vals, false)
239+
f.flatten(buf, vals, false, false) // keys are ours, no need to escape
240240
continuing := len(builtins) > 0
241241
if len(f.valuesStr) > 0 {
242242
if continuing {
@@ -253,7 +253,7 @@ func (f Formatter) render(builtins, args []interface{}) string {
253253
if hook := f.opts.RenderArgsHook; hook != nil {
254254
vals = hook(f.sanitize(vals))
255255
}
256-
f.flatten(buf, vals, continuing)
256+
f.flatten(buf, vals, continuing, true) // escape user-provided keys
257257
if f.outputFormat == outputJSON {
258258
buf.WriteByte('}')
259259
}
@@ -263,10 +263,13 @@ func (f Formatter) render(builtins, args []interface{}) string {
263263
// flatten renders a list of key-value pairs into a buffer. If continuing is
264264
// true, it assumes that the buffer has previous values and will emit a
265265
// separator (which depends on the output format) before the first pair it
266-
// writes. This also returns a potentially modified version of kvList, which
266+
// writes. If escapeKeys is true, the keys are assumed to have
267+
// non-JSON-compatible characters in them and must be evaluated for escapes.
268+
//
269+
// This function returns a potentially modified version of kvList, which
267270
// ensures that there is a value for every key (adding a value if needed) and
268271
// that each key is a string (substituting a key if needed).
269-
func (f Formatter) flatten(buf *bytes.Buffer, kvList []interface{}, continuing bool) []interface{} {
272+
func (f Formatter) flatten(buf *bytes.Buffer, kvList []interface{}, continuing bool, escapeKeys bool) []interface{} {
270273
// This logic overlaps with sanitize() but saves one type-cast per key,
271274
// which can be measurable.
272275
if len(kvList)%2 != 0 {
@@ -290,9 +293,14 @@ func (f Formatter) flatten(buf *bytes.Buffer, kvList []interface{}, continuing b
290293
}
291294
}
292295

293-
buf.WriteByte('"')
294-
buf.WriteString(k)
295-
buf.WriteByte('"')
296+
if escapeKeys {
297+
buf.WriteString(prettyString(k))
298+
} else {
299+
// this is faster
300+
buf.WriteByte('"')
301+
buf.WriteString(k)
302+
buf.WriteByte('"')
303+
}
296304
if f.outputFormat == outputJSON {
297305
buf.WriteByte(':')
298306
} else {
@@ -308,8 +316,7 @@ func (f Formatter) pretty(value interface{}) string {
308316
}
309317

310318
const (
311-
flagRawString = 0x1 // do not print quotes on strings
312-
flagRawStruct = 0x2 // do not print braces on structs
319+
flagRawStruct = 0x1 // do not print braces on structs
313320
)
314321

315322
// TODO: This is not fast. Most of the overhead goes here.
@@ -334,11 +341,7 @@ func (f Formatter) prettyWithFlags(value interface{}, flags uint32) string {
334341
case bool:
335342
return strconv.FormatBool(v)
336343
case string:
337-
if flags&flagRawString > 0 {
338-
return v
339-
}
340-
// This is empirically faster than strings.Builder.
341-
return strconv.Quote(v)
344+
return prettyString(v)
342345
case int:
343346
return strconv.FormatInt(int64(v), 10)
344347
case int8:
@@ -379,9 +382,8 @@ func (f Formatter) prettyWithFlags(value interface{}, flags uint32) string {
379382
if i > 0 {
380383
buf.WriteByte(',')
381384
}
382-
buf.WriteByte('"')
383-
buf.WriteString(v[i].(string))
384-
buf.WriteByte('"')
385+
// arbitrary keys might need escaping
386+
buf.WriteString(prettyString(v[i].(string)))
385387
buf.WriteByte(':')
386388
buf.WriteString(f.pretty(v[i+1]))
387389
}
@@ -401,11 +403,7 @@ func (f Formatter) prettyWithFlags(value interface{}, flags uint32) string {
401403
case reflect.Bool:
402404
return strconv.FormatBool(v.Bool())
403405
case reflect.String:
404-
if flags&flagRawString > 0 {
405-
return v.String()
406-
}
407-
// This is empirically faster than strings.Builder.
408-
return `"` + v.String() + `"`
406+
return prettyString(v.String())
409407
case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
410408
return strconv.FormatInt(int64(v.Int()), 10)
411409
case reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64, reflect.Uintptr:
@@ -463,6 +461,7 @@ func (f Formatter) prettyWithFlags(value interface{}, flags uint32) string {
463461
if name == "" {
464462
name = fld.Name
465463
}
464+
// field names can't contain characters which need escaping
466465
buf.WriteByte('"')
467466
buf.WriteString(name)
468467
buf.WriteByte('"')
@@ -493,10 +492,14 @@ func (f Formatter) prettyWithFlags(value interface{}, flags uint32) string {
493492
if i > 0 {
494493
buf.WriteByte(',')
495494
}
496-
// JSON only does string keys.
497-
buf.WriteByte('"')
498-
buf.WriteString(f.prettyWithFlags(it.Key().Interface(), flagRawString))
499-
buf.WriteByte('"')
495+
// prettyWithFlags will produce already-escaped values
496+
keystr := f.prettyWithFlags(it.Key().Interface(), 0)
497+
if t.Key().Kind() != reflect.String {
498+
// JSON only does string keys. Unlike Go's standard JSON, we'll
499+
// convert just about anything to a string.
500+
keystr = prettyString(keystr)
501+
}
502+
buf.WriteString(keystr)
500503
buf.WriteByte(':')
501504
buf.WriteString(f.pretty(it.Value().Interface()))
502505
i++
@@ -512,6 +515,29 @@ func (f Formatter) prettyWithFlags(value interface{}, flags uint32) string {
512515
return fmt.Sprintf(`"<unhandled-%s>"`, t.Kind().String())
513516
}
514517

518+
func prettyString(s string) string {
519+
// Avoid escaping (which does allocations) if we can.
520+
if needsEscape(s) {
521+
return strconv.Quote(s)
522+
}
523+
b := bytes.NewBuffer(make([]byte, 0, 1024))
524+
b.WriteByte('"')
525+
b.WriteString(s)
526+
b.WriteByte('"')
527+
return b.String()
528+
}
529+
530+
// needsEscape determines whether the input string needs to be escaped or not,
531+
// without doing any allocations.
532+
func needsEscape(s string) bool {
533+
for _, r := range s {
534+
if !strconv.IsPrint(r) || r == '\\' || r == '"' {
535+
return true
536+
}
537+
}
538+
return false
539+
}
540+
515541
func isEmpty(v reflect.Value) bool {
516542
switch v.Kind() {
517543
case reflect.Array, reflect.Map, reflect.Slice, reflect.String:
@@ -683,7 +709,7 @@ func (f *Formatter) AddValues(kvList []interface{}) {
683709

684710
// Pre-render values, so we don't have to do it on each Info/Error call.
685711
buf := bytes.NewBuffer(make([]byte, 0, 1024))
686-
f.values = f.flatten(buf, vals, false)
712+
f.values = f.flatten(buf, vals, false, true) // escape user-provided keys
687713
f.valuesStr = buf.String()
688714
}
689715

funcr/funcr_test.go

Lines changed: 66 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,7 @@ import (
2727
"github.com/go-logr/logr"
2828
)
2929

30+
// Will be handled via reflection instead of type assertions.
3031
type substr string
3132

3233
func ptrint(i int) *int {
@@ -198,7 +199,9 @@ func TestPretty(t *testing.T) {
198199
exp string // used in cases where JSON can't handle it
199200
}{
200201
{val: "strval"},
202+
{val: "strval\nwith\t\"escapes\""},
201203
{val: substr("substrval")},
204+
{val: substr("substrval\nwith\t\"escapes\"")},
202205
{val: true},
203206
{val: false},
204207
{val: int(93)},
@@ -235,7 +238,11 @@ func TestPretty(t *testing.T) {
235238
exp: `[]`,
236239
},
237240
{val: []int{9, 3, 7, 6}},
241+
{val: []string{"str", "with\tescape"}},
242+
{val: []substr{"substr", "with\tescape"}},
238243
{val: [4]int{9, 3, 7, 6}},
244+
{val: [2]string{"str", "with\tescape"}},
245+
{val: [2]substr{"substr", "with\tescape"}},
239246
{
240247
val: struct {
241248
Int int
@@ -255,11 +262,32 @@ func TestPretty(t *testing.T) {
255262
"nine": 3,
256263
},
257264
},
265+
{
266+
val: map[string]int{
267+
"with\tescape": 76,
268+
},
269+
},
258270
{
259271
val: map[substr]int{
260272
"nine": 3,
261273
},
262274
},
275+
{
276+
val: map[substr]int{
277+
"with\tescape": 76,
278+
},
279+
},
280+
{
281+
val: map[int]int{
282+
9: 3,
283+
},
284+
},
285+
{
286+
val: map[float64]int{
287+
9.5: 3,
288+
},
289+
exp: `{"9.5":3}`,
290+
},
263291
{
264292
val: struct {
265293
X int `json:"x"`
@@ -283,6 +311,7 @@ func TestPretty(t *testing.T) {
283311
val: []struct{ X, Y string }{
284312
{"nine", "three"},
285313
{"seven", "six"},
314+
{"with\t", "\tescapes"},
286315
},
287316
},
288317
{
@@ -438,6 +467,24 @@ func TestPretty(t *testing.T) {
438467
val: PseudoStruct(makeKV("f1", 1, "f2", true, "f3", []int{})),
439468
exp: `{"f1":1,"f2":true,"f3":[]}`,
440469
},
470+
{
471+
val: map[TjsontagsString]int{
472+
{String1: `"quoted"`, String4: `unquoted`}: 1,
473+
},
474+
exp: `{"{\"string1\":\"\\\"quoted\\\"\",\"-\":\"\",\"string4\":\"unquoted\",\"String5\":\"\"}":1}`,
475+
},
476+
{
477+
val: map[TjsontagsInt]int{
478+
{Int1: 1, Int2: 2}: 3,
479+
},
480+
exp: `{"{\"int1\":1,\"-\":0,\"Int5\":0}":3}`,
481+
},
482+
{
483+
val: map[[2]struct{ S string }]int{
484+
{{S: `"quoted"`}, {S: "unquoted"}}: 1,
485+
},
486+
exp: `{"[{\"S\":\"\\\"quoted\\\"\"},{\"S\":\"unquoted\"}]":1}`,
487+
},
441488
}
442489

443490
f := NewFormatter(Options{})
@@ -449,7 +496,7 @@ func TestPretty(t *testing.T) {
449496
} else {
450497
jb, err := json.Marshal(tc.val)
451498
if err != nil {
452-
t.Errorf("[%d]: unexpected error: %v", i, err)
499+
t.Fatalf("[%d]: unexpected error: %v", i, err)
453500
}
454501
want = string(jb)
455502
}
@@ -496,6 +543,13 @@ func TestRender(t *testing.T) {
496543
args: makeKV("bool", PseudoStruct(makeKV("boolsub", true))),
497544
expectKV: `"int"={"intsub":1} "str"={"strsub":"2"} "bool"={"boolsub":true}`,
498545
expectJSON: `{"int":{"intsub":1},"str":{"strsub":"2"},"bool":{"boolsub":true}}`,
546+
}, {
547+
name: "escapes",
548+
builtins: makeKV("\"1\"", 1), // will not be escaped, but should never happen
549+
values: makeKV("\tstr", "ABC"), // escaped
550+
args: makeKV("bool\n", true), // escaped
551+
expectKV: `""1""=1 "\tstr"="ABC" "bool\n"=true`,
552+
expectJSON: `{""1"":1,"\tstr":"ABC","bool\n":true}`,
499553
}, {
500554
name: "missing value",
501555
builtins: makeKV("builtin"),
@@ -505,27 +559,27 @@ func TestRender(t *testing.T) {
505559
expectJSON: `{"builtin":"<no-value>","value":"<no-value>","arg":"<no-value>"}`,
506560
}, {
507561
name: "non-string key int",
508-
args: makeKV(123, "val"),
562+
builtins: makeKV(123, "val"), // should never happen
509563
values: makeKV(456, "val"),
510-
builtins: makeKV(789, "val"),
511-
expectKV: `"<non-string-key: 789>"="val" "<non-string-key: 456>"="val" "<non-string-key: 123>"="val"`,
512-
expectJSON: `{"<non-string-key: 789>":"val","<non-string-key: 456>":"val","<non-string-key: 123>":"val"}`,
564+
args: makeKV(789, "val"),
565+
expectKV: `"<non-string-key: 123>"="val" "<non-string-key: 456>"="val" "<non-string-key: 789>"="val"`,
566+
expectJSON: `{"<non-string-key: 123>":"val","<non-string-key: 456>":"val","<non-string-key: 789>":"val"}`,
513567
}, {
514568
name: "non-string key struct",
515-
args: makeKV(struct {
569+
builtins: makeKV(struct { // will not be escaped, but should never happen
516570
F1 string
517571
F2 int
518-
}{"arg", 123}, "val"),
572+
}{"builtin", 123}, "val"),
519573
values: makeKV(struct {
520574
F1 string
521575
F2 int
522576
}{"value", 456}, "val"),
523-
builtins: makeKV(struct {
577+
args: makeKV(struct {
524578
F1 string
525579
F2 int
526-
}{"builtin", 789}, "val"),
527-
expectKV: `"<non-string-key: {"F1":"builtin",>"="val" "<non-string-key: {"F1":"value","F>"="val" "<non-string-key: {"F1":"arg","F2">"="val"`,
528-
expectJSON: `{"<non-string-key: {"F1":"builtin",>":"val","<non-string-key: {"F1":"value","F>":"val","<non-string-key: {"F1":"arg","F2">":"val"}`,
580+
}{"arg", 789}, "val"),
581+
expectKV: `"<non-string-key: {"F1":"builtin",>"="val" "<non-string-key: {\"F1\":\"value\",\"F>"="val" "<non-string-key: {\"F1\":\"arg\",\"F2\">"="val"`,
582+
expectJSON: `{"<non-string-key: {"F1":"builtin",>":"val","<non-string-key: {\"F1\":\"value\",\"F>":"val","<non-string-key: {\"F1\":\"arg\",\"F2\">":"val"}`,
529583
}}
530584

531585
for _, tc := range testCases {
@@ -534,7 +588,7 @@ func TestRender(t *testing.T) {
534588
formatter.AddValues(tc.values)
535589
r := formatter.render(tc.builtins, tc.args)
536590
if r != expect {
537-
t.Errorf("wrong output:\nexpected %q\n got %q", expect, r)
591+
t.Errorf("wrong output:\nexpected %v\n got %v", expect, r)
538592
}
539593
}
540594
t.Run("KV", func(t *testing.T) {

0 commit comments

Comments
 (0)