Skip to content

Commit 1215c81

Browse files
pohlypohly
authored
Merge pull request #113 from thockin/funcr-qstr
funcr: Escape strings when they are not known-safe.
2 parents 561de0c + c9b68e6 commit 1215c81

File tree

2 files changed

+159
-40
lines changed

2 files changed

+159
-40
lines changed

‎funcr/funcr.go

Lines changed: 68 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,7 @@ package funcr
3636

3737
import (
3838
"bytes"
39+
"encoding"
3940
"fmt"
4041
"path/filepath"
4142
"reflect"
@@ -236,7 +237,7 @@ func (f Formatter) render(builtins, args []interface{}) string {
236237
if hook := f.opts.RenderBuiltinsHook; hook != nil {
237238
vals = hook(f.sanitize(vals))
238239
}
239-
f.flatten(buf, vals, false)
240+
f.flatten(buf, vals, false, false) // keys are ours, no need to escape
240241
continuing := len(builtins) > 0
241242
if len(f.valuesStr) > 0 {
242243
if continuing {
@@ -253,7 +254,7 @@ func (f Formatter) render(builtins, args []interface{}) string {
253254
if hook := f.opts.RenderArgsHook; hook != nil {
254255
vals = hook(f.sanitize(vals))
255256
}
256-
f.flatten(buf, vals, continuing)
257+
f.flatten(buf, vals, continuing, true) // escape user-provided keys
257258
if f.outputFormat == outputJSON {
258259
buf.WriteByte('}')
259260
}
@@ -263,10 +264,13 @@ func (f Formatter) render(builtins, args []interface{}) string {
263264
// flatten renders a list of key-value pairs into a buffer. If continuing is
264265
// true, it assumes that the buffer has previous values and will emit a
265266
// separator (which depends on the output format) before the first pair it
266-
// writes. This also returns a potentially modified version of kvList, which
267+
// writes. If escapeKeys is true, the keys are assumed to have
268+
// non-JSON-compatible characters in them and must be evaluated for escapes.
269+
//
270+
// This function returns a potentially modified version of kvList, which
267271
// ensures that there is a value for every key (adding a value if needed) and
268272
// that each key is a string (substituting a key if needed).
269-
func (f Formatter) flatten(buf *bytes.Buffer, kvList []interface{}, continuing bool) []interface{} {
273+
func (f Formatter) flatten(buf *bytes.Buffer, kvList []interface{}, continuing bool, escapeKeys bool) []interface{} {
270274
// This logic overlaps with sanitize() but saves one type-cast per key,
271275
// which can be measurable.
272276
if len(kvList)%2 != 0 {
@@ -290,9 +294,14 @@ func (f Formatter) flatten(buf *bytes.Buffer, kvList []interface{}, continuing b
290294
}
291295
}
292296

293-
buf.WriteByte('"')
294-
buf.WriteString(k)
295-
buf.WriteByte('"')
297+
if escapeKeys {
298+
buf.WriteString(prettyString(k))
299+
} else {
300+
// this is faster
301+
buf.WriteByte('"')
302+
buf.WriteString(k)
303+
buf.WriteByte('"')
304+
}
296305
if f.outputFormat == outputJSON {
297306
buf.WriteByte(':')
298307
} else {
@@ -308,8 +317,7 @@ func (f Formatter) pretty(value interface{}) string {
308317
}
309318

310319
const (
311-
flagRawString = 0x1 // do not print quotes on strings
312-
flagRawStruct = 0x2 // do not print braces on structs
320+
flagRawStruct = 0x1 // do not print braces on structs
313321
)
314322

315323
// TODO: This is not fast. Most of the overhead goes here.
@@ -334,11 +342,7 @@ func (f Formatter) prettyWithFlags(value interface{}, flags uint32) string {
334342
case bool:
335343
return strconv.FormatBool(v)
336344
case string:
337-
if flags&flagRawString > 0 {
338-
return v
339-
}
340-
// This is empirically faster than strings.Builder.
341-
return strconv.Quote(v)
345+
return prettyString(v)
342346
case int:
343347
return strconv.FormatInt(int64(v), 10)
344348
case int8:
@@ -379,9 +383,8 @@ func (f Formatter) prettyWithFlags(value interface{}, flags uint32) string {
379383
if i > 0 {
380384
buf.WriteByte(',')
381385
}
382-
buf.WriteByte('"')
383-
buf.WriteString(v[i].(string))
384-
buf.WriteByte('"')
386+
// arbitrary keys might need escaping
387+
buf.WriteString(prettyString(v[i].(string)))
385388
buf.WriteByte(':')
386389
buf.WriteString(f.pretty(v[i+1]))
387390
}
@@ -401,11 +404,7 @@ func (f Formatter) prettyWithFlags(value interface{}, flags uint32) string {
401404
case reflect.Bool:
402405
return strconv.FormatBool(v.Bool())
403406
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() + `"`
407+
return prettyString(v.String())
409408
case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
410409
return strconv.FormatInt(int64(v.Int()), 10)
411410
case reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64, reflect.Uintptr:
@@ -463,6 +462,7 @@ func (f Formatter) prettyWithFlags(value interface{}, flags uint32) string {
463462
if name == "" {
464463
name = fld.Name
465464
}
465+
// field names can't contain characters which need escaping
466466
buf.WriteByte('"')
467467
buf.WriteString(name)
468468
buf.WriteByte('"')
@@ -493,10 +493,26 @@ func (f Formatter) prettyWithFlags(value interface{}, flags uint32) string {
493493
if i > 0 {
494494
buf.WriteByte(',')
495495
}
496-
// JSON only does string keys.
497-
buf.WriteByte('"')
498-
buf.WriteString(f.prettyWithFlags(it.Key().Interface(), flagRawString))
499-
buf.WriteByte('"')
496+
// If a map key supports TextMarshaler, use it.
497+
keystr := ""
498+
if m, ok := it.Key().Interface().(encoding.TextMarshaler); ok {
499+
txt, err := m.MarshalText()
500+
if err != nil {
501+
keystr = fmt.Sprintf("<error-MarshalText: %s>", err.Error())
502+
} else {
503+
keystr = string(txt)
504+
}
505+
keystr = prettyString(keystr)
506+
} else {
507+
// prettyWithFlags will produce already-escaped values
508+
keystr = f.prettyWithFlags(it.Key().Interface(), 0)
509+
if t.Key().Kind() != reflect.String {
510+
// JSON only does string keys. Unlike Go's standard JSON, we'll
511+
// convert just about anything to a string.
512+
keystr = prettyString(keystr)
513+
}
514+
}
515+
buf.WriteString(keystr)
500516
buf.WriteByte(':')
501517
buf.WriteString(f.pretty(it.Value().Interface()))
502518
i++
@@ -512,6 +528,29 @@ func (f Formatter) prettyWithFlags(value interface{}, flags uint32) string {
512528
return fmt.Sprintf(`"<unhandled-%s>"`, t.Kind().String())
513529
}
514530

531+
func prettyString(s string) string {
532+
// Avoid escaping (which does allocations) if we can.
533+
if needsEscape(s) {
534+
return strconv.Quote(s)
535+
}
536+
b := bytes.NewBuffer(make([]byte, 0, 1024))
537+
b.WriteByte('"')
538+
b.WriteString(s)
539+
b.WriteByte('"')
540+
return b.String()
541+
}
542+
543+
// needsEscape determines whether the input string needs to be escaped or not,
544+
// without doing any allocations.
545+
func needsEscape(s string) bool {
546+
for _, r := range s {
547+
if !strconv.IsPrint(r) || r == '\\' || r == '"' {
548+
return true
549+
}
550+
}
551+
return false
552+
}
553+
515554
func isEmpty(v reflect.Value) bool {
516555
switch v.Kind() {
517556
case reflect.Array, reflect.Map, reflect.Slice, reflect.String:
@@ -675,14 +714,15 @@ func (f *Formatter) AddName(name string) {
675714
func (f *Formatter) AddValues(kvList []interface{}) {
676715
// Three slice args forces a copy.
677716
n := len(f.values)
678-
vals := append(f.values[:n:n], kvList...)
717+
vals := f.values[:n:n]
718+
vals = append(vals, kvList...)
679719
if hook := f.opts.RenderValuesHook; hook != nil {
680720
vals = hook(f.sanitize(vals))
681721
}
682722

683723
// Pre-render values, so we don't have to do it on each Info/Error call.
684724
buf := bytes.NewBuffer(make([]byte, 0, 1024))
685-
f.values = f.flatten(buf, vals, false)
725+
f.values = f.flatten(buf, vals, false, true) // escape user-provided keys
686726
f.valuesStr = buf.String()
687727
}
688728

‎funcr/funcr_test.go

Lines changed: 91 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 {
@@ -36,6 +37,20 @@ func ptrstr(s string) *string {
3637
return &s
3738
}
3839

40+
// point implements encoding.TextMarshaler and can be used as a map key.
41+
type point struct{ x, y int }
42+
43+
func (p point) MarshalText() ([]byte, error) {
44+
return []byte(fmt.Sprintf("(%d, %d)", p.x, p.y)), nil
45+
}
46+
47+
// pointErr implements encoding.TextMarshaler but returns an error.
48+
type pointErr struct{ x, y int }
49+
50+
func (p pointErr) MarshalText() ([]byte, error) {
51+
return nil, fmt.Errorf("uh oh: %d, %d", p.x, p.y)
52+
}
53+
3954
// Logging this should result in the MarshalLog() value.
4055
type Tmarshaler string
4156

@@ -198,7 +213,9 @@ func TestPretty(t *testing.T) {
198213
exp string // used in cases where JSON can't handle it
199214
}{
200215
{val: "strval"},
216+
{val: "strval\nwith\t\"escapes\""},
201217
{val: substr("substrval")},
218+
{val: substr("substrval\nwith\t\"escapes\"")},
202219
{val: true},
203220
{val: false},
204221
{val: int(93)},
@@ -235,7 +252,11 @@ func TestPretty(t *testing.T) {
235252
exp: `[]`,
236253
},
237254
{val: []int{9, 3, 7, 6}},
255+
{val: []string{"str", "with\tescape"}},
256+
{val: []substr{"substr", "with\tescape"}},
238257
{val: [4]int{9, 3, 7, 6}},
258+
{val: [2]string{"str", "with\tescape"}},
259+
{val: [2]substr{"substr", "with\tescape"}},
239260
{
240261
val: struct {
241262
Int int
@@ -255,11 +276,43 @@ func TestPretty(t *testing.T) {
255276
"nine": 3,
256277
},
257278
},
279+
{
280+
val: map[string]int{
281+
"with\tescape": 76,
282+
},
283+
},
258284
{
259285
val: map[substr]int{
260286
"nine": 3,
261287
},
262288
},
289+
{
290+
val: map[substr]int{
291+
"with\tescape": 76,
292+
},
293+
},
294+
{
295+
val: map[int]int{
296+
9: 3,
297+
},
298+
},
299+
{
300+
val: map[float64]int{
301+
9.5: 3,
302+
},
303+
exp: `{"9.5":3}`,
304+
},
305+
{
306+
val: map[point]int{
307+
{x: 1, y: 2}: 3,
308+
},
309+
},
310+
{
311+
val: map[pointErr]int{
312+
{x: 1, y: 2}: 3,
313+
},
314+
exp: `{"<error-MarshalText: uh oh: 1, 2>":3}`,
315+
},
263316
{
264317
val: struct {
265318
X int `json:"x"`
@@ -283,6 +336,7 @@ func TestPretty(t *testing.T) {
283336
val: []struct{ X, Y string }{
284337
{"nine", "three"},
285338
{"seven", "six"},
339+
{"with\t", "\tescapes"},
286340
},
287341
},
288342
{
@@ -438,6 +492,24 @@ func TestPretty(t *testing.T) {
438492
val: PseudoStruct(makeKV("f1", 1, "f2", true, "f3", []int{})),
439493
exp: `{"f1":1,"f2":true,"f3":[]}`,
440494
},
495+
{
496+
val: map[TjsontagsString]int{
497+
{String1: `"quoted"`, String4: `unquoted`}: 1,
498+
},
499+
exp: `{"{\"string1\":\"\\\"quoted\\\"\",\"-\":\"\",\"string4\":\"unquoted\",\"String5\":\"\"}":1}`,
500+
},
501+
{
502+
val: map[TjsontagsInt]int{
503+
{Int1: 1, Int2: 2}: 3,
504+
},
505+
exp: `{"{\"int1\":1,\"-\":0,\"Int5\":0}":3}`,
506+
},
507+
{
508+
val: map[[2]struct{ S string }]int{
509+
{{S: `"quoted"`}, {S: "unquoted"}}: 1,
510+
},
511+
exp: `{"[{\"S\":\"\\\"quoted\\\"\"},{\"S\":\"unquoted\"}]":1}`,
512+
},
441513
}
442514

443515
f := NewFormatter(Options{})
@@ -449,7 +521,7 @@ func TestPretty(t *testing.T) {
449521
} else {
450522
jb, err := json.Marshal(tc.val)
451523
if err != nil {
452-
t.Errorf("[%d]: unexpected error: %v", i, err)
524+
t.Fatalf("[%d]: unexpected error: %v\ngot: %q", i, err, ours)
453525
}
454526
want = string(jb)
455527
}
@@ -496,6 +568,13 @@ func TestRender(t *testing.T) {
496568
args: makeKV("bool", PseudoStruct(makeKV("boolsub", true))),
497569
expectKV: `"int"={"intsub":1} "str"={"strsub":"2"} "bool"={"boolsub":true}`,
498570
expectJSON: `{"int":{"intsub":1},"str":{"strsub":"2"},"bool":{"boolsub":true}}`,
571+
}, {
572+
name: "escapes",
573+
builtins: makeKV("\"1\"", 1), // will not be escaped, but should never happen
574+
values: makeKV("\tstr", "ABC"), // escaped
575+
args: makeKV("bool\n", true), // escaped
576+
expectKV: `""1""=1 "\tstr"="ABC" "bool\n"=true`,
577+
expectJSON: `{""1"":1,"\tstr":"ABC","bool\n":true}`,
499578
}, {
500579
name: "missing value",
501580
builtins: makeKV("builtin"),
@@ -505,27 +584,27 @@ func TestRender(t *testing.T) {
505584
expectJSON: `{"builtin":"<no-value>","value":"<no-value>","arg":"<no-value>"}`,
506585
}, {
507586
name: "non-string key int",
508-
args: makeKV(123, "val"),
587+
builtins: makeKV(123, "val"), // should never happen
509588
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"}`,
589+
args: makeKV(789, "val"),
590+
expectKV: `"<non-string-key: 123>"="val" "<non-string-key: 456>"="val" "<non-string-key: 789>"="val"`,
591+
expectJSON: `{"<non-string-key: 123>":"val","<non-string-key: 456>":"val","<non-string-key: 789>":"val"}`,
513592
}, {
514593
name: "non-string key struct",
515-
args: makeKV(struct {
594+
builtins: makeKV(struct { // will not be escaped, but should never happen
516595
F1 string
517596
F2 int
518-
}{"arg", 123}, "val"),
597+
}{"builtin", 123}, "val"),
519598
values: makeKV(struct {
520599
F1 string
521600
F2 int
522601
}{"value", 456}, "val"),
523-
builtins: makeKV(struct {
602+
args: makeKV(struct {
524603
F1 string
525604
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"}`,
605+
}{"arg", 789}, "val"),
606+
expectKV: `"<non-string-key: {"F1":"builtin",>"="val" "<non-string-key: {\"F1\":\"value\",\"F>"="val" "<non-string-key: {\"F1\":\"arg\",\"F2\">"="val"`,
607+
expectJSON: `{"<non-string-key: {"F1":"builtin",>":"val","<non-string-key: {\"F1\":\"value\",\"F>":"val","<non-string-key: {\"F1\":\"arg\",\"F2\">":"val"}`,
529608
}}
530609

531610
for _, tc := range testCases {
@@ -534,7 +613,7 @@ func TestRender(t *testing.T) {
534613
formatter.AddValues(tc.values)
535614
r := formatter.render(tc.builtins, tc.args)
536615
if r != expect {
537-
t.Errorf("wrong output:\nexpected %q\n got %q", expect, r)
616+
t.Errorf("wrong output:\nexpected %v\n got %v", expect, r)
538617
}
539618
}
540619
t.Run("KV", func(t *testing.T) {

0 commit comments

Comments
 (0)