gzhttp: Support Flush always (#386)

* gzhttp: Support Flush always

Fixes nytimes/gziphandler#87

gzhttp/README.md

@@ -194,6 +194,15 @@ When replacing, this can be used to find a replacement.
 By default, some mime types will now be excluded.
 To re-enable compression of all types, use the `ContentTypeFilter(gzhttp.CompressAllContentTypeFilter)` option.
 
+### Flushing data
+
+The wrapper supports the [http.Flusher](https://golang.org/pkg/net/http/#Flusher) interface.
+
+The only caveat is that the writer may not yet have received enough bytes to determine if `MinSize`
+has been reached. In this case it will assume that the minimum size has been reached.
+
+If nothing has been written to the response writer, nothing will be flushed.
+
 ## License
 
 [Apache 2.0](LICENSE)

gzhttp/gzip.go

@@ -118,10 +118,10 @@ func (w *GzipResponseWriter) Write(b []byte) (int, error) {
 
 			// Handles the intended case of setting a nil Content-Type (as for http/server or http/fs)
 			// Set the header only if the key does not exist
-			_, haveType := w.Header()["Content-Type"]
-			if !haveType {
+			if _, ok := w.Header()[contentType]; !ok {
 				w.Header().Set(contentType, ct)
 			}
+
 			// If the Content-Type is acceptable to GZIP, initialize the GZIP writer.
 			if w.contentTypeFilter(ct) {
 				if err := w.startGzip(); err != nil {
@@ -253,13 +253,42 @@ func (w *GzipResponseWriter) Close() error {
 // Flush flushes the underlying *gzip.Writer and then the underlying
 // http.ResponseWriter if it is an http.Flusher. This makes GzipResponseWriter
 // an http.Flusher.
+// If not enough bytes has been written to determine if we have reached minimum size,
+// this will be ignored.
+// If nothing has been written yet, nothing will be flushed.
 func (w *GzipResponseWriter) Flush() {
 	if w.gw == nil && !w.ignore {
-		// Only flush once startGzip or startPlain has been called.
-		//
-		// Flush is thus a no-op until we're certain whether a plain
-		// or gzipped response will be served.
-		return
+		if len(w.buf) == 0 {
+			// Nothing written yet.
+			return
+		}
+		var (
+			cl, _ = atoi(w.Header().Get(contentLength))
+			ct    = w.Header().Get(contentType)
+			ce    = w.Header().Get(contentEncoding)
+			cr    = w.Header().Get(contentRange)
+		)
+
+		if ct == "" {
+			ct = http.DetectContentType(w.buf)
+
+			// Handles the intended case of setting a nil Content-Type (as for http/server or http/fs)
+			// Set the header only if the key does not exist
+			if _, ok := w.Header()[contentType]; !ok {
+				w.Header().Set(contentType, ct)
+			}
+		}
+		if cl == 0 {
+			// Assume minSize.
+			cl = w.minSize
+		}
+
+		// See if we should compress...
+		if ce == "" && cr == "" && cl >= w.minSize && w.contentTypeFilter(ct) {
+			w.startGzip()
+		} else {
+			w.startPlain()
+		}
 	}
 
 	if w.gw != nil {

gzhttp/gzip_test.go

@@ -677,6 +677,59 @@ func TestContentTypes(t *testing.T) {
 	}
 }
 
+func TestFlush(t *testing.T) {
+	for _, tt := range contentTypeTests {
+		t.Run(tt.name, func(t *testing.T) {
+			handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+				w.WriteHeader(http.StatusOK)
+				w.Header().Set("Content-Type", tt.contentType)
+				tb := testBody
+				for len(tb) > 0 {
+					// Write 100 bytes per run
+					// Detection should not be affected (we send 100 bytes)
+					toWrite := 100
+					if toWrite > len(tb) {
+						toWrite = len(tb)
+					}
+					_, err := w.Write(tb[:toWrite])
+					if err != nil {
+						t.Fatal(err)
+					}
+					// Flush between each write
+					w.(http.Flusher).Flush()
+					tb = tb[toWrite:]
+				}
+			})
+
+			wrapper, err := NewWrapper(ContentTypes(tt.acceptedContentTypes))
+			assertNil(t, err)
+
+			req, _ := http.NewRequest("GET", "/whatever", nil)
+			req.Header.Set("Accept-Encoding", "gzip")
+			// This doesn't allow checking flushes, but we validate if content is correct.
+			resp := httptest.NewRecorder()
+			wrapper(handler).ServeHTTP(resp, req)
+			res := resp.Result()
+
+			assertEqual(t, 200, res.StatusCode)
+			if tt.expectedGzip {
+				assertEqual(t, "gzip", res.Header.Get("Content-Encoding"))
+				zr, err := gzip.NewReader(resp.Body)
+				assertNil(t, err)
+				got, err := ioutil.ReadAll(zr)
+				assertNil(t, err)
+				assertEqual(t, testBody, got)
+
+			} else {
+				assertNotEqual(t, "gzip", res.Header.Get("Content-Encoding"))
+				got, err := ioutil.ReadAll(resp.Body)
+				assertNil(t, err)
+				assertEqual(t, testBody, got)
+			}
+		})
+	}
+}
+
 var contentTypeTest2 = []struct {
 	name         string
 	contentType  string
@@ -856,6 +909,67 @@ func TestContentTypeDetect(t *testing.T) {
 				assertNotEqual(t, "gzip", res.Header.Get("Content-Encoding"))
 			}
 		})
+		t.Run(tt.desc+"empty", func(t *testing.T) {
+			handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+				w.Header().Set("Content-Type", "")
+				w.WriteHeader(http.StatusOK)
+				for i := range tt.data {
+					// Do one byte writes...
+					w.Write([]byte{tt.data[i]})
+				}
+				w.Write(testBody)
+			})
+
+			wrapper, err := NewWrapper()
+			assertNil(t, err)
+
+			req, _ := http.NewRequest("GET", "/whatever", nil)
+			req.Header.Set("Accept-Encoding", "gzip")
+			resp := httptest.NewRecorder()
+			wrapper(handler).ServeHTTP(resp, req)
+			res := resp.Result()
+
+			assertEqual(t, 200, res.StatusCode)
+			// Is Content-Type still empty?
+			assertEqual(t, "", res.Header.Get("Content-Type"))
+			shouldGZ := DefaultContentTypeFilter(tt.contentType)
+			if shouldGZ {
+				assertEqual(t, "gzip", res.Header.Get("Content-Encoding"))
+			} else {
+				assertNotEqual(t, "gzip", res.Header.Get("Content-Encoding"))
+			}
+		})
+		t.Run(tt.desc+"flush", func(t *testing.T) {
+			handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+				w.Header().Set("Content-Type", "")
+				w.WriteHeader(http.StatusOK)
+				for i := range tt.data {
+					// Do one byte writes...
+					w.Write([]byte{tt.data[i]})
+				}
+				w.(http.Flusher).Flush()
+				w.Write(testBody)
+			})
+
+			wrapper, err := NewWrapper()
+			assertNil(t, err)
+
+			req, _ := http.NewRequest("GET", "/whatever", nil)
+			req.Header.Set("Accept-Encoding", "gzip")
+			resp := httptest.NewRecorder()
+			wrapper(handler).ServeHTTP(resp, req)
+			res := resp.Result()
+
+			assertEqual(t, 200, res.StatusCode)
+			// Is Content-Type still empty?
+			assertEqual(t, "", res.Header.Get("Content-Type"))
+			shouldGZ := DefaultContentTypeFilter(tt.contentType)
+			if shouldGZ {
+				assertEqual(t, "gzip", res.Header.Get("Content-Encoding"))
+			} else {
+				assertNotEqual(t, "gzip", res.Header.Get("Content-Encoding"))
+			}
+		})
 	}
 }