gziphandler.go

package gziphandler

import (
	"compress/gzip"
	"io"
	"net/http"
	"strings"
)

type gzipResponseWriter struct {
	io.Writer
	http.ResponseWriter

	statusCode     int
	contentTypeSet bool
}

func (w gzipResponseWriter) WriteHeader(statusCode int) {
	w.statusCode = statusCode
	w.ResponseWriter.WriteHeader(statusCode)
}

func (w gzipResponseWriter) Write(b []byte) (int, error) {
	if w.contentTypeSet == false && w.Header().Get("Content-Type") == "" {
		// If no content type, apply sniffing algorithm to un-gzipped body.
		w.Header().Set("Content-Type", http.DetectContentType(b))
		w.contentTypeSet = true
	}

	return w.Writer.Write(b)
}

func makeGzipHandler(fn http.HandlerFunc) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		if !strings.Contains(r.Header.Get("Accept-Encoding"), "gzip") {
			fn(w, r)
			return
		}

		h := w.Header()
		h["Content-Encoding"] = []string{"gzip"}
		h["Vary"] = []string{"Accept-Encoding"}

		// We can ignore the error as only the level argument can produce an error.
		gz, _ := gzip.NewWriterLevel(w, gzip.BestSpeed)
		gzr := gzipResponseWriter{Writer: gz, ResponseWriter: w}

		defer func() {
			// gz.Close will write a footer even if no data has been written.
			// 304 and 204 expect an empty body so don't close it.
			if gzr.statusCode != 304 && gzr.statusCode != 204 {
				gz.Close()
			}
		}()

		fn(gzr, r)
	}
}

func handler(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", "text/plain")
	w.Write([]byte("This is a test."))
}

func main() {
	http.ListenAndServe(":1113", makeGzipHandler(handler))
}

gziphandler_test.go

package gziphandler

import (
	"compress/gzip"
	"net/http"
	"net/http/httptest"
	"testing"
)

func test(w http.ResponseWriter, r *http.Request) {
	w.Write([]byte("test"))
}

func TestGZip(t *testing.T) {
	req, err := http.NewRequest("GET", "http://example.com/", nil)
	if err != nil {
		t.Error(err)
	}
	req.Header.Set("Accept-Encoding", "gzip, deflate")

	w := httptest.NewRecorder()
	makeGzipHandler(test)(w, req)

	if w.Header().Get("Content-Encoding") != "gzip" {
		t.Errorf("expected Content-Encoding: gzip got %d", w.Header().Get("Content-Encoding"))
	}

	r, err := gzip.NewReader(w.Body)
	if err != nil {
		t.Error(err)
	}

	b := make([]byte, 32)
	if n, err := r.Read(b); err != nil {
		t.Error(err)
	} else if n != 4 {
		t.Errorf(`expected len("test") == 4, got %d`, n)
	}

	if string(b[:4]) != "test" {
		t.Errorf(`expected "test" go "%s"`, string(b[:4]))
	}
}