From 9bd0fff31975dd856836855770862ac18823e0f2 Mon Sep 17 00:00:00 2001
From: SALLEYRON Julien <julien.salleyron@gmail.com>
Date: Thu, 9 Nov 2017 00:48:03 +0100
Subject: [PATCH] Keep status when stream mode and compress

---
 .../anonymize/anonymize_doOnJSON_test.go      |  3 ++-
 glide.lock                                    |  2 +-
 middlewares/compress_test.go                  | 25 +++++++++++++++++++
 middlewares/headers.go                        |  3 ++-
 middlewares/headers_test.go                   |  5 ++--
 vendor/github.com/NYTimes/gziphandler/gzip.go |  8 +++++-
 6 files changed, 40 insertions(+), 6 deletions(-)

diff --git a/cmd/traefik/anonymize/anonymize_doOnJSON_test.go b/cmd/traefik/anonymize/anonymize_doOnJSON_test.go
index 87f43d125..91ebbf05f 100644
--- a/cmd/traefik/anonymize/anonymize_doOnJSON_test.go
+++ b/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) {
diff --git a/glide.lock b/glide.lock
index 806ea99f0..560b517e4 100644
--- a/glide.lock
+++ b/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
diff --git a/middlewares/compress_test.go b/middlewares/compress_test.go
index 19689dd96..605349e50 100644
--- a/middlewares/compress_test.go
+++ b/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)
 
diff --git a/middlewares/headers.go b/middlewares/headers.go
index ee2368949..de4a251fb 100644
--- a/middlewares/headers.go
+++ b/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.
diff --git a/middlewares/headers_test.go b/middlewares/headers_test.go
index 54dc1e4de..a78a0c2ca 100644
--- a/middlewares/headers_test.go
+++ b/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) {
diff --git a/vendor/github.com/NYTimes/gziphandler/gzip.go b/vendor/github.com/NYTimes/gziphandler/gzip.go
index b6af9115a..91130c424 100644
--- a/vendor/github.com/NYTimes/gziphandler/gzip.go
+++ b/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
 	}
 }