Skip to content

Commit

Permalink
Merge pull request #113 from thockin/funcr-qstr
Browse files Browse the repository at this point in the history 
funcr: Escape strings when they are not known-safe.
  • Loading branch information
@pohly
pohly authored Oct 14, 2021
2 parents 561de0c + c9b68e6 commit 1215c81
Show file tree
Hide file tree
Showing 2 changed files with 159 additions and 40 deletions.
96 changes: 68 additions & 28 deletions funcr/funcr.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -36,6 +36,7 @@ package funcr

import (
"bytes"
"encoding"
"fmt"
"path/filepath"
"reflect"
Expand Down Expand Up @@ -236,7 +237,7 @@ func (f Formatter) render(builtins, args []interface{}) string {
if hook := f.opts.RenderBuiltinsHook; hook != nil {
vals = hook(f.sanitize(vals))
}
f.flatten(buf, vals, false)
f.flatten(buf, vals, false, false) // keys are ours, no need to escape
continuing := len(builtins) > 0
if len(f.valuesStr) > 0 {
if continuing {
Expand All @@ -253,7 +254,7 @@ func (f Formatter) render(builtins, args []interface{}) string {
if hook := f.opts.RenderArgsHook; hook != nil {
vals = hook(f.sanitize(vals))
}
f.flatten(buf, vals, continuing)
f.flatten(buf, vals, continuing, true) // escape user-provided keys
if f.outputFormat == outputJSON {
buf.WriteByte('}')
}
Expand All @@ -263,10 +264,13 @@ func (f Formatter) render(builtins, args []interface{}) string {
// flatten renders a list of key-value pairs into a buffer. If continuing is
// true, it assumes that the buffer has previous values and will emit a
// separator (which depends on the output format) before the first pair it
// writes. This also returns a potentially modified version of kvList, which
// writes. If escapeKeys is true, the keys are assumed to have
// non-JSON-compatible characters in them and must be evaluated for escapes.
//
// This function returns a potentially modified version of kvList, which
// ensures that there is a value for every key (adding a value if needed) and
// that each key is a string (substituting a key if needed).
func (f Formatter) flatten(buf *bytes.Buffer, kvList []interface{}, continuing bool) []interface{} {
func (f Formatter) flatten(buf *bytes.Buffer, kvList []interface{}, continuing bool, escapeKeys bool) []interface{} {
// This logic overlaps with sanitize() but saves one type-cast per key,
// which can be measurable.
if len(kvList)%2 != 0 {
Expand All @@ -290,9 +294,14 @@ func (f Formatter) flatten(buf *bytes.Buffer, kvList []interface{}, continuing b
}
}

buf.WriteByte('"')
buf.WriteString(k)
buf.WriteByte('"')
if escapeKeys {
buf.WriteString(prettyString(k))
} else {
// this is faster
buf.WriteByte('"')
buf.WriteString(k)
buf.WriteByte('"')
}
if f.outputFormat == outputJSON {
buf.WriteByte(':')
} else {
Expand All @@ -308,8 +317,7 @@ func (f Formatter) pretty(value interface{}) string {
}

const (
flagRawString = 0x1 // do not print quotes on strings
flagRawStruct = 0x2 // do not print braces on structs
flagRawStruct = 0x1 // do not print braces on structs
)

// TODO: This is not fast. Most of the overhead goes here.
Expand All @@ -334,11 +342,7 @@ func (f Formatter) prettyWithFlags(value interface{}, flags uint32) string {
case bool:
return strconv.FormatBool(v)
case string:
if flags&flagRawString > 0 {
return v
}
// This is empirically faster than strings.Builder.
return strconv.Quote(v)
return prettyString(v)
case int:
return strconv.FormatInt(int64(v), 10)
case int8:
Expand Down Expand Up @@ -379,9 +383,8 @@ func (f Formatter) prettyWithFlags(value interface{}, flags uint32) string {
if i > 0 {
buf.WriteByte(',')
}
buf.WriteByte('"')
buf.WriteString(v[i].(string))
buf.WriteByte('"')
// arbitrary keys might need escaping
buf.WriteString(prettyString(v[i].(string)))
buf.WriteByte(':')
buf.WriteString(f.pretty(v[i+1]))
}
Expand All @@ -401,11 +404,7 @@ func (f Formatter) prettyWithFlags(value interface{}, flags uint32) string {
case reflect.Bool:
return strconv.FormatBool(v.Bool())
case reflect.String:
if flags&flagRawString > 0 {
return v.String()
}
// This is empirically faster than strings.Builder.
return `"` + v.String() + `"`
return prettyString(v.String())
case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
return strconv.FormatInt(int64(v.Int()), 10)
case reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64, reflect.Uintptr:
Expand Down Expand Up @@ -463,6 +462,7 @@ func (f Formatter) prettyWithFlags(value interface{}, flags uint32) string {
if name == "" {
name = fld.Name
}
// field names can't contain characters which need escaping
buf.WriteByte('"')
buf.WriteString(name)
buf.WriteByte('"')
Expand Down Expand Up @@ -493,10 +493,26 @@ func (f Formatter) prettyWithFlags(value interface{}, flags uint32) string {
if i > 0 {
buf.WriteByte(',')
}
// JSON only does string keys.
buf.WriteByte('"')
buf.WriteString(f.prettyWithFlags(it.Key().Interface(), flagRawString))
buf.WriteByte('"')
// If a map key supports TextMarshaler, use it.
keystr := ""
if m, ok := it.Key().Interface().(encoding.TextMarshaler); ok {
txt, err := m.MarshalText()
if err != nil {
keystr = fmt.Sprintf("<error-MarshalText: %s>", err.Error())
} else {
keystr = string(txt)
}
keystr = prettyString(keystr)
} else {
// prettyWithFlags will produce already-escaped values
keystr = f.prettyWithFlags(it.Key().Interface(), 0)
if t.Key().Kind() != reflect.String {
// JSON only does string keys. Unlike Go's standard JSON, we'll
// convert just about anything to a string.
keystr = prettyString(keystr)
}
}
buf.WriteString(keystr)
buf.WriteByte(':')
buf.WriteString(f.pretty(it.Value().Interface()))
i++
Expand All @@ -512,6 +528,29 @@ func (f Formatter) prettyWithFlags(value interface{}, flags uint32) string {
return fmt.Sprintf(`"<unhandled-%s>"`, t.Kind().String())
}

func prettyString(s string) string {
// Avoid escaping (which does allocations) if we can.
if needsEscape(s) {
return strconv.Quote(s)
}
b := bytes.NewBuffer(make([]byte, 0, 1024))
b.WriteByte('"')
b.WriteString(s)
b.WriteByte('"')
return b.String()
}

// needsEscape determines whether the input string needs to be escaped or not,
// without doing any allocations.
func needsEscape(s string) bool {
for _, r := range s {
if !strconv.IsPrint(r) || r == '\\' || r == '"' {
return true
}
}
return false
}

func isEmpty(v reflect.Value) bool {
switch v.Kind() {
case reflect.Array, reflect.Map, reflect.Slice, reflect.String:
Expand Down Expand Up @@ -675,14 +714,15 @@ func (f *Formatter) AddName(name string) {
func (f *Formatter) AddValues(kvList []interface{}) {
// Three slice args forces a copy.
n := len(f.values)
vals := append(f.values[:n:n], kvList...)
vals := f.values[:n:n]
vals = append(vals, kvList...)
if hook := f.opts.RenderValuesHook; hook != nil {
vals = hook(f.sanitize(vals))
}

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

Expand Down
103 changes: 91 additions & 12 deletions funcr/funcr_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,7 @@ import (
"github.com/go-logr/logr"
)

// Will be handled via reflection instead of type assertions.
type substr string

func ptrint(i int) *int {
Expand All @@ -36,6 +37,20 @@ func ptrstr(s string) *string {
return &s
}

// point implements encoding.TextMarshaler and can be used as a map key.
type point struct{ x, y int }

func (p point) MarshalText() ([]byte, error) {
return []byte(fmt.Sprintf("(%d, %d)", p.x, p.y)), nil
}

// pointErr implements encoding.TextMarshaler but returns an error.
type pointErr struct{ x, y int }

func (p pointErr) MarshalText() ([]byte, error) {
return nil, fmt.Errorf("uh oh: %d, %d", p.x, p.y)
}

// Logging this should result in the MarshalLog() value.
type Tmarshaler string

Expand Down Expand Up @@ -198,7 +213,9 @@ func TestPretty(t *testing.T) {
exp string // used in cases where JSON can't handle it
}{
{val: "strval"},
{val: "strval\nwith\t\"escapes\""},
{val: substr("substrval")},
{val: substr("substrval\nwith\t\"escapes\"")},
{val: true},
{val: false},
{val: int(93)},
Expand Down Expand Up @@ -235,7 +252,11 @@ func TestPretty(t *testing.T) {
exp: `[]`,
},
{val: []int{9, 3, 7, 6}},
{val: []string{"str", "with\tescape"}},
{val: []substr{"substr", "with\tescape"}},
{val: [4]int{9, 3, 7, 6}},
{val: [2]string{"str", "with\tescape"}},
{val: [2]substr{"substr", "with\tescape"}},
{
val: struct {
Int int
Expand All @@ -255,11 +276,43 @@ func TestPretty(t *testing.T) {
"nine": 3,
},
},
{
val: map[string]int{
"with\tescape": 76,
},
},
{
val: map[substr]int{
"nine": 3,
},
},
{
val: map[substr]int{
"with\tescape": 76,
},
},
{
val: map[int]int{
9: 3,
},
},
{
val: map[float64]int{
9.5: 3,
},
exp: `{"9.5":3}`,
},
{
val: map[point]int{
{x: 1, y: 2}: 3,
},
},
{
val: map[pointErr]int{
{x: 1, y: 2}: 3,
},
exp: `{"<error-MarshalText: uh oh: 1, 2>":3}`,
},
{
val: struct {
X int `json:"x"`
Expand All @@ -283,6 +336,7 @@ func TestPretty(t *testing.T) {
val: []struct{ X, Y string }{
{"nine", "three"},
{"seven", "six"},
{"with\t", "\tescapes"},
},
},
{
Expand Down Expand Up @@ -438,6 +492,24 @@ func TestPretty(t *testing.T) {
val: PseudoStruct(makeKV("f1", 1, "f2", true, "f3", []int{})),
exp: `{"f1":1,"f2":true,"f3":[]}`,
},
{
val: map[TjsontagsString]int{
{String1: `"quoted"`, String4: `unquoted`}: 1,
},
exp: `{"{\"string1\":\"\\\"quoted\\\"\",\"-\":\"\",\"string4\":\"unquoted\",\"String5\":\"\"}":1}`,
},
{
val: map[TjsontagsInt]int{
{Int1: 1, Int2: 2}: 3,
},
exp: `{"{\"int1\":1,\"-\":0,\"Int5\":0}":3}`,
},
{
val: map[[2]struct{ S string }]int{
{{S: `"quoted"`}, {S: "unquoted"}}: 1,
},
exp: `{"[{\"S\":\"\\\"quoted\\\"\"},{\"S\":\"unquoted\"}]":1}`,
},
}

f := NewFormatter(Options{})
Expand All @@ -449,7 +521,7 @@ func TestPretty(t *testing.T) {
} else {
jb, err := json.Marshal(tc.val)
if err != nil {
t.Errorf("[%d]: unexpected error: %v", i, err)
t.Fatalf("[%d]: unexpected error: %v\ngot: %q", i, err, ours)
}
want = string(jb)
}
Expand Down Expand Up @@ -496,6 +568,13 @@ func TestRender(t *testing.T) {
args: makeKV("bool", PseudoStruct(makeKV("boolsub", true))),
expectKV: `"int"={"intsub":1} "str"={"strsub":"2"} "bool"={"boolsub":true}`,
expectJSON: `{"int":{"intsub":1},"str":{"strsub":"2"},"bool":{"boolsub":true}}`,
}, {
name: "escapes",
builtins: makeKV("\"1\"", 1), // will not be escaped, but should never happen
values: makeKV("\tstr", "ABC"), // escaped
args: makeKV("bool\n", true), // escaped
expectKV: `""1""=1 "\tstr"="ABC" "bool\n"=true`,
expectJSON: `{""1"":1,"\tstr":"ABC","bool\n":true}`,
}, {
name: "missing value",
builtins: makeKV("builtin"),
Expand All @@ -505,27 +584,27 @@ func TestRender(t *testing.T) {
expectJSON: `{"builtin":"<no-value>","value":"<no-value>","arg":"<no-value>"}`,
}, {
name: "non-string key int",
args: makeKV(123, "val"),
builtins: makeKV(123, "val"), // should never happen
values: makeKV(456, "val"),
builtins: makeKV(789, "val"),
expectKV: `"<non-string-key: 789>"="val" "<non-string-key: 456>"="val" "<non-string-key: 123>"="val"`,
expectJSON: `{"<non-string-key: 789>":"val","<non-string-key: 456>":"val","<non-string-key: 123>":"val"}`,
args: makeKV(789, "val"),
expectKV: `"<non-string-key: 123>"="val" "<non-string-key: 456>"="val" "<non-string-key: 789>"="val"`,
expectJSON: `{"<non-string-key: 123>":"val","<non-string-key: 456>":"val","<non-string-key: 789>":"val"}`,
}, {
name: "non-string key struct",
args: makeKV(struct {
builtins: makeKV(struct { // will not be escaped, but should never happen
F1 string
F2 int
}{"arg", 123}, "val"),
}{"builtin", 123}, "val"),
values: makeKV(struct {
F1 string
F2 int
}{"value", 456}, "val"),
builtins: makeKV(struct {
args: makeKV(struct {
F1 string
F2 int
}{"builtin", 789}, "val"),
expectKV: `"<non-string-key: {"F1":"builtin",>"="val" "<non-string-key: {"F1":"value","F>"="val" "<non-string-key: {"F1":"arg","F2">"="val"`,
expectJSON: `{"<non-string-key: {"F1":"builtin",>":"val","<non-string-key: {"F1":"value","F>":"val","<non-string-key: {"F1":"arg","F2">":"val"}`,
}{"arg", 789}, "val"),
expectKV: `"<non-string-key: {"F1":"builtin",>"="val" "<non-string-key: {\"F1\":\"value\",\"F>"="val" "<non-string-key: {\"F1\":\"arg\",\"F2\">"="val"`,
expectJSON: `{"<non-string-key: {"F1":"builtin",>":"val","<non-string-key: {\"F1\":\"value\",\"F>":"val","<non-string-key: {\"F1\":\"arg\",\"F2\">":"val"}`,
}}

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

0 comments on commit 1215c81

Please sign in to comment.