From 4324668c5125c02964542a0b6d076b6bb28799a1 Mon Sep 17 00:00:00 2001
From: Tom Thorogood <me+github@tomthorogood.co.uk>
Date: Tue, 4 Apr 2017 10:13:59 +0930
Subject: [PATCH] Fix Content-Type detection w/ buffered response

---
 gzip.go | 41 ++++++++++++++++++++++++++++++-----------
 1 file changed, 30 insertions(+), 11 deletions(-)

diff --git a/gzip.go b/gzip.go
index 0345d47..83be372 100644
--- a/gzip.go
+++ b/gzip.go
@@ -35,14 +35,6 @@ type responseWriter struct {
 
 // Write appends data to the gzip writer.
 func (w *responseWriter) Write(b []byte) (int, error) {
-	h := w.Header()
-
-	// If content type is not set.
-	if _, ok := h["Content-Type"]; !ok {
-		// It infer it from the uncompressed body.
-		h["Content-Type"] = []string{http.DetectContentType(b)}
-	}
-
 	// GZIP responseWriter is initialized. Use the GZIP
 	// responseWriter.
 	if w.gw != nil {
@@ -58,11 +50,15 @@ func (w *responseWriter) Write(b []byte) (int, error) {
 		// responseWriter.
 		w.buf = append(w.buf, b...)
 		return len(b), nil
-	} else if err := w.startGzip(); err != nil {
+	}
+
+	w.inferContentType(b)
+
+	if err := w.startGzip(); err != nil {
 		return 0, err
-	} else {
-		return w.gw.Write(b)
 	}
+
+	return w.gw.Write(b)
 }
 
 // startGzip initialize any GZIP specific informations.
@@ -101,6 +97,29 @@ func (w *responseWriter) startGzip() error {
 	return err
 }
 
+func (w *responseWriter) inferContentType(b []byte) {
+	h := w.Header()
+
+	// If content type is not set.
+	if _, ok := h["Content-Type"]; ok {
+		return
+	}
+
+	if len(w.buf) != 0 {
+		const sniffLen = 512
+		if len(w.buf) >= sniffLen {
+			b = w.buf
+		} else if len(w.buf)+len(b) > sniffLen {
+			b = append(w.buf, b[:sniffLen-len(w.buf)]...)
+		} else {
+			b = append(w.buf, b...)
+		}
+	}
+
+	// It infer it from the uncompressed body.
+	h["Content-Type"] = []string{http.DetectContentType(b)}
+}
+
 // WriteHeader just saves the response code until close or
 // GZIP effective writes.
 func (w *responseWriter) WriteHeader(code int) {