compress: preserve status code

glide.lock

@@ -1,4 +1,4 @@
-hash: c7f74f8b5eac556a531d2cb3bae212e32785c1128d416c75fabedc31d27f0e03
+hash: b056b388f961ddc3509b19b4546d8dfc70fd50ebd1ca47f2fd2de67bf49ba01e
 updated: 2017-08-12T14:15:06.346751095+02:00
 imports:
 - name: cloud.google.com/go
@@ -377,7 +377,7 @@ imports:
   repo: https://github.com/ijc25/Gotty.git
   vcs: git
 - name: github.com/NYTimes/gziphandler
-  version: 316adfc72ed3b0157975917adf62ba2dc31842ce
+  version: 824b33f2a7457025697878c865c323f801118043
   repo: https://github.com/containous/gziphandler.git
   vcs: git
 - name: github.com/ogier/pflag

glide.yaml

@@ -81,6 +81,7 @@ import:
 - package: github.com/NYTimes/gziphandler
   repo: https://github.com/containous/gziphandler.git
   vcs: git
+  version: ^v1002.0.0
 - package: github.com/docker/leadership
 - package: github.com/satori/go.uuid
   version: ^1.1.0

middlewares/compress.go

@@ -17,7 +17,10 @@ func (c *Compress) ServeHTTP(rw http.ResponseWriter, r *http.Request, next http.
 }
 
 func gzipHandler(h http.Handler) http.Handler {
-	wrapper, err := gziphandler.NewGzipHandler(gzip.DefaultCompression, gziphandler.DefaultMinSize, &gziphandler.GzipResponseWriterWrapper{})
+	wrapper, err := gziphandler.GzipHandlerWithOpts(
+		&gziphandler.GzipResponseWriterWrapper{},
+		gziphandler.CompressionLevel(gzip.DefaultCompression),
+		gziphandler.MinSize(gziphandler.DefaultMinSize))
 	if err != nil {
 		log.Error(err)
 	}

middlewares/compress_test.go

@@ -9,6 +9,7 @@ import (
 	"github.com/NYTimes/gziphandler"
 	"github.com/containous/traefik/testhelpers"
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 	"github.com/urfave/negroni"
 )
 
@@ -80,64 +81,115 @@ func TestShouldNotCompressWhenNoAcceptEncodingHeader(t *testing.T) {
 	assert.EqualValues(t, rw.Body.Bytes(), fakeBody)
 }
 
-func TestIntegrationShouldNotCompressWhenContentAlreadyCompressed(t *testing.T) {
+func TestIntegrationShouldNotCompress(t *testing.T) {
 	fakeCompressedBody := generateBytes(100000)
+	comp := &Compress{}
 
-	handler := func(rw http.ResponseWriter, r *http.Request) {
+	testCases := []struct {
+		name               string
+		handler            func(rw http.ResponseWriter, r *http.Request)
+		expectedStatusCode int
+	}{
+		{
+			name: "when content already compressed",
+			handler: func(rw http.ResponseWriter, r *http.Request) {
 				rw.Header().Add(contentEncodingHeader, gzipValue)
 				rw.Header().Add(varyHeader, acceptEncodingHeader)
 				rw.Write(fakeCompressedBody)
+			},
+			expectedStatusCode: http.StatusOK,
+		},
+		{
+			name: "when content already compressed and status code Created",
+			handler: func(rw http.ResponseWriter, r *http.Request) {
+				rw.Header().Add(contentEncodingHeader, gzipValue)
+				rw.Header().Add(varyHeader, acceptEncodingHeader)
+				rw.WriteHeader(http.StatusCreated)
+				rw.Write(fakeCompressedBody)
+			},
+			expectedStatusCode: http.StatusCreated,
+		},
 	}
 
-	comp := &Compress{}
+	for _, test := range testCases {
 
+		t.Run(test.name, func(t *testing.T) {
 			negro := negroni.New(comp)
-	negro.UseHandlerFunc(handler)
+			negro.UseHandlerFunc(test.handler)
 			ts := httptest.NewServer(negro)
 			defer ts.Close()
 
-	client := &http.Client{}
 			req := testhelpers.MustNewRequest(http.MethodGet, ts.URL, nil)
 			req.Header.Add(acceptEncodingHeader, gzipValue)
 
-	resp, err := client.Do(req)
+			resp, err := http.DefaultClient.Do(req)
-	assert.NoError(t, err, "there should be no error")
+			require.NoError(t, err)
+
+			assert.Equal(t, test.expectedStatusCode, resp.StatusCode)
 
 			assert.Equal(t, gzipValue, resp.Header.Get(contentEncodingHeader))
 			assert.Equal(t, acceptEncodingHeader, resp.Header.Get(varyHeader))
 
 			body, err := ioutil.ReadAll(resp.Body)
+			require.NoError(t, err)
 			assert.EqualValues(t, fakeCompressedBody, body)
+		})
+	}
 }
 
-func TestIntegrationShouldCompressWhenAcceptEncodingHeaderIsPresent(t *testing.T) {
+func TestIntegrationShouldCompress(t *testing.T) {
 	fakeBody := generateBytes(100000)
 
-	handler := func(rw http.ResponseWriter, r *http.Request) {
+	testCases := []struct {
+		name               string
+		handler            func(rw http.ResponseWriter, r *http.Request)
+		expectedStatusCode int
+	}{
+		{
+			name: "when AcceptEncoding header is present",
+			handler: func(rw http.ResponseWriter, r *http.Request) {
 				rw.Write(fakeBody)
+			},
+			expectedStatusCode: http.StatusOK,
+		},
+		{
+			name: "when AcceptEncoding header is present and status code Created",
+			handler: func(rw http.ResponseWriter, r *http.Request) {
+				rw.WriteHeader(http.StatusCreated)
+				rw.Write(fakeBody)
+			},
+			expectedStatusCode: http.StatusCreated,
+		},
 	}
 
+	for _, test := range testCases {
+
+		t.Run(test.name, func(t *testing.T) {
 			comp := &Compress{}
 
 			negro := negroni.New(comp)
-	negro.UseHandlerFunc(handler)
+			negro.UseHandlerFunc(test.handler)
 			ts := httptest.NewServer(negro)
 			defer ts.Close()
 
-	client := &http.Client{}
 			req := testhelpers.MustNewRequest(http.MethodGet, ts.URL, nil)
 			req.Header.Add(acceptEncodingHeader, gzipValue)
 
-	resp, err := client.Do(req)
+			resp, err := http.DefaultClient.Do(req)
-	assert.NoError(t, err, "there should be no error")
+			require.NoError(t, err)
+
+			assert.Equal(t, test.expectedStatusCode, resp.StatusCode)
 
 			assert.Equal(t, gzipValue, resp.Header.Get(contentEncodingHeader))
 			assert.Equal(t, acceptEncodingHeader, resp.Header.Get(varyHeader))
 
 			body, err := ioutil.ReadAll(resp.Body)
+			require.NoError(t, err)
 			if assert.ObjectsAreEqualValues(body, fakeBody) {
 				assert.Fail(t, "expected a compressed body", "got %v", body)
 			}
+		})
+	}
 }
 
 func generateBytes(len int) []byte {

vendor/github.com/NYTimes/gziphandler/gzip.go

@@ -81,6 +81,8 @@ type GzipResponseWriter struct {
 
 	minSize int    // Specifed the minimum response size to gzip. If the response length is bigger than this value, it is compressed.
 	buf     []byte // Holds the first part of the write before reaching the minSize or the end of the write.
+
+	contentTypes []string // Only compress if the response is one of these content-types. All are accepted if empty.
 }
 
 // Write appends data to the gzip writer.
@@ -101,8 +103,10 @@ func (w *GzipResponseWriter) Write(b []byte) (int, error) {
 	// On the first write, w.buf changes from nil to a valid slice
 	w.buf = append(w.buf, b...)
 
-	// If the global writes are bigger than the minSize, compression is enable.
+	// If the global writes are bigger than the minSize and we're about to write
-	if len(w.buf) >= w.minSize {
+	// a response containing a content type we want to handle, enable
+	// compression.
+	if len(w.buf) >= w.minSize && handleContentType(w.contentTypes, w) {
 		err := w.startGzip()
 		if err != nil {
 			return 0, err
@@ -231,24 +235,29 @@ func NewGzipLevelHandler(level int) (func(http.Handler) http.Handler, error) {
 // NewGzipLevelAndMinSize behave as NewGzipLevelHandler except it let the caller
 // specify the minimum size before compression.
 func NewGzipLevelAndMinSize(level, minSize int) (func(http.Handler) http.Handler, error) {
-	return NewGzipHandler(level, minSize, &GzipResponseWriter{})
+	return GzipHandlerWithOpts(&GzipResponseWriter{}, CompressionLevel(level), MinSize(minSize))
 }
 
-// NewGzipHandler behave as NewGzipLevelHandler except it let the caller
+func GzipHandlerWithOpts(gw GzipWriter, opts ...option) (func(http.Handler) http.Handler, error) {
-// specify the minimum size before compression and a GzipWriter.
-func NewGzipHandler(level, minSize int, gw GzipWriter) (func(http.Handler) http.Handler, error) {
-	if level != gzip.DefaultCompression && (level < gzip.BestSpeed || level > gzip.BestCompression) {
-		return nil, fmt.Errorf("invalid compression level requested: %d", level)
-	}
-	if minSize < 0 {
-		return nil, errors.New("minimum size must be more than zero")
-	}
 	if gw == nil {
 		return nil, errors.New("the GzipWriter must be defined")
 	}
 
+	c := &config{
+		level:   gzip.DefaultCompression,
+		minSize: DefaultMinSize,
+	}
+
+	for _, o := range opts {
+		o(c)
+	}
+
+	if err := c.validate(); err != nil {
+		return nil, err
+	}
+
 	return func(h http.Handler) http.Handler {
-		index := poolIndex(level)
+		index := poolIndex(c.level)
 
 		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 			w.Header().Add(vary, acceptEncoding)
@@ -256,7 +265,8 @@ func NewGzipHandler(level, minSize int, gw GzipWriter) (func(http.Handler) http.
 			if acceptsGzip(r) {
 				gw.SetResponseWriter(w)
 				gw.setIndex(index)
-				gw.setMinSize(minSize)
+				gw.setMinSize(c.minSize)
+				gw.setContentTypes(c.contentTypes)
 				defer gw.Close()
 
 				h.ServeHTTP(gw, r)
@@ -267,6 +277,48 @@ func NewGzipHandler(level, minSize int, gw GzipWriter) (func(http.Handler) http.
 	}, nil
 }
 
+// Used for functional configuration.
+type config struct {
+	minSize      int
+	level        int
+	contentTypes []string
+}
+
+func (c *config) validate() error {
+	if c.level != gzip.DefaultCompression && (c.level < gzip.BestSpeed || c.level > gzip.BestCompression) {
+		return fmt.Errorf("invalid compression level requested: %d", c.level)
+	}
+
+	if c.minSize < 0 {
+		return fmt.Errorf("minimum size must be more than zero")
+	}
+
+	return nil
+}
+
+type option func(c *config)
+
+func MinSize(size int) option {
+	return func(c *config) {
+		c.minSize = size
+	}
+}
+
+func CompressionLevel(level int) option {
+	return func(c *config) {
+		c.level = level
+	}
+}
+
+func ContentTypes(types []string) option {
+	return func(c *config) {
+		c.contentTypes = []string{}
+		for _, v := range types {
+			c.contentTypes = append(c.contentTypes, strings.ToLower(v))
+		}
+	}
+}
+
 // GzipHandler wraps an HTTP handler, to transparently gzip the response body if
 // the client supports it (via the Accept-Encoding header). This will compress at
 // the default compression level.
@@ -282,6 +334,23 @@ func acceptsGzip(r *http.Request) bool {
 	return acceptedEncodings["gzip"] > 0.0
 }
 
+// returns true if we've been configured to compress the specific content type.
+func handleContentType(contentTypes []string, w http.ResponseWriter) bool {
+	// If contentTypes is empty we handle all content types.
+	if len(contentTypes) == 0 {
+		return true
+	}
+
+	ct := strings.ToLower(w.Header().Get(contentType))
+	for _, c := range contentTypes {
+		if c == ct {
+			return true
+		}
+	}
+
+	return false
+}
+
 // parseEncodings attempts to parse a list of codings, per RFC 2616, as might
 // appear in an Accept-Encoding header. It returns a map of content-codings to
 // quality values, and an error containing the errors encountered. It's probably

vendor/github.com/NYTimes/gziphandler/wrapper.go

@@ -23,6 +23,7 @@ type GzipWriter interface {
 	SetResponseWriter(http.ResponseWriter)
 	setIndex(int)
 	setMinSize(int)
+	setContentTypes([]string)
 }
 
 func (w *GzipResponseWriter) SetResponseWriter(rw http.ResponseWriter) {
@@ -37,6 +38,10 @@ func (w *GzipResponseWriter) setMinSize(minSize int) {
 	w.minSize = minSize
 }
 
+func (w *GzipResponseWriter) setContentTypes(contentTypes []string) {
+	w.contentTypes = contentTypes
+}
+
 // --------
 
 type GzipResponseWriterWrapper struct {
@@ -45,6 +50,9 @@ type GzipResponseWriterWrapper struct {
 
 func (g *GzipResponseWriterWrapper) Write(b []byte) (int, error) {
 	if g.gw == nil && isEncoded(g.Header()) {
+		if g.code != 0 {
+			g.ResponseWriter.WriteHeader(g.code)
+		}
 		return g.ResponseWriter.Write(b)
 	}
 	return g.GzipResponseWriter.Write(b)