2 files changed, 43 insertions, 2 deletions

diff --git a/vendor/github.com/NYTimes/gziphandler/gzip.go b/vendor/github.com/NYTimes/gziphandler/gzip.go
index b6af9115a..13ee34a42 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.
@@ -150,7 +151,9 @@ func (w *GzipResponseWriter) startGzip() error {
 
 // WriteHeader just saves the response code until close or GZIP effective writes.
 func (w *GzipResponseWriter) WriteHeader(code int) {
-	w.code = code
+	if w.code == 0 {
+		w.code = code
+	}
 }
 
 // init graps a new gzip writer from the gzipWriterPool and writes the correct
@@ -167,7 +170,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 +199,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
 	}
 }
 
diff --git a/vendor/github.com/NYTimes/gziphandler/gzip_test.go b/vendor/github.com/NYTimes/gziphandler/gzip_test.go
index 655a19373..7decfd17b 100644
--- a/vendor/github.com/NYTimes/gziphandler/gzip_test.go
+++ b/vendor/github.com/NYTimes/gziphandler/gzip_test.go
@@ -306,6 +306,39 @@ func TestStatusCodes(t *testing.T) {
 	}
 }
 
+func TestStatusCodesFlushed(t *testing.T) {
+	handler := GzipHandler(http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
+		rw.WriteHeader(http.StatusNotFound)
+		rw.(http.Flusher).Flush()
+		rw.Write([]byte("Not found"))
+	}))
+	r := httptest.NewRequest(http.MethodGet, "/", nil)
+	r.Header.Set(acceptEncoding, "gzip")
+	w := httptest.NewRecorder()
+	handler.ServeHTTP(w, r)
+
+	result := w.Result()
+	if result.StatusCode != http.StatusNotFound {
+		t.Errorf("StatusCode should have been 404 but was %d", result.StatusCode)
+	}
+}
+
+func TestIgnoreSubsequentWriteHeader(t *testing.T) {
+	handler := GzipHandler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.WriteHeader(500)
+		w.WriteHeader(404)
+	}))
+	r := httptest.NewRequest("GET", "/", nil)
+	r.Header.Set("Accept-Encoding", "gzip")
+	w := httptest.NewRecorder()
+	handler.ServeHTTP(w, r)
+
+	result := w.Result()
+	if result.StatusCode != 500 {
+		t.Errorf("StatusCode should have been 500 but was %d", result.StatusCode)
+	}
+}
+
 func TestDontWriteWhenNotWrittenTo(t *testing.T) {
 	// When using gzip as middleware without ANY writes in the handler,
 	// ensure the gzip middleware doesn't touch the actual ResponseWriter