Check Content-Length against minSize and passthrough writes if no gzip

gzip.go

@@ -29,6 +29,7 @@ const (
 	// The examples seem to indicate that it is.
 	DefaultQValue = 1.0
 
+	// DefaultMinSize is the default minimum size until we enable gzip compression.
 	// 1500 bytes is the MTU size for the internet since that is the largest size allowed at the network layer.
 	// If you take a file that is 1300 bytes and compress it to 800 bytes, it’s still transmitted in that same 1500 byte packet regardless, so you’ve gained nothing.
 	// That being the case, you should restrict the gzip compression to files with a size greater than a single packet, 1400 bytes (1.4KB) is a safe value.
@@ -82,6 +83,7 @@ type GzipResponseWriter struct {
 
 	minSize int    // Specifed the minimum response size to gzip. If the response length is bigger than this value, it is compressed.
 	buf     []byte // Holds the first part of the write before reaching the minSize or the end of the write.
+	ignore  bool   // If true, then we immediately passthru writes to the underlying ResponseWriter.
 
 	contentTypes []parsedContentType // Only compress if the response is one of these content-types. All are accepted if empty.
 }
@@ -96,38 +98,56 @@ func (w GzipResponseWriterWithCloseNotify) CloseNotify() <-chan bool {
 
 // Write appends data to the gzip writer.
 func (w *GzipResponseWriter) Write(b []byte) (int, error) {
-	// If content type is not set.
-	if _, ok := w.Header()[contentType]; !ok {
-		// It infer it from the uncompressed body.
-		w.Header().Set(contentType, http.DetectContentType(b))
-	}
-
 	// GZIP responseWriter is initialized. Use the GZIP responseWriter.
 	if w.gw != nil {
-		n, err := w.gw.Write(b)
-		return n, err
+		return w.gw.Write(b)
+	}
+
+	// If we have already decided not to use GZIP, immediately passthrough.
+	if w.ignore {
+		return w.ResponseWriter.Write(b)
 	}
 
 	// Save the write into a buffer for later use in GZIP responseWriter (if content is long enough) or at close with regular responseWriter.
 	// On the first write, w.buf changes from nil to a valid slice
 	w.buf = append(w.buf, b...)
 
-	// If the global writes are bigger than the minSize and we're about to write
-	// a response containing a content type we want to handle, enable
-	// compression.
-	if len(w.buf) >= w.minSize && handleContentType(w.contentTypes, w) && w.Header().Get(contentEncoding) == "" {
-		err := w.startGzip()
-		if err != nil {
-			return 0, err
-		}
+	// If they provided a Content-Length, store that to check against minSize.
+	var cl int
+	if cls := w.Header().Get(contentLength); cls != "" {
+		cl, _ = strconv.Atoi(cls)
 	}
 
+	// Only continue if they didn't already choose an encoding.
+	if w.Header().Get(contentEncoding) == "" {
+		// If the current buffer is less than minSize and a Content-Length isn't set, then wait until we have more data.
+		if len(w.buf) < w.minSize && cl == 0 {
+			return len(b), nil
+		}
+		// If the Content-Length is larger than minSize or the current buffer is larger than minSize, then continue.
+		if cl >= w.minSize || len(w.buf) >= w.minSize {
+			// If a Content-Type wasn't specified, infer it from the current buffer.
+			if _, ok := w.Header()[contentType]; !ok {
+				w.Header().Set(contentType, http.DetectContentType(w.buf))
+			}
+			// If the Content-Type is acceptable to GZIP, initialize the GZIP writer.
+			if handleContentType(w.contentTypes, w) {
+				if err := w.startGzip(); err != nil {
+					return 0, err
+				}
+				return len(b), nil
+			}
+		}
+	}
+	// If we got here, we should not GZIP this response.
+	if err := w.startPlain(); err != nil {
+		return 0, err
+	}
 	return len(b), nil
 }
 
-// startGzip initialize any GZIP specific informations.
+// startGzip initializes a GZIP writer and writes the buffer.
 func (w *GzipResponseWriter) startGzip() error {
-
 	// Set the GZIP header.
 	w.Header().Set(contentEncoding, "gzip")
 
@@ -141,20 +161,43 @@ func (w *GzipResponseWriter) startGzip() error {
 		w.ResponseWriter.WriteHeader(w.code)
 	}
 
-	// Initialize the GZIP response.
-	w.init()
+	// Initialize and flush the buffer into the gzip response if there are any bytes.
+	// If there aren't any, we shouldn't initialize it yet because on Close it will
+	// write the gzip header even if nothing was ever written.
+	if len(w.buf) > 0 {
+		// Initialize the GZIP response.
+		w.init()
+		n, err := w.gw.Write(w.buf)
 
-	// Flush the buffer into the gzip response.
-	n, err := w.gw.Write(w.buf)
+		// This should never happen (per io.Writer docs), but if the write didn't
+		// accept the entire buffer but returned no specific error, we have no clue
+		// what's going on, so abort just to be safe.
+		if err == nil && n < len(w.buf) {
+			err = io.ErrShortWrite
+		}
+		return err
+	}
+	return nil
+}
 
+// startPlain writes to sent bytes and buffer the underlying ResponseWriter without gzip.
+func (w *GzipResponseWriter) startPlain() error {
+	if w.code != 0 {
+		w.ResponseWriter.WriteHeader(w.code)
+	}
+	w.ignore = true
+	// If Write was never called then don't call Write on the underlying ResponseWriter.
+	if w.buf == nil {
+		return nil
+	}
+	n, err := w.ResponseWriter.Write(w.buf)
+	w.buf = nil
 	// This should never happen (per io.Writer docs), but if the write didn't
 	// accept the entire buffer but returned no specific error, we have no clue
 	// what's going on, so abort just to be safe.
 	if err == nil && n < len(w.buf) {
-		return io.ErrShortWrite
+		err = io.ErrShortWrite
 	}
-
-	w.buf = nil
 	return err
 }
 
@@ -177,19 +220,18 @@ func (w *GzipResponseWriter) init() {
 
 // Close will close the gzip.Writer and will put it back in the gzipWriterPool.
 func (w *GzipResponseWriter) Close() error {
+	if w.ignore {
+		return nil
+	}
+
 	if w.gw == nil {
-		// Gzip not trigged yet, write out regular response.
-		if w.code != 0 {
-			w.ResponseWriter.WriteHeader(w.code)
-		}
-		if w.buf != nil {
-			_, writeErr := w.ResponseWriter.Write(w.buf)
-			// Returns the error if any at write.
-			if writeErr != nil {
-				return fmt.Errorf("gziphandler: write to regular responseWriter at close gets error: %q", writeErr.Error())
-			}
+		// GZIP not triggered yet, write out regular response.
+		err := w.startPlain()
+		// Returns the error if any at write.
+		if err != nil {
+			err = fmt.Errorf("gziphandler: write to regular responseWriter at close gets error: %q", err.Error())
 		}
-		return nil
+		return err
 	}
 
 	err := w.gw.Close()

gzip_test.go

@@ -25,9 +25,9 @@ func TestParseEncodings(t *testing.T) {
 	examples := map[string]codings{
 
 		// Examples from RFC 2616
-		"compress, gzip": {"compress": 1.0, "gzip": 1.0},
-		"":               {},
-		"*":              {"*": 1.0},
+		"compress, gzip":                     {"compress": 1.0, "gzip": 1.0},
+		"":                                   {},
+		"*":                                  {"*": 1.0},
 		"compress;q=0.5, gzip;q=1.0":         {"compress": 0.5, "gzip": 1.0},
 		"gzip;q=1.0, identity; q=0.5, *;q=0": {"gzip": 1.0, "identity": 0.5, "*": 0.0},
 
@@ -158,18 +158,23 @@ func TestGzipHandlerNoBody(t *testing.T) {
 	tests := []struct {
 		statusCode      int
 		contentEncoding string
-		bodyLen         int
+		emptyBody       bool
+		body            []byte
 	}{
 		// Body must be empty.
-		{http.StatusNoContent, "", 0},
-		{http.StatusNotModified, "", 0},
+		{http.StatusNoContent, "", true, nil},
+		{http.StatusNotModified, "", true, nil},
 		// Body is going to get gzip'd no matter what.
-		{http.StatusOK, "", 0},
+		{http.StatusOK, "", true, []byte{}},
+		{http.StatusOK, "gzip", false, []byte(testBody)},
 	}
 
 	for num, test := range tests {
 		handler := GzipHandler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 			w.WriteHeader(test.statusCode)
+			if test.body != nil {
+				w.Write(test.body)
+			}
 		}))
 
 		rec := httptest.NewRecorder()
@@ -194,16 +199,31 @@ func TestGzipHandlerNoBody(t *testing.T) {
 		header := rec.Header()
 		assert.Equal(t, test.contentEncoding, header.Get("Content-Encoding"), fmt.Sprintf("for test iteration %d", num))
 		assert.Equal(t, "Accept-Encoding", header.Get("Vary"), fmt.Sprintf("for test iteration %d", num))
-		assert.Equal(t, test.bodyLen, len(body), fmt.Sprintf("for test iteration %d", num))
+		if test.emptyBody {
+			assert.Empty(t, body, fmt.Sprintf("for test iteration %d", num))
+		} else {
+			assert.NotEmpty(t, body, fmt.Sprintf("for test iteration %d", num))
+			assert.NotEqual(t, test.body, body, fmt.Sprintf("for test iteration %d", num))
+		}
 	}
 }
 
 func TestGzipHandlerContentLength(t *testing.T) {
-	b := []byte(testBody)
-	handler := GzipHandler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-		w.Header().Set("Content-Length", strconv.Itoa(len(b)))
-		w.Write(b)
-	}))
+	testBodyBytes := []byte(testBody)
+	tests := []struct {
+		bodyLen   int
+		bodies    [][]byte
+		emptyBody bool
+	}{
+		{len(testBody), [][]byte{testBodyBytes}, false},
+		// each of these writes is less than the DefaultMinSize
+		{len(testBody), [][]byte{testBodyBytes[:200], testBodyBytes[200:]}, false},
+		// without a defined Content-Length it should still gzip
+		{0, [][]byte{testBodyBytes[:200], testBodyBytes[200:]}, false},
+		// simulate a HEAD request with an empty write (to populate headers)
+		{len(testBody), [][]byte{nil}, true},
+	}
+
 	// httptest.NewRecorder doesn't give you access to the Content-Length
 	// header so instead, we create a server on a random port and make
 	// a request to that instead
@@ -213,35 +233,50 @@ func TestGzipHandlerContentLength(t *testing.T) {
 	}
 	defer ln.Close()
 	srv := &http.Server{
-		Handler: handler,
+		Handler: nil,
 	}
 	go srv.Serve(ln)
 
-	req := &http.Request{
-		Method: "GET",
-		URL:    &url.URL{Path: "/", Scheme: "http", Host: ln.Addr().String()},
-		Header: make(http.Header),
-		Close:  true,
-	}
-	req.Header.Set("Accept-Encoding", "gzip")
-	res, err := http.DefaultClient.Do(req)
-	if err != nil {
-		t.Fatalf("Unexpected error making http request: %v", err)
-	}
-	defer res.Body.Close()
+	for num, test := range tests {
+		srv.Handler = GzipHandler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			if test.bodyLen > 0 {
+				w.Header().Set("Content-Length", strconv.Itoa(test.bodyLen))
+			}
+			for _, b := range test.bodies {
+				w.Write(b)
+			}
+		}))
+		req := &http.Request{
+			Method: "GET",
+			URL:    &url.URL{Path: "/", Scheme: "http", Host: ln.Addr().String()},
+			Header: make(http.Header),
+			Close:  true,
+		}
+		req.Header.Set("Accept-Encoding", "gzip")
+		res, err := http.DefaultClient.Do(req)
+		if err != nil {
+			t.Fatalf("Unexpected error making http request in test iteration %d: %v", num, err)
+		}
+		defer res.Body.Close()
 
-	body, err := ioutil.ReadAll(res.Body)
-	if err != nil {
-		t.Fatalf("Unexpected error reading response body: %v", err)
-	}
+		body, err := ioutil.ReadAll(res.Body)
+		if err != nil {
+			t.Fatalf("Unexpected error reading response body in test iteration %d: %v", num, err)
+		}
 
-	l, err := strconv.Atoi(res.Header.Get("Content-Length"))
-	if err != nil {
-		t.Fatalf("Unexpected error parsing Content-Length: %v", err)
+		l, err := strconv.Atoi(res.Header.Get("Content-Length"))
+		if err != nil {
+			t.Fatalf("Unexpected error parsing Content-Length in test iteration %d: %v", num, err)
+		}
+		if test.emptyBody {
+			assert.Empty(t, body, fmt.Sprintf("for test iteration %d", num))
+			assert.Equal(t, 0, l, fmt.Sprintf("for test iteration %d", num))
+		} else {
+			assert.Len(t, body, l, fmt.Sprintf("for test iteration %d", num))
+		}
+		assert.Equal(t, "gzip", res.Header.Get("Content-Encoding"), fmt.Sprintf("for test iteration %d", num))
+		assert.NotEqual(t, test.bodyLen, l, fmt.Sprintf("for test iteration %d", num))
 	}
-	assert.Len(t, body, l)
-	assert.Equal(t, "gzip", res.Header.Get("Content-Encoding"))
-	assert.NotEqual(t, b, body)
 }
 
 func TestGzipHandlerMinSizeMustBePositive(t *testing.T) {