From 3fc300cd779d9ce775eb299d0f0dbde7a620ce75 Mon Sep 17 00:00:00 2001 From: Mihail Stoykov Date: Fri, 12 Apr 2019 11:49:59 +0300 Subject: [PATCH] Add ability to compress request body (#989) This adds `compression` param to all http methods with that will compress the body set the correct `Content-Encoding` and `Content-Length` and will send the request with the compressed body. The current compressions supported are `gzip` and `deflate` and their combinations separated by a comma (`,`). If Content-Length of Content-Encoding has been set manually by the user, they won't be reset but instead a warning will be printed. --- js/modules/k6/http/request.go | 22 +++- js/modules/k6/http/request_test.go | 142 +++++++++++++++++++++ lib/netext/httpext/compression_type_gen.go | 49 +++++++ lib/netext/httpext/request.go | 114 ++++++++++++++++- lib/netext/httpext/request_test.go | 84 ++++++++++++ release notes/upcoming.md | 6 +- 6 files changed, 402 insertions(+), 15 deletions(-) create mode 100644 lib/netext/httpext/compression_type_gen.go create mode 100644 lib/netext/httpext/request_test.go diff --git a/js/modules/k6/http/request.go b/js/modules/k6/http/request.go index 36fcdab0eb8..d448c3a0118 100644 --- a/js/modules/k6/http/request.go +++ b/js/modules/k6/http/request.go @@ -24,7 +24,6 @@ import ( "bytes" "context" "fmt" - "io/ioutil" "mime/multipart" "net/http" "net/textproto" @@ -235,11 +234,6 @@ func (h *HTTP) parseRequest( } } - if result.Body != nil { - result.Req.Body = ioutil.NopCloser(result.Body) - result.Req.ContentLength = int64(result.Body.Len()) - } - if userAgent := state.Options.UserAgent; userAgent.String != "" { result.Req.Header.Set("User-Agent", userAgent.String) } @@ -310,6 +304,22 @@ func (h *HTTP) parseRequest( case *HTTPCookieJar: result.ActiveJar = v.jar } + case "compression": + var algosString = strings.TrimSpace(params.Get(k).ToString().String()) + if algosString == "" { + continue + } + var algos = strings.Split(algosString, ",") + var err error + result.Compressions = make([]httpext.CompressionType, len(algos)) + for index, algo := range algos { + algo = strings.TrimSpace(algo) + result.Compressions[index], err = httpext.CompressionTypeString(algo) + if err != nil { + return nil, fmt.Errorf("unknown compression algorithm %s, supported algorithms are %s", + algo, httpext.CompressionTypeValues()) + } + } case "redirects": result.Redirects = null.IntFrom(params.Get(k).ToInteger()) case "tags": diff --git a/js/modules/k6/http/request_test.go b/js/modules/k6/http/request_test.go index 764ce8b7915..2e48e67b5ab 100644 --- a/js/modules/k6/http/request_test.go +++ b/js/modules/k6/http/request_test.go @@ -22,8 +22,11 @@ package http import ( "bytes" + "compress/gzip" + "compress/zlib" "context" "fmt" + "io" "io/ioutil" "net/http" "net/http/cookiejar" @@ -1208,6 +1211,145 @@ func TestSystemTags(t *testing.T) { } } +func TestRequestCompression(t *testing.T) { + t.Parallel() + tb, state, _, rt, _ := newRuntime(t) + defer tb.Cleanup() + + // We don't expect any failed requests + state.Options.Throw = null.BoolFrom(true) + + var text = ` + Lorem ipsum dolor sit amet, consectetur adipiscing elit. + Maecenas sed pharetra sapien. Nunc laoreet molestie ante ac gravida. + Etiam interdum dui viverra posuere egestas. Pellentesque at dolor tristique, + mattis turpis eget, commodo purus. Nunc orci aliquam.` + + var decompress = func(algo string, input io.Reader) io.Reader { + switch algo { + case "gzip": + w, err := gzip.NewReader(input) + if err != nil { + t.Fatal(err) + } + return w + case "deflate": + w, err := zlib.NewReader(input) + if err != nil { + t.Fatal(err) + } + return w + default: + t.Fatal("unknown algorithm " + algo) + } + return nil // unreachable + } + + var ( + expectedEncoding string + actualEncoding string + ) + tb.Mux.HandleFunc("/compressed-text", http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + require.Equal(t, expectedEncoding, r.Header.Get("Content-Encoding")) + + expectedLength, err := strconv.Atoi(r.Header.Get("Content-Length")) + require.NoError(t, err) + var algos = strings.Split(actualEncoding, ", ") + var compressedBuf = new(bytes.Buffer) + n, err := io.Copy(compressedBuf, r.Body) + require.Equal(t, int(n), expectedLength) + require.NoError(t, err) + var prev io.Reader = compressedBuf + + if expectedEncoding != "" { + for i := len(algos) - 1; i >= 0; i-- { + prev = decompress(algos[i], prev) + } + } + + var buf bytes.Buffer + _, err = io.Copy(&buf, prev) + require.NoError(t, err) + require.Equal(t, text, buf.String()) + })) + + var testCases = []struct { + name string + compression string + expectedError string + }{ + {compression: ""}, + {compression: " "}, + {compression: "gzip"}, + {compression: "gzip, gzip"}, + {compression: "gzip, gzip "}, + {compression: "gzip,gzip"}, + {compression: "gzip, gzip, gzip, gzip, gzip, gzip, gzip"}, + {compression: "deflate"}, + {compression: "deflate, gzip"}, + {compression: "gzip,deflate, gzip"}, + { + compression: "George", + expectedError: `unknown compression algorithm George`, + }, + { + compression: "gzip, George", + expectedError: `unknown compression algorithm George`, + }, + } + for _, testCase := range testCases { + testCase := testCase + t.Run(testCase.compression, func(t *testing.T) { + var algos = strings.Split(testCase.compression, ",") + for i, algo := range algos { + algos[i] = strings.TrimSpace(algo) + } + expectedEncoding = strings.Join(algos, ", ") + actualEncoding = expectedEncoding + _, err := common.RunString(rt, tb.Replacer.Replace(` + http.post("HTTPBIN_URL/compressed-text", `+"`"+text+"`"+`, {"compression": "`+testCase.compression+`"}); + `)) + if testCase.expectedError == "" { + require.NoError(t, err) + } else { + require.Error(t, err) + require.Contains(t, err.Error(), testCase.expectedError) + } + + }) + } + + t.Run("custom set header", func(t *testing.T) { + expectedEncoding = "not, valid" + actualEncoding = "gzip, deflate" + t.Run("encoding", func(t *testing.T) { + _, err := common.RunString(rt, tb.Replacer.Replace(` + http.post("HTTPBIN_URL/compressed-text", `+"`"+text+"`"+`, + {"compression": "`+actualEncoding+`", + "headers": {"Content-Encoding": "`+expectedEncoding+`"} + } + ); + `)) + require.NoError(t, err) + + }) + + t.Run("encoding and length", func(t *testing.T) { + _, err := common.RunString(rt, tb.Replacer.Replace(` + http.post("HTTPBIN_URL/compressed-text", `+"`"+text+"`"+`, + {"compression": "`+actualEncoding+`", + "headers": {"Content-Encoding": "`+expectedEncoding+`", + "Content-Length": "12"} + } + ); + `)) + require.Error(t, err) + // TODO: This probably shouldn't be like this + require.Contains(t, err.Error(), "http: ContentLength=12 with Body length 211") + }) + }) +} + func TestResponseTypes(t *testing.T) { t.Parallel() tb, state, _, rt, _ := newRuntime(t) diff --git a/lib/netext/httpext/compression_type_gen.go b/lib/netext/httpext/compression_type_gen.go new file mode 100644 index 00000000000..c2ee00e79ce --- /dev/null +++ b/lib/netext/httpext/compression_type_gen.go @@ -0,0 +1,49 @@ +// Code generated by "enumer -type=CompressionType -transform=snake -trimprefix CompressionType -output compression_type_gen.go"; DO NOT EDIT. + +package httpext + +import ( + "fmt" +) + +const _CompressionTypeName = "gzipdeflate" + +var _CompressionTypeIndex = [...]uint8{0, 4, 11} + +func (i CompressionType) String() string { + if i >= CompressionType(len(_CompressionTypeIndex)-1) { + return fmt.Sprintf("CompressionType(%d)", i) + } + return _CompressionTypeName[_CompressionTypeIndex[i]:_CompressionTypeIndex[i+1]] +} + +var _CompressionTypeValues = []CompressionType{0, 1} + +var _CompressionTypeNameToValueMap = map[string]CompressionType{ + _CompressionTypeName[0:4]: 0, + _CompressionTypeName[4:11]: 1, +} + +// CompressionTypeString retrieves an enum value from the enum constants string name. +// Throws an error if the param is not part of the enum. +func CompressionTypeString(s string) (CompressionType, error) { + if val, ok := _CompressionTypeNameToValueMap[s]; ok { + return val, nil + } + return 0, fmt.Errorf("%s does not belong to CompressionType values", s) +} + +// CompressionTypeValues returns all values of the enum +func CompressionTypeValues() []CompressionType { + return _CompressionTypeValues +} + +// IsACompressionType returns "true" if the value is listed in the enum definition. "false" otherwise +func (i CompressionType) IsACompressionType() bool { + for _, v := range _CompressionTypeValues { + if i == v { + return true + } + } + return false +} diff --git a/lib/netext/httpext/request.go b/lib/netext/httpext/request.go index a36cbc1b01f..5d3dc701812 100644 --- a/lib/netext/httpext/request.go +++ b/lib/netext/httpext/request.go @@ -45,6 +45,10 @@ import ( null "gopkg.in/guregu/null.v3" ) +const compressionHeaderOverwriteMessage = "Both compression and the `%s` header were specified " + + "in the %s request for '%s', the custom header has precedence and won't be overwritten. " + + "This will likely result in invalid data being sent to the server." + // HTTPRequestCookie is a representation of a cookie used for request objects type HTTPRequestCookie struct { Name, Value string @@ -70,6 +74,22 @@ func (u URL) GetURL() *url.URL { return u.u } +// CompressionType is used to specify what compression is to be used to compress the body of a +// request +// The conversion and validation methods are auto-generated with https://github.com/alvaroloes/enumer: +//nolint: lll +//go:generate enumer -type=CompressionType -transform=snake -trimprefix CompressionType -output compression_type_gen.go +type CompressionType uint + +const ( + // CompressionTypeGzip compresses through gzip + CompressionTypeGzip CompressionType = iota + // CompressionTypeDeflate compresses through flate + CompressionTypeDeflate + // TODO: add compress(lzw), brotli maybe bzip2 and others listed at + // https://en.wikipedia.org/wiki/HTTP_compression#Content-Encoding_tokens +) + // Request represent an http request type Request struct { Method string `json:"method"` @@ -88,6 +108,7 @@ type ParsedHTTPRequest struct { Auth string Throw bool ResponseType ResponseType + Compressions []CompressionType Redirects null.Int ActiveJar *cookiejar.Jar Cookies map[string]*HTTPRequestCookie @@ -103,6 +124,44 @@ func stdCookiesToHTTPRequestCookies(cookies []*http.Cookie) map[string][]*HTTPRe return result } +func compressBody(algos []CompressionType, body io.ReadCloser) (io.Reader, int64, string, error) { + var contentEncoding string + var prevBuf io.Reader = body + var buf *bytes.Buffer + for _, compressionType := range algos { + if buf != nil { + prevBuf = buf + } + buf = new(bytes.Buffer) + + if contentEncoding != "" { + contentEncoding += ", " + } + contentEncoding += compressionType.String() + var w io.WriteCloser + switch compressionType { + case CompressionTypeGzip: + w = gzip.NewWriter(buf) + case CompressionTypeDeflate: + w = zlib.NewWriter(buf) + default: + return nil, 0, "", fmt.Errorf("unknown compressionType %s", compressionType) + } + // we don't close in defer because zlib will write it's checksum again if it closes twice :( + var _, err = io.Copy(w, prevBuf) + if err != nil { + _ = w.Close() + return nil, 0, "", err + } + + if err = w.Close(); err != nil { + return nil, 0, "", err + } + } + + return buf, int64(buf.Len()), contentEncoding, body.Close() +} + // MakeRequest makes http request for tor the provided ParsedHTTPRequest //TODO break this function up //nolint: gocyclo @@ -115,8 +174,45 @@ func MakeRequest(ctx context.Context, preq *ParsedHTTPRequest) (*Response, error Cookies: stdCookiesToHTTPRequestCookies(preq.Req.Cookies()), Headers: preq.Req.Header, } + + if contentLength := preq.Req.Header.Get("Content-Length"); contentLength != "" { + length, err := strconv.Atoi(contentLength) + if err == nil { + preq.Req.ContentLength = int64(length) + } + // TODO: maybe do something in the other case ... but no error + } + if preq.Body != nil { + preq.Req.Body = ioutil.NopCloser(preq.Body) + + // TODO: maybe hide this behind of flag in order for this to not happen for big post/puts? + // should we set this after the compression? what will be the point ? respReq.Body = preq.Body.String() + + switch { + case len(preq.Compressions) > 0: + compressedBody, length, contentEncoding, err := compressBody(preq.Compressions, preq.Req.Body) + if err != nil { + return nil, err + } + + preq.Req.Body = ioutil.NopCloser(compressedBody) + if preq.Req.Header.Get("Content-Length") == "" { + preq.Req.ContentLength = length + } else { + state.Logger.Warningf(compressionHeaderOverwriteMessage, "Content-Length", preq.Req.Method, preq.Req.URL) + } + if preq.Req.Header.Get("Content-Encoding") == "" { + preq.Req.Header.Set("Content-Encoding", contentEncoding) + } else { + state.Logger.Warningf(compressionHeaderOverwriteMessage, "Content-Encoding", preq.Req.Method, preq.Req.URL) + } + case preq.Req.Header.Get("Content-Length") == "": + preq.Req.ContentLength = int64(preq.Body.Len()) + } + // TODO: print some message in case we have Content-Length set so that we can warn users + // that setting it manually can lead to bad requests } tags := state.Options.RunTags.CloneTags() @@ -244,11 +340,19 @@ func MakeRequest(ctx context.Context, preq *ParsedHTTPRequest) (*Response, error resp.Error = tracerTransport.errorMsg resp.ErrorCode = int(tracerTransport.errorCode) if resErr == nil && res != nil { - switch res.Header.Get("Content-Encoding") { - case "deflate": - res.Body, resErr = zlib.NewReader(res.Body) - case "gzip": - res.Body, resErr = gzip.NewReader(res.Body) + compression, err := CompressionTypeString(strings.TrimSpace(res.Header.Get("Content-Encoding"))) + if err == nil { // in case of error we just won't uncompress + switch compression { + case CompressionTypeDeflate: + res.Body, resErr = zlib.NewReader(res.Body) + case CompressionTypeGzip: + res.Body, resErr = gzip.NewReader(res.Body) + default: + // We have not implemented a compression ... :( + resErr = fmt.Errorf( + "unsupported compressionType %s for uncompression. This is a bug in k6 please report it", + compression) + } } } if resErr == nil && res != nil { diff --git a/lib/netext/httpext/request_test.go b/lib/netext/httpext/request_test.go new file mode 100644 index 00000000000..b0a5fde8682 --- /dev/null +++ b/lib/netext/httpext/request_test.go @@ -0,0 +1,84 @@ +package httpext + +import ( + "bytes" + "context" + "io" + "io/ioutil" + "net/http" + "testing" + + "github.com/pkg/errors" + "github.com/stretchr/testify/require" +) + +type reader func([]byte) (int, error) + +func (r reader) Read(a []byte) (int, error) { + return ((func([]byte) (int, error))(r))(a) +} + +const badReadMsg = "bad read error for test" +const badCloseMsg = "bad close error for test" + +func badReadBody() io.Reader { + return reader(func(_ []byte) (int, error) { + return 0, errors.New(badReadMsg) + }) +} + +type closer func() error + +func (c closer) Close() error { + return ((func() error)(c))() +} + +func badCloseBody() io.ReadCloser { + return struct { + io.Reader + io.Closer + }{ + Reader: reader(func(_ []byte) (int, error) { + return 0, io.EOF + }), + Closer: closer(func() error { + return errors.New(badCloseMsg) + }), + } +} + +func TestCompressionBodyError(t *testing.T) { + var algos = []CompressionType{CompressionTypeGzip} + t.Run("bad read body", func(t *testing.T) { + _, _, _, err := compressBody(algos, ioutil.NopCloser(badReadBody())) + require.Error(t, err) + require.Equal(t, err.Error(), badReadMsg) + }) + + t.Run("bad close body", func(t *testing.T) { + _, _, _, err := compressBody(algos, badCloseBody()) + require.Error(t, err) + require.Equal(t, err.Error(), badCloseMsg) + }) +} + +func TestMakeRequestError(t *testing.T) { + var ctx, cancel = context.WithCancel(context.Background()) + defer cancel() + + t.Run("bad compression algorithm body", func(t *testing.T) { + var req, err = http.NewRequest("GET", "https://wont.be.used", nil) + + require.NoError(t, err) + var badCompressionType = CompressionType(13) + require.False(t, badCompressionType.IsACompressionType()) + var preq = &ParsedHTTPRequest{ + Req: req, + Body: new(bytes.Buffer), + Compressions: []CompressionType{badCompressionType}, + } + _, err = MakeRequest(ctx, preq) + require.Error(t, err) + require.Equal(t, err.Error(), "unknown compressionType CompressionType(13)") + }) +} diff --git a/release notes/upcoming.md b/release notes/upcoming.md index 0ad973442e2..bb1607bee23 100644 --- a/release notes/upcoming.md +++ b/release notes/upcoming.md @@ -2,11 +2,9 @@ TODO: Intro ## New Features! -### Category: Title (#533) +### HTTP: request body compression (#988) -Description of feature. - -**Docs**: [Title](http://k6.readme.io/docs/TODO) +Now all http methods have an additional param called `compression` that will make k6 compress the body before sending it. It will also correctly set both `Content-Encoding` and `Content-Length`, unless they were manually set in the request `headers` by the user. The current supported algorithms are `deflate` and `gzip` and any combination of the two separated by a comma (`,`). ## Bugs fixed!