Keep status when stream mode and compress

cmd/traefik/anonymize/anonymize_doOnJSON_test.go

@@ -1,8 +1,9 @@
 package anonymize
 
 import (
-	"github.com/stretchr/testify/assert"
 	"testing"
+
+	"github.com/stretchr/testify/assert"
 )
 
 func Test_doOnJSON(t *testing.T) {

glide.lock

@@ -383,7 +383,7 @@ imports:
   repo: https://github.com/ijc25/Gotty.git
   vcs: git
 - name: github.com/NYTimes/gziphandler
-  version: 97ae7fbaf81620fe97840685304a78a306a39c64
+  version: 0f67f3f25d3b17590ee0ab93fcb216363ee30967
 - name: github.com/ogier/pflag
   version: 45c278ab3607870051a2ea9040bb85fcb8557481
 - name: github.com/opencontainers/go-digest

middlewares/compress_test.go

@@ -137,6 +137,31 @@ func TestIntegrationShouldNotCompress(t *testing.T) {
 	}
 }
 
+func TestShouldWriteHeaderWhenFlush(t *testing.T) {
+	comp := &Compress{}
+	negro := negroni.New(comp)
+	negro.UseHandlerFunc(func(rw http.ResponseWriter, r *http.Request) {
+		rw.Header().Add(contentEncodingHeader, gzipValue)
+		rw.Header().Add(varyHeader, acceptEncodingHeader)
+		rw.WriteHeader(http.StatusUnauthorized)
+		rw.(http.Flusher).Flush()
+		rw.Write([]byte("short"))
+	})
+	ts := httptest.NewServer(negro)
+	defer ts.Close()
+
+	req := testhelpers.MustNewRequest(http.MethodGet, ts.URL, nil)
+	req.Header.Add(acceptEncodingHeader, gzipValue)
+
+	resp, err := http.DefaultClient.Do(req)
+	require.NoError(t, err)
+
+	assert.Equal(t, http.StatusUnauthorized, resp.StatusCode)
+
+	assert.Equal(t, gzipValue, resp.Header.Get(contentEncodingHeader))
+	assert.Equal(t, acceptEncodingHeader, resp.Header.Get(varyHeader))
+}
+
 func TestIntegrationShouldCompress(t *testing.T) {
 	fakeBody := generateBytes(100000)
 

middlewares/headers.go

@@ -3,8 +3,9 @@ package middlewares
 //Middleware based on https://github.com/unrolled/secure
 
 import (
-	"github.com/containous/traefik/types"
 	"net/http"
+
+	"github.com/containous/traefik/types"
 )
 
 // HeaderOptions is a struct for specifying configuration options for the headers middleware.

middlewares/headers_test.go

@@ -3,11 +3,12 @@ package middlewares
 //Middleware tests based on https://github.com/unrolled/secure
 
 import (
-	"github.com/containous/traefik/testhelpers"
-	"github.com/stretchr/testify/assert"
 	"net/http"
 	"net/http/httptest"
 	"testing"
+
+	"github.com/containous/traefik/testhelpers"
+	"github.com/stretchr/testify/assert"
 )
 
 var myHandler = http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {

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

@@ -82,6 +82,7 @@ type GzipResponseWriter struct {
 	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.
+	flushed      bool     // Indicate if the stream was already flushed
 }
 
 // Write appends data to the gzip writer.
@@ -167,7 +168,8 @@ func (w *GzipResponseWriter) init() {
 func (w *GzipResponseWriter) Close() error {
 	if w.gw == nil {
 		// Gzip not trigged yet, write out regular response.
-		if w.code != 0 {
+		// WriteHeader only if it wasn't already wrote by a Flush
+		if !w.flushed && w.code != 0 {
 			w.ResponseWriter.WriteHeader(w.code)
 		}
 		if w.buf != nil {
@@ -195,7 +197,11 @@ func (w *GzipResponseWriter) Flush() {
 	}
 
 	if fw, ok := w.ResponseWriter.(http.Flusher); ok {
+		if !w.flushed && w.code != 0 {
+			w.ResponseWriter.WriteHeader(w.code)
+		}
 		fw.Flush()
+		w.flushed = true
 	}
 }