Add ability to compress request body

js/modules/k6/http/request.go

@@ -30,6 +30,7 @@ import (
 	"net/textproto"
 	"net/url"
 	"reflect"
+	"strconv"
 	"strings"
 	"sync"
 	"time"
@@ -235,11 +236,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 +306,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":
@@ -340,6 +352,23 @@ func (h *HTTP) parseRequest(
 		}
 	}
 
+	if contentLength := result.Req.Header.Get("Content-Length"); contentLength != "" {
+		length, err := strconv.Atoi(contentLength)
+		if err == nil {
+			result.Req.ContentLength = int64(length)
+		}
+		// TODO: maybe do something in the other case ... but no error
+	}
+	if result.Body != nil {
+		result.Req.Body = ioutil.NopCloser(result.Body)
+		if result.Req.Header.Get("Content-Length") == "" {
+			result.Req.ContentLength = int64(result.Body.Len())
+		} else {
+			// TODO: print line number, maybe don't print this at all ?
+			state.Logger.Warningf("Content-Length is specifically set won't reset it based on body length")
+		}
+	}
+
 	if result.ActiveJar != nil {
 		httpext.SetRequestCookies(result.Req, result.ActiveJar, result.Cookies)
 	}

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"
@@ -1212,6 +1215,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)

lib/netext/httpext/request.go

@@ -70,6 +70,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 +104,7 @@ type ParsedHTTPRequest struct {
 	Auth         string
 	Throw        bool
 	ResponseType ResponseType
+	Compressions []CompressionType
 	Redirects    null.Int
 	ActiveJar    *cookiejar.Jar
 	Cookies      map[string]*HTTPRequestCookie
@@ -103,6 +120,43 @@ func stdCookiesToHTTPRequestCookies(cookies []*http.Cookie) map[string][]*HTTPRe
 	return result
 }
 
+func compressBody(algos []CompressionType, body io.Reader) (io.Reader, int64, string, error) {
+	var contentEncoding string
+	var prevBuf = 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, nil
+}
+
 // MakeRequest makes http request for tor the provided ParsedHTTPRequest
 //TODO break this function up
 //nolint: gocyclo
@@ -116,6 +170,26 @@ func MakeRequest(ctx context.Context, preq *ParsedHTTPRequest) (*Response, error
 		Headers: preq.Req.Header,
 	}
 	if preq.Body != nil {
+		if 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 {
+				// TODO: print line
+				state.Logger.Warning("Content-Length is specifically set - won't be reset due to compression being specified")
+			}
+			if preq.Req.Header.Get("Content-Encoding") == "" {
+				preq.Req.Header.Set("Content-Encoding", contentEncoding)
+			} else {
+				// TODO: print line
+				state.Logger.Warning("Content-Encoding is specifically set - won't be reset due to compression being specified")
+			}
+		}
+		// TODO: maybe hide this behind of flag in order for this to not happen for big post/puts?
 		respReq.Body = preq.Body.String()
 	}
 

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!