From edf8477e0cb3b93a0a1e0550cb0869e1f64f6398 Mon Sep 17 00:00:00 2001
From: James Hartig <fastest963@gmail.com>
Date: Thu, 16 Aug 2018 16:01:59 -0400
Subject: [PATCH] Fixed double WriteHeader calls caused by #71 (#73)

---
 gzip.go      |  4 ++++
 gzip_test.go | 52 ++++++++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 56 insertions(+)

diff --git a/gzip.go b/gzip.go
index 8989961..6b71623 100644
--- a/gzip.go
+++ b/gzip.go
@@ -159,6 +159,8 @@ func (w *GzipResponseWriter) startGzip() error {
 	// Write the header to gzip response.
 	if w.code != 0 {
 		w.ResponseWriter.WriteHeader(w.code)
+		// Ensure that no other WriteHeader's happen
+		w.code = 0
 	}
 
 	// Initialize and flush the buffer into the gzip response if there are any bytes.
@@ -184,6 +186,8 @@ func (w *GzipResponseWriter) startGzip() error {
 func (w *GzipResponseWriter) startPlain() error {
 	if w.code != 0 {
 		w.ResponseWriter.WriteHeader(w.code)
+		// Ensure that no other WriteHeader's happen
+		w.code = 0
 	}
 	w.ignore = true
 	// If Write was never called then don't call Write on the underlying ResponseWriter.
diff --git a/gzip_test.go b/gzip_test.go
index f97ea4e..bed7f52 100644
--- a/gzip_test.go
+++ b/gzip_test.go
@@ -346,6 +346,58 @@ func TestGzipDoubleClose(t *testing.T) {
 	assert.False(t, w1 == w2)
 }
 
+type panicOnSecondWriteHeaderWriter struct {
+	http.ResponseWriter
+	headerWritten bool
+}
+
+func (w *panicOnSecondWriteHeaderWriter) WriteHeader(s int) {
+	if w.headerWritten {
+		panic("header already written")
+	}
+	w.headerWritten = true
+	w.ResponseWriter.WriteHeader(s)
+}
+
+func TestGzipHandlerDoubleWriteHeader(t *testing.T) {
+	handler := GzipHandler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.Header().Set("Content-Length", "15000")
+		// Specifically write the header here
+		w.WriteHeader(304)
+		// Ensure that after a Write the header isn't triggered again on close
+		w.Write(nil)
+	}))
+	wrapper := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w = &panicOnSecondWriteHeaderWriter{
+			ResponseWriter: w,
+		}
+		handler.ServeHTTP(w, r)
+	})
+
+	rec := httptest.NewRecorder()
+	// TODO: in Go1.7 httptest.NewRequest was introduced this should be used
+	// once 1.6 is not longer supported.
+	req := &http.Request{
+		Method:     "GET",
+		URL:        &url.URL{Path: "/"},
+		Proto:      "HTTP/1.1",
+		ProtoMinor: 1,
+		RemoteAddr: "192.0.2.1:1234",
+		Header:     make(http.Header),
+	}
+	req.Header.Set("Accept-Encoding", "gzip")
+	wrapper.ServeHTTP(rec, req)
+	body, err := ioutil.ReadAll(rec.Body)
+	if err != nil {
+		t.Fatalf("Unexpected error reading response body: %v", err)
+	}
+	assert.Empty(t, body)
+	header := rec.Header()
+	assert.Equal(t, "gzip", header.Get("Content-Encoding"))
+	assert.Equal(t, "Accept-Encoding", header.Get("Vary"))
+	assert.Equal(t, 304, rec.Code)
+}
+
 func TestStatusCodes(t *testing.T) {
 	handler := GzipHandler(http.NotFoundHandler())
 	r := httptest.NewRequest("GET", "/", nil)