Set headers on write

gzip/gzip.go

@@ -28,34 +28,85 @@ const (
 	NoCompression      = gzip.NoCompression
 )
 
+type Options struct {
+	Level               int
+	ExcludeContentTypes []string
+}
+
 // gzipResponseWriter is the ResponseWriter that negroni.ResponseWriter is
 // wrapped in.
 type gzipResponseWriter struct {
-	w *gzip.Writer
+	w                   *gzip.Writer
+	excludeContentTypes []string
+	skipCompression     bool
+	wroteHeader         bool
 	negroni.ResponseWriter
 }
 
 // Write writes bytes to the gzip.Writer. It will also set the Content-Type
 // header using the net/http library content type detection if the Content-Type
 // header was not set yet.
-func (grw gzipResponseWriter) Write(b []byte) (int, error) {
-	if len(grw.Header().Get(headerContentType)) == 0 {
-		grw.Header().Set(headerContentType, http.DetectContentType(b))
+func (grw *gzipResponseWriter) Write(b []byte) (int, error) {
+	if !grw.wroteHeader {
+		// It is not too late to set Content-Type!
+		contentType := grw.Header().Get(headerContentType)
+		if len(contentType) == 0 {
+			contentType = http.DetectContentType(b)
+			grw.Header().Set(headerContentType, contentType)
+		}
+
+		for _, ct := range grw.excludeContentTypes {
+			if contentType == ct {
+				grw.skipCompression = true
+				break
+			}
+		}
+
+		if !grw.skipCompression {
+			grw.Header().Set(headerContentEncoding, encodingGzip)
+			grw.Header().Set(headerVary, headerAcceptEncoding)
+		}
+
+		grw.wroteHeader = true
+	}
+
+	if grw.skipCompression {
+		return grw.ResponseWriter.Write(b) // bypass
+	} else {
+		return grw.w.Write(b)
 	}
-	return grw.w.Write(b)
+}
+
+func (grw *gzipResponseWriter) WriteHeader(s int) {
+	contentType := grw.Header().Get(headerContentType)
+	for _, ct := range grw.excludeContentTypes {
+		if contentType == ct {
+			grw.skipCompression = true
+			break
+		}
+	}
+
+	if !grw.skipCompression {
+		grw.Header().Set(headerContentEncoding, encodingGzip)
+		grw.Header().Set(headerVary, headerAcceptEncoding)
+	}
+
+	grw.wroteHeader = true
+	grw.ResponseWriter.WriteHeader(s)
 }
 
 // handler struct contains the ServeHTTP method
 type handler struct {
 	pool sync.Pool
+	opt  *Options
 }
 
 // Gzip returns a handler which will handle the Gzip compression in ServeHTTP.
 // Valid values for level are identical to those in the compress/gzip package.
-func Gzip(level int) *handler {
-	h := &handler{}
+func Gzip(opt *Options) *handler {
+	h := &handler{opt: opt}
 	h.pool.New = func() interface{} {
-		gz, err := gzip.NewWriterLevel(ioutil.Discard, level)
+		gz, err := gzip.NewWriterLevel(ioutil.Discard, h.opt.Level)
 		if err != nil {
 			panic(err)
 		}
@@ -94,25 +145,24 @@ func (h *handler) ServeHTTP(w http.ResponseWriter, r *http.Request, next http.Ha
 	defer h.pool.Put(gz)
 	gz.Reset(w)
 
-	// Set the appropriate gzip headers.
-	headers := w.Header()
-	headers.Set(headerContentEncoding, encodingGzip)
-	headers.Set(headerVary, headerAcceptEncoding)
-
 	// Wrap the original http.ResponseWriter with negroni.ResponseWriter
 	// and create the gzipResponseWriter.
 	nrw := negroni.NewResponseWriter(w)
+
 	grw := gzipResponseWriter{
-		gz,
-		nrw,
+		w:                   gz,
+		excludeContentTypes: h.opt.ExcludeContentTypes,
+		ResponseWriter:      nrw,
 	}
 
 	// Call the next handler supplying the gzipResponseWriter instead of
 	// the original.
-	next(grw, r)
+	next(&grw, r)
 
 	// Delete the content length after we know we have been written to.
 	grw.Header().Del(headerContentLength)
 
-	gz.Close()
+	if !grw.skipCompression {
+		gz.Close()
+	}
 }