Merge pull request #1082 from imiric/feat/http-zstd-brotli-compression

Add support for Brotli and Zstandard compression

Gopkg.toml

@@ -94,3 +94,11 @@
   [[prune.project]]
     name = "github.com/spf13/cobra"
     unused-packages = false
+
+[[constraint]]
+  name = "github.com/klauspost/compress"
+  version = "1.7.2"
+
+[[constraint]]
+  branch = "master"
+  name = "github.com/andybalholm/brotli"

js/modules/k6/http/request_test.go

@@ -37,7 +37,9 @@ import (
 	"testing"
 	"time"
 
+	"github.com/andybalholm/brotli"
 	"github.com/dop251/goja"
+	"github.com/klauspost/compress/zstd"
 	"github.com/loadimpact/k6/js/common"
 	"github.com/loadimpact/k6/lib"
 	"github.com/loadimpact/k6/lib/metrics"
@@ -314,6 +316,24 @@ func TestRequestAndBatch(t *testing.T) {
 			`))
 			assert.NoError(t, err)
 		})
+		t.Run("zstd", func(t *testing.T) {
+			_, err := common.RunString(rt, sr(`
+				let res = http.get("HTTPSBIN_IP_URL/zstd");
+				if (res.json()['compression'] != 'zstd') {
+					throw new Error("unexpected body data: " + res.json()['compression'])
+				}
+			`))
+			assert.NoError(t, err)
+		})
+		t.Run("brotli", func(t *testing.T) {
+			_, err := common.RunString(rt, sr(`
+				let res = http.get("HTTPSBIN_IP_URL/brotli");
+				if (res.json()['compression'] != 'br') {
+					throw new Error("unexpected body data: " + res.json()['compression'])
+				}
+			`))
+			assert.NoError(t, err)
+		})
 	})
 	t.Run("CompressionWithAcceptEncodingHeader", func(t *testing.T) {
 		t.Run("gzip", func(t *testing.T) {
@@ -1227,6 +1247,9 @@ func TestRequestCompression(t *testing.T) {
 
 	var decompress = func(algo string, input io.Reader) io.Reader {
 		switch algo {
+		case "br":
+			w := brotli.NewReader(input)
+			return w
 		case "gzip":
 			w, err := gzip.NewReader(input)
 			if err != nil {
@@ -1239,6 +1262,12 @@ func TestRequestCompression(t *testing.T) {
 				t.Fatal(err)
 			}
 			return w
+		case "zstd":
+			w, err := zstd.NewReader(input)
+			if err != nil {
+				t.Fatal(err)
+			}
+			return w
 		default:
 			t.Fatal("unknown algorithm " + algo)
 		}
@@ -1288,6 +1317,10 @@ func TestRequestCompression(t *testing.T) {
 		{compression: "deflate"},
 		{compression: "deflate, gzip"},
 		{compression: "gzip,deflate, gzip"},
+		{compression: "zstd"},
+		{compression: "zstd, gzip, deflate"},
+		{compression: "br"},
+		{compression: "br, gzip, deflate"},
 		{
 			compression:   "George",
 			expectedError: `unknown compression algorithm George`,

lib/netext/httpext/request.go

@@ -39,6 +39,8 @@ import (
 
 	ntlmssp "github.com/Azure/go-ntlmssp"
 	digest "github.com/Soontao/goHttpDigestClient"
+	"github.com/andybalholm/brotli"
+	"github.com/klauspost/compress/zstd"
 	"github.com/loadimpact/k6/lib"
 	"github.com/loadimpact/k6/stats"
 	log "github.com/sirupsen/logrus"
@@ -86,7 +88,11 @@ const (
 	CompressionTypeGzip CompressionType = iota
 	// CompressionTypeDeflate compresses through flate
 	CompressionTypeDeflate
-	// TODO: add compress(lzw), brotli maybe bzip2 and others listed at
+	// CompressionTypeZstd compresses through zstd
+	CompressionTypeZstd
+	// CompressionTypeBr compresses through brotli
+	CompressionTypeBr
+	// TODO: add compress(lzw), maybe bzip2 and others listed at
 	// https://en.wikipedia.org/wiki/HTTP_compression#Content-Encoding_tokens
 )
 
@@ -115,6 +121,27 @@ type ParsedHTTPRequest struct {
 	Tags         map[string]string
 }
 
+// Matches non-compliant io.Closer implementations (e.g. zstd.Decoder)
+type ncloser interface {
+	Close()
+}
+
+type readCloser struct {
+	io.Reader
+}
+
+// Close readers with differing Close() implementations
+func (r readCloser) Close() error {
+	var err error
+	switch v := r.Reader.(type) {
+	case io.Closer:
+		err = v.Close()
+	case ncloser:
+		v.Close()
+	}
+	return err
+}
+
 func stdCookiesToHTTPRequestCookies(cookies []*http.Cookie) map[string][]*HTTPRequestCookie {
 	var result = make(map[string][]*HTTPRequestCookie, len(cookies))
 	for _, cookie := range cookies {
@@ -144,6 +171,10 @@ func compressBody(algos []CompressionType, body io.ReadCloser) (io.Reader, int64
 			w = gzip.NewWriter(buf)
 		case CompressionTypeDeflate:
 			w = zlib.NewWriter(buf)
+		case CompressionTypeZstd:
+			w, _ = zstd.NewWriter(buf)
+		case CompressionTypeBr:
+			w = brotli.NewWriter(buf)
 		default:
 			return nil, 0, "", fmt.Errorf("unknown compressionType %s", compressionType)
 		}
@@ -179,15 +210,33 @@ func readResponseBody(
 		return nil, respErr
 	}
 
-	// Transperently decompress the body if it's has a content-encoding we
+	rc := &readCloser{resp.Body}
+	// Ensure that the entire response body is read and closed, e.g. in case of decoding errors
+	defer func(respBody io.ReadCloser) {
+		_, _ = io.Copy(ioutil.Discard, respBody)
+		_ = respBody.Close()
+	}(resp.Body)
+
+	// Transparently decompress the body if it's has a content-encoding we
 	// support. If not, simply return it as it is.
 	contentEncoding := strings.TrimSpace(resp.Header.Get("Content-Encoding"))
 	if compression, err := CompressionTypeString(contentEncoding); err == nil {
+		var decoder io.ReadCloser
 		switch compression {
 		case CompressionTypeDeflate:
-			resp.Body, respErr = zlib.NewReader(resp.Body)
+			decoder, respErr = zlib.NewReader(resp.Body)
+			rc = &readCloser{decoder}
 		case CompressionTypeGzip:
-			resp.Body, respErr = gzip.NewReader(resp.Body)
+			decoder, respErr = gzip.NewReader(resp.Body)
+			rc = &readCloser{decoder}
+		case CompressionTypeZstd:
+			var zstdecoder *zstd.Decoder
+			zstdecoder, respErr = zstd.NewReader(resp.Body)
+			rc = &readCloser{zstdecoder}
+		case CompressionTypeBr:
+			var brdecoder *brotli.Reader
+			brdecoder = brotli.NewReader(resp.Body)
+			rc = &readCloser{brdecoder}
 		default:
 			// We have not implemented a compression ... :(
 			respErr = fmt.Errorf(
@@ -200,8 +249,11 @@ func readResponseBody(
 	buf := state.BPool.Get()
 	defer state.BPool.Put(buf)
 	buf.Reset()
-	_, err := io.Copy(buf, resp.Body)
-	_ = resp.Body.Close()
+	_, err := io.Copy(buf, rc.Reader)
+	if err != nil {
+		respErr = err
+	}
+	err = rc.Close()
 	if err != nil {
 		respErr = err
 	}

lib/testutils/httpmultibin.go

@@ -25,7 +25,9 @@ import (
 	"context"
 	"crypto/tls"
 	"crypto/x509"
+	"encoding/json"
 	"fmt"
+	"io"
 	"net"
 	"net/http"
 	"net/http/httptest"
@@ -35,9 +37,13 @@ import (
 	"testing"
 	"time"
 
+	"github.com/andybalholm/brotli"
 	"github.com/gorilla/websocket"
+	"github.com/klauspost/compress/zstd"
 	"github.com/loadimpact/k6/lib/netext"
+	"github.com/loadimpact/k6/lib/netext/httpext"
 	"github.com/mccutchen/go-httpbin/httpbin"
+	"github.com/pkg/errors"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"golang.org/x/net/http2"
@@ -89,6 +95,11 @@ type HTTPMultiBin struct {
 	Cleanup         func()
 }
 
+type jsonBody struct {
+	Header      http.Header `json:"headers"`
+	Compression string      `json:"compression"`
+}
+
 func getWebsocketEchoHandler(t testing.TB) http.Handler {
 	return http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
 		t.Logf("[%p %s] Upgrading to websocket connection...", req, req.URL)
@@ -116,12 +127,51 @@ func getWebsocketCloserHandler(t testing.TB) http.Handler {
 	})
 }
 
+func writeJSON(w io.Writer, v interface{}) error {
+	e := json.NewEncoder(w)
+	e.SetIndent("", "  ")
+	return errors.Wrap(e.Encode(v), "failed to encode JSON")
+}
+
+func getEncodedHandler(t testing.TB, compressionType httpext.CompressionType) http.Handler {
+	return http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
+		var (
+			encoding string
+			err      error
+			encw     io.WriteCloser
+		)
+
+		switch compressionType {
+		case httpext.CompressionTypeBr:
+			encw = brotli.NewWriter(rw)
+			encoding = "br"
+		case httpext.CompressionTypeZstd:
+			encw, _ = zstd.NewWriter(rw)
+			encoding = "zstd"
+		}
+
+		rw.Header().Set("Content-Type", "application/json")
+		rw.Header().Add("Content-Encoding", encoding)
+		data := jsonBody{
+			Header:      req.Header,
+			Compression: encoding,
+		}
+		err = writeJSON(encw, data)
+		_ = encw.Close()
+		if !assert.NoError(t, err) {
+			return
+		}
+	})
+}
+
 // NewHTTPMultiBin returns a fully configured and running HTTPMultiBin
 func NewHTTPMultiBin(t testing.TB) *HTTPMultiBin {
 	// Create a http.ServeMux and set the httpbin handler as the default
 	mux := http.NewServeMux()
+	mux.Handle("/brotli", getEncodedHandler(t, httpext.CompressionTypeBr))
 	mux.Handle("/ws-echo", getWebsocketEchoHandler(t))
 	mux.Handle("/ws-close", getWebsocketCloserHandler(t))
+	mux.Handle("/zstd", getEncodedHandler(t, httpext.CompressionTypeZstd))
 	mux.Handle("/", httpbin.New().Handler())
 
 	// Initialize the HTTP server and get its details