From d9b795aec8585a0fb435072f68d842d596c332de Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Tom=C3=A1s=20Senart?= Date: Wed, 25 Mar 2020 11:59:52 +0100 Subject: [PATCH] results: Fix Headers encoding (#507) * Revert "results: Faster JSON encoding and decoding (#506)" This reverts commit 3629602dcd773a4db3d41d1addf10e3706a2a6ba. * results: Fix headers encoding --- .gitignore | 1 - go.mod | 2 +- go.sum | 4 +- lib/results.go | 48 ++++++- lib/results_easyjson.go | 223 +++++++++++++++++++++++++++++++++ lib/results_json.go | 268 ---------------------------------------- lib/results_test.go | 60 +++++---- 7 files changed, 304 insertions(+), 302 deletions(-) create mode 100644 lib/results_easyjson.go delete mode 100644 lib/results_json.go diff --git a/.gitignore b/.gitignore index 295f2d43..d0c0e875 100644 --- a/.gitignore +++ b/.gitignore @@ -32,4 +32,3 @@ vendor *.lz .DS_Store -.idea diff --git a/go.mod b/go.mod index 86de5473..d1edc42f 100644 --- a/go.mod +++ b/go.mod @@ -23,6 +23,6 @@ require ( github.com/miekg/dns v1.1.17 github.com/streadway/quantile v0.0.0-20150917103942-b0c588724d25 github.com/tsenart/go-tsz v0.0.0-20180814232043-cdeb9e1e981e - github.com/valyala/fastjson v1.5.0 golang.org/x/net v0.0.0-20190827160401-ba9fcec4b297 + pgregory.net/rapid v0.3.3 ) diff --git a/go.sum b/go.sum index db9a8d1e..514ad3a6 100644 --- a/go.sum +++ b/go.sum @@ -38,8 +38,6 @@ github.com/streadway/quantile v0.0.0-20150917103942-b0c588724d25 h1:7z3LSn867ex6 github.com/streadway/quantile v0.0.0-20150917103942-b0c588724d25/go.mod h1:lbP8tGiBjZ5YWIc2fzuRpTaz0b/53vT6PEs3QuAWzuU= github.com/tsenart/go-tsz v0.0.0-20180814232043-cdeb9e1e981e h1:bB5SXzQmSUsJCmjPDN9fKYx3SSDER5diSjlN6TefTCc= github.com/tsenart/go-tsz v0.0.0-20180814232043-cdeb9e1e981e/go.mod h1:SWZznP1z5Ki7hDT2ioqiFKEse8K9tU2OUvaRI0NeGQo= -github.com/valyala/fastjson v1.5.0 h1:DGrb4wEYso2HdGLyLmNoyNCQnCWfjd8yhghPv5/5YQg= -github.com/valyala/fastjson v1.5.0/go.mod h1:CLCAqky6SMuOcxStkYQvblddUtoRxhYMGLrsQns1aXY= golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w= golang.org/x/crypto v0.0.0-20190829043050-9756ffdc2472 h1:Gv7RPwsi3eZ2Fgewe3CBsuOebPwO27PoXzRpJPsvSSM= golang.org/x/crypto v0.0.0-20190829043050-9756ffdc2472/go.mod h1:yigFU9vqHzYiE8UmvKecakEJjdnWj3jj499lnFckfCI= @@ -59,3 +57,5 @@ golang.org/x/text v0.3.2/go.mod h1:bEr9sfX3Q8Zfm5fL9x+3itogRgK3+ptLWKqgva+5dAk= golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ= golang.org/x/tools v0.0.0-20190907020128-2ca718005c18/go.mod h1:b+2E5dAYhXwXZwtnZ6UAqBI28+e2cm9otk0dWdXHAEo= golang.org/x/xerrors v0.0.0-20190717185122-a985d3407aa7/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= +pgregory.net/rapid v0.3.3 h1:jCjBsY4ln4Atz78QoBWxUEvAHaFyNDQg9+WU62aCn1U= +pgregory.net/rapid v0.3.3/go.mod h1:UYpPVyjFHzYBGHIxLFoupi8vwk6rXNzRY9OMvVxFIOU= diff --git a/lib/results.go b/lib/results.go index 15aa3ef5..eacd42d0 100644 --- a/lib/results.go +++ b/lib/results.go @@ -13,6 +13,9 @@ import ( "strconv" "strings" "time" + + "github.com/mailru/easyjson/jlexer" + "github.com/mailru/easyjson/jwriter" ) func init() { @@ -252,15 +255,48 @@ func NewCSVDecoder(r io.Reader) Decoder { r.Method = rec[9] r.URL = rec[10] - pr := textproto.NewReader(bufio.NewReader( - base64.NewDecoder(base64.StdEncoding, strings.NewReader(rec[11])))) - hdr, err := pr.ReadMIMEHeader() - if err != nil { - return err + if rec[11] != "" { + pr := textproto.NewReader(bufio.NewReader( + base64.NewDecoder(base64.StdEncoding, strings.NewReader(rec[11])))) + hdr, err := pr.ReadMIMEHeader() + if err != nil { + return err + } + r.Headers = http.Header(hdr) } - r.Headers = http.Header(hdr) return err } } +//go:generate easyjson -no_std_marshalers -output_filename results_easyjson.go results.go +//easyjson:json +type jsonResult Result + +// NewJSONEncoder returns an Encoder that dumps the given *Results as a JSON +// object. +func NewJSONEncoder(w io.Writer) Encoder { + var jw jwriter.Writer + return func(r *Result) error { + (*jsonResult)(r).MarshalEasyJSON(&jw) + if jw.Error != nil { + return jw.Error + } + jw.RawByte('\n') + _, err := jw.DumpTo(w) + return err + } +} + +// NewJSONDecoder returns a Decoder that decodes JSON encoded Results. +func NewJSONDecoder(r io.Reader) Decoder { + rd := bufio.NewReader(r) + return func(r *Result) (err error) { + var jl jlexer.Lexer + if jl.Data, err = rd.ReadBytes('\n'); err != nil { + return err + } + (*jsonResult)(r).UnmarshalEasyJSON(&jl) + return jl.Error() + } +} diff --git a/lib/results_easyjson.go b/lib/results_easyjson.go new file mode 100644 index 00000000..3beb22f7 --- /dev/null +++ b/lib/results_easyjson.go @@ -0,0 +1,223 @@ +// Code generated by easyjson for marshaling/unmarshaling. DO NOT EDIT. + +package vegeta + +import ( + json "encoding/json" + easyjson "github.com/mailru/easyjson" + jlexer "github.com/mailru/easyjson/jlexer" + jwriter "github.com/mailru/easyjson/jwriter" + http "net/http" + time "time" +) + +// suppress unused package warning +var ( + _ *json.RawMessage + _ *jlexer.Lexer + _ *jwriter.Writer + _ easyjson.Marshaler +) + +func easyjsonBd1621b8DecodeGithubComTsenartVegetaV12Lib(in *jlexer.Lexer, out *jsonResult) { + isTopLevel := in.IsStart() + if in.IsNull() { + if isTopLevel { + in.Consumed() + } + in.Skip() + return + } + in.Delim('{') + for !in.IsDelim('}') { + key := in.UnsafeString() + in.WantColon() + if in.IsNull() { + in.Skip() + in.WantComma() + continue + } + switch key { + case "attack": + out.Attack = string(in.String()) + case "seq": + out.Seq = uint64(in.Uint64()) + case "code": + out.Code = uint16(in.Uint16()) + case "timestamp": + if data := in.Raw(); in.Ok() { + in.AddError((out.Timestamp).UnmarshalJSON(data)) + } + case "latency": + out.Latency = time.Duration(in.Int64()) + case "bytes_out": + out.BytesOut = uint64(in.Uint64()) + case "bytes_in": + out.BytesIn = uint64(in.Uint64()) + case "error": + out.Error = string(in.String()) + case "body": + if in.IsNull() { + in.Skip() + out.Body = nil + } else { + out.Body = in.Bytes() + } + case "method": + out.Method = string(in.String()) + case "url": + out.URL = string(in.String()) + case "headers": + if in.IsNull() { + in.Skip() + } else { + in.Delim('{') + if !in.IsDelim('}') { + out.Headers = make(http.Header) + } else { + out.Headers = nil + } + for !in.IsDelim('}') { + key := string(in.String()) + in.WantColon() + var v2 []string + if in.IsNull() { + in.Skip() + v2 = nil + } else { + in.Delim('[') + if v2 == nil { + if !in.IsDelim(']') { + v2 = make([]string, 0, 4) + } else { + v2 = []string{} + } + } else { + v2 = (v2)[:0] + } + for !in.IsDelim(']') { + var v3 string + v3 = string(in.String()) + v2 = append(v2, v3) + in.WantComma() + } + in.Delim(']') + } + (out.Headers)[key] = v2 + in.WantComma() + } + in.Delim('}') + } + default: + in.SkipRecursive() + } + in.WantComma() + } + in.Delim('}') + if isTopLevel { + in.Consumed() + } +} +func easyjsonBd1621b8EncodeGithubComTsenartVegetaV12Lib(out *jwriter.Writer, in jsonResult) { + out.RawByte('{') + first := true + _ = first + { + const prefix string = ",\"attack\":" + out.RawString(prefix[1:]) + out.String(string(in.Attack)) + } + { + const prefix string = ",\"seq\":" + out.RawString(prefix) + out.Uint64(uint64(in.Seq)) + } + { + const prefix string = ",\"code\":" + out.RawString(prefix) + out.Uint16(uint16(in.Code)) + } + { + const prefix string = ",\"timestamp\":" + out.RawString(prefix) + out.Raw((in.Timestamp).MarshalJSON()) + } + { + const prefix string = ",\"latency\":" + out.RawString(prefix) + out.Int64(int64(in.Latency)) + } + { + const prefix string = ",\"bytes_out\":" + out.RawString(prefix) + out.Uint64(uint64(in.BytesOut)) + } + { + const prefix string = ",\"bytes_in\":" + out.RawString(prefix) + out.Uint64(uint64(in.BytesIn)) + } + { + const prefix string = ",\"error\":" + out.RawString(prefix) + out.String(string(in.Error)) + } + { + const prefix string = ",\"body\":" + out.RawString(prefix) + out.Base64Bytes(in.Body) + } + { + const prefix string = ",\"method\":" + out.RawString(prefix) + out.String(string(in.Method)) + } + { + const prefix string = ",\"url\":" + out.RawString(prefix) + out.String(string(in.URL)) + } + { + const prefix string = ",\"headers\":" + out.RawString(prefix) + if in.Headers == nil && (out.Flags&jwriter.NilMapAsEmpty) == 0 { + out.RawString(`null`) + } else { + out.RawByte('{') + v6First := true + for v6Name, v6Value := range in.Headers { + if v6First { + v6First = false + } else { + out.RawByte(',') + } + out.String(string(v6Name)) + out.RawByte(':') + if v6Value == nil && (out.Flags&jwriter.NilSliceAsEmpty) == 0 { + out.RawString("null") + } else { + out.RawByte('[') + for v7, v8 := range v6Value { + if v7 > 0 { + out.RawByte(',') + } + out.String(string(v8)) + } + out.RawByte(']') + } + } + out.RawByte('}') + } + } + out.RawByte('}') +} + +// MarshalEasyJSON supports easyjson.Marshaler interface +func (v jsonResult) MarshalEasyJSON(w *jwriter.Writer) { + easyjsonBd1621b8EncodeGithubComTsenartVegetaV12Lib(w, v) +} + +// UnmarshalEasyJSON supports easyjson.Unmarshaler interface +func (v *jsonResult) UnmarshalEasyJSON(l *jlexer.Lexer) { + easyjsonBd1621b8DecodeGithubComTsenartVegetaV12Lib(l, v) +} diff --git a/lib/results_json.go b/lib/results_json.go deleted file mode 100644 index 95c294ea..00000000 --- a/lib/results_json.go +++ /dev/null @@ -1,268 +0,0 @@ -package vegeta - -import ( - "bufio" - "encoding/base64" - "io" - "net/http" - "reflect" - "strconv" - "strings" - "time" - "unsafe" - - "github.com/valyala/fastjson" -) - -// NewJSONEncoder returns an Encoder that dumps the given *Results as a JSON -// object. -func NewJSONEncoder(w io.Writer) Encoder { - buf := make([]byte, 0, 4096) - return func(r *Result) error { - buf = buf[:0] - buf = append(buf, `{"attack":"`...) - if len(r.Attack) > 0 { - buf = appendJSONString(buf, r.Attack) - } - - buf = append(buf, `","seq":`...) - buf = strconv.AppendUint(buf, r.Seq, 10) - buf = append(buf, `,"code":`...) - buf = strconv.AppendUint(buf, uint64(r.Code), 10) - buf = append(buf, `,"timestamp":"`...) - buf = r.Timestamp.AppendFormat(buf, time.RFC3339Nano) - buf = append(buf, `","latency":`...) - buf = strconv.AppendInt(buf, int64(r.Latency), 10) - buf = append(buf, `,"bytes_out":`...) - buf = strconv.AppendUint(buf, r.BytesOut, 10) - buf = append(buf, `,"bytes_in":`...) - buf = strconv.AppendUint(buf, r.BytesIn, 10) - - buf = append(buf, `,"error":"`...) - if len(r.Error) > 0 { - buf = appendJSONString(buf, r.Error) - } - - buf = append(buf, `","body":"`...) - if len(r.Body) > 0 { - buf = appendBase64(buf, r.Body) - } - - buf = append(buf, `","method":"`...) - buf = append(buf, r.Method...) - buf = append(buf, `","url":"`...) - buf = appendJSONString(buf, r.URL) - - buf = append(buf, `","headers":{`...) - for k, vs := range r.Headers { - buf = append(buf, '"') - buf = appendJSONString(buf, k) - buf = append(buf, `":[`...) - for _, v := range vs { - buf = append(buf, '"') - buf = appendJSONString(buf, v) - buf = append(buf, `",`...) - } - if len(vs) > 0 { - buf = buf[:len(buf)-1] - } - buf = append(buf, `],`...) - } - if len(r.Headers) > 0 { - buf = buf[:len(buf)-1] - } - buf = append(buf, "}}\n"...) - - _, err := w.Write(buf) - return err - } -} - -// NewJSONDecoder returns a Decoder that decodes JSON encoded Results. -func NewJSONDecoder(r io.Reader) Decoder { - var p fastjson.Parser - rd := bufio.NewReader(r) - return func(r *Result) (err error) { - line, err := rd.ReadBytes('\n') - if err != nil { - return err - } - - v, err := p.ParseBytes(line) - if err != nil { - return err - } - - r.Attack = string(v.GetStringBytes("attack")) - r.Seq = v.GetUint64("seq") - r.Code = uint16(v.GetUint("code")) - - r.Timestamp, err = time.Parse(time.RFC3339Nano, string(v.GetStringBytes("timestamp"))) - if err != nil { - return err - } - - r.Latency = time.Duration(v.GetInt64("latency")) - r.BytesIn = v.GetUint64("bytes_in") - r.BytesOut = v.GetUint64("bytes_out") - r.Error = string(v.GetStringBytes("error")) - - body := v.GetStringBytes("body") - r.Body = make([]byte, base64.StdEncoding.DecodedLen(len(body))) - n, err := base64.StdEncoding.Decode(r.Body, body) - if err != nil { - return err - } - r.Body = r.Body[:n] - - r.Method = string(v.GetStringBytes("method")) - r.URL = string(v.GetStringBytes("url")) - - headers, err := v.Get("headers").Object() - if err != nil { - return err - } - - r.Headers = make(http.Header, headers.Len()) - headers.Visit(func(key []byte, v *fastjson.Value) { - if err != nil { // Previous visit errored - return - } - - var vs []*fastjson.Value - if vs, err = v.Array(); err != nil { - return - } - - k := string(key) - for _, v := range vs { - r.Headers[k] = append(r.Headers[k], string(v.GetStringBytes())) - } - }) - - return err - } -} - -func appendBase64(buf []byte, bs []byte) []byte { - n := base64.StdEncoding.EncodedLen(len(bs)) - buf = expand(buf, n) - base64.StdEncoding.Encode(buf[len(buf):len(buf)+n], bs) - return buf[:len(buf)+n] -} - -// expand grows the given buf to have enough capacity to hold n -// extra bytes beyond the current len -func expand(buf []byte, n int) []byte { - l := len(buf) - free := cap(buf) - l - grow := n - free - if grow > 0 { - buf = append(buf[:cap(buf)], make([]byte, grow)...)[:l] - } - return buf -} - -// The following code was copied and adapted from https://github.com/valyala/quicktemplate - -func s2b(s string) []byte { - sh := (*reflect.StringHeader)(unsafe.Pointer(&s)) - bh := reflect.SliceHeader{ - Data: sh.Data, - Len: sh.Len, - Cap: sh.Len, - } - return *(*[]byte)(unsafe.Pointer(&bh)) -} - -func b2s(z []byte) string { - return *(*string)(unsafe.Pointer(&z)) -} - -func appendJSONString(buf []byte, s string) []byte { - if len(s) > 24 && - strings.IndexByte(s, '"') < 0 && - strings.IndexByte(s, '\\') < 0 && - strings.IndexByte(s, '\n') < 0 && - strings.IndexByte(s, '\r') < 0 && - strings.IndexByte(s, '\t') < 0 && - strings.IndexByte(s, '\f') < 0 && - strings.IndexByte(s, '\b') < 0 && - strings.IndexByte(s, '<') < 0 && - strings.IndexByte(s, '\'') < 0 && - strings.IndexByte(s, 0) < 0 { - - // fast path - nothing to escape - return append(buf, s2b(s)...) - } - - // slow path - write := func(bs []byte) { buf = append(buf, bs...) } - b := s2b(s) - j := 0 - n := len(b) - if n > 0 { - // Hint the compiler to remove bounds checks in the loop below. - _ = b[n-1] - } - for i := 0; i < n; i++ { - switch b[i] { - case '"': - write(b[j:i]) - write(strBackslashQuote) - j = i + 1 - case '\\': - write(b[j:i]) - write(strBackslashBackslash) - j = i + 1 - case '\n': - write(b[j:i]) - write(strBackslashN) - j = i + 1 - case '\r': - write(b[j:i]) - write(strBackslashR) - j = i + 1 - case '\t': - write(b[j:i]) - write(strBackslashT) - j = i + 1 - case '\f': - write(b[j:i]) - write(strBackslashF) - j = i + 1 - case '\b': - write(b[j:i]) - write(strBackslashB) - j = i + 1 - case '<': - write(b[j:i]) - write(strBackslashLT) - j = i + 1 - case '\'': - write(b[j:i]) - write(strBackslashQ) - j = i + 1 - case 0: - write(b[j:i]) - write(strBackslashZero) - j = i + 1 - } - } - write(b[j:]) - - return buf -} - -var ( - strBackslashQuote = []byte(`\"`) - strBackslashBackslash = []byte(`\\`) - strBackslashN = []byte(`\n`) - strBackslashR = []byte(`\r`) - strBackslashT = []byte(`\t`) - strBackslashF = []byte(`\u000c`) - strBackslashB = []byte(`\u0008`) - strBackslashLT = []byte(`\u003c`) - strBackslashQ = []byte(`\u0027`) - strBackslashZero = []byte(`\u0000`) -) diff --git a/lib/results_test.go b/lib/results_test.go index 8df292be..27e61342 100644 --- a/lib/results_test.go +++ b/lib/results_test.go @@ -9,8 +9,9 @@ import ( "net/http" "reflect" "testing" - "testing/quick" "time" + + "pgregory.net/rapid" ) func TestResultDecoding(t *testing.T) { @@ -80,20 +81,35 @@ func TestResultEncoding(t *testing.T) { t.Run(tc.encoding, func(t *testing.T) { t.Parallel() - err := quick.Check(func(code uint16, ts uint32, latency time.Duration, seq, bsIn, bsOut uint64, body []byte, attack, e string) bool { + rapid.Check(t, func(t *rapid.T) { + hdrs := rapid.MapOf( + rapid.StringMatching(`([\w-]+)`), + rapid.SliceOfN(rapid.StringMatching(`\S`), 1, -1), + ).Draw(t, "headers").(map[string][]string) + want := Result{ - Attack: attack, - Seq: seq, - Code: code, - Timestamp: time.Unix(int64(ts), 0), - Latency: latency, - BytesIn: bsIn, - BytesOut: bsOut, - Error: e, - Body: body, - Method: "GET", - URL: "http://vegeta.test", - Headers: http.Header{"Foo": []string{"bar"}}, + Attack: rapid.String().Draw(t, "attack").(string), + Seq: rapid.Uint64().Draw(t, "seq").(uint64), + Code: rapid.Uint16().Draw(t, "code").(uint16), + Timestamp: time.Unix(rapid.Int64Range(0, 1e8).Draw(t, "timestamp").(int64), 0), + Latency: time.Duration(rapid.Int64Min(0).Draw(t, "latency").(int64)), + BytesIn: rapid.Uint64().Draw(t, "bytes_in").(uint64), + BytesOut: rapid.Uint64().Draw(t, "bytes_out").(uint64), + Error: rapid.String().Draw(t, "error").(string), + Body: rapid.SliceOf(rapid.Byte()).Draw(t, "body").([]byte), + Method: rapid.StringMatching("^(GET|PUT|POST|DELETE|HEAD|OPTIONS)$"). + Draw(t, "method").(string), + URL: rapid.String().Draw(t, "url").(string), + } + + if len(hdrs) > 0 { + want.Headers = make(http.Header, len(hdrs)) + } + + for k, vs := range hdrs { + for _, v := range vs { + want.Headers.Add(k, v) + } } var buf bytes.Buffer @@ -104,6 +120,8 @@ func TestResultEncoding(t *testing.T) { } } + encoded := buf.String() + dec := tc.dec(&buf) if dec == nil { t.Fatal("Cannot get decoder") @@ -111,21 +129,15 @@ func TestResultEncoding(t *testing.T) { for j := 0; j < 2; j++ { var got Result if err := dec(&got); err != nil { - t.Fatalf("err: %q buffer: %s", err, buf.String()) + t.Fatalf("err: %q buffer: %s", err, encoded) } if !got.Equal(want) { - t.Logf("\ngot: %#v\nwant: %#v\n", got, want) - return false + t.Logf("encoded: %s", encoded) + t.Fatalf("\ngot: %#v\nwant: %#v\n", got, want) } } - - return true - }, nil) - - if err != nil { - t.Fatal(err) - } + }) }) } }