Improve gzip middleware (goadesign#1368)

* Improve gzip middleware

Change gzip middleware to only be employed if:

* Response is above specified size (Default: 256 bytes)
* Status code is one of. (Default: OK, Created, Accepted)
* Content type is one of. (Default: Common types+ all `application/vnd.`)

Options added are `AddContentTypes`, "OnlyContentTypes", "AddStatusCodes", "OnlyStatusCodes" and "MinSize".

Expanded tests to cover these.

* Export option type for linter.

* Comment wording.

Author: klauspost

middleware/gzip/middleware.go

@@ -1,14 +1,14 @@
 package gzip
 
 import (
+	"bytes"
 	"compress/gzip"
+	"context"
 	"io/ioutil"
 	"net/http"
 	"strings"
 	"sync"
 
-	"context"
-
 	"github.com/goadesign/goa"
 )
 
@@ -28,28 +28,226 @@ const (
 // capabilities.
 type gzipResponseWriter struct {
 	http.ResponseWriter
-	gzw *gzip.Writer
+	gzw            *gzip.Writer
+	buf            *bytes.Buffer
+	pool           *sync.Pool
+	statusCode     int
+	shouldCompress *bool
+	o              options
 }
 
 // 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) {
+func (grw *gzipResponseWriter) Write(b []byte) (int, error) {
 	if len(grw.Header().Get(headerContentType)) == 0 {
 		grw.Header().Set(headerContentType, http.DetectContentType(b))
 	}
-	return grw.gzw.Write(b)
+
+	// If we already decided to gzip, do that.
+	if grw.gzw != nil {
+		return grw.gzw.Write(b)
+	}
+
+	// If we have already decided not to gzip, do that.
+	if grw.shouldCompress != nil && !*grw.shouldCompress {
+		return grw.ResponseWriter.Write(b)
+	}
+
+	// Detect types, check status code.
+	if grw.shouldCompress == nil {
+		s := grw.o.shouldCompress(grw.Header().Get(headerContentType), grw.statusCode)
+		grw.shouldCompress = &s
+		if !s {
+			return grw.ResponseWriter.Write(b)
+		}
+	}
+
+	// Check if length is above minimum,
+	// if not save to buffer.
+	size := len(b) + grw.buf.Len()
+	if size < grw.o.minSize {
+		return grw.buf.Write(b)
+	}
+
+	// Reset our gzip writer to use the http.ResponseWriter
+	// Retrieve gzip writer from the pool. Reset it to use the ResponseWriter.
+	// This allows us to re-use an already allocated buffer rather than
+	// allocating a new buffer for every request.
+	grw.Header().Set(headerContentEncoding, encodingGzip)
+	grw.Header().Set(headerVary, headerAcceptEncoding)
+
+	gz := grw.pool.Get().(*gzip.Writer)
+	gz.Reset(grw.ResponseWriter)
+	grw.gzw = gz
+
+	// Write buffer
+	if grw.buf.Len() > 0 {
+		_, err := gz.Write(grw.buf.Bytes())
+		if err != nil {
+			return 0, err
+		}
+		grw.buf.Reset()
+	}
+	return gz.Write(b)
+}
+
+func (grw *gzipResponseWriter) WriteHeader(n int) {
+	grw.statusCode = n
+	grw.ResponseWriter.WriteHeader(n)
+}
+
+type (
+	// Option allows to override default parameters.
+	Option func(*options) error
+
+	// options contains final options
+	options struct {
+		minSize      int
+		contentTypes []string
+		statusCodes  map[int]struct{}
+	}
+)
+
+var defaultOptions = options{
+	minSize:      256,
+	contentTypes: defaultContentTypes,
+}
+
+func init() {
+	defaultOptions.statusCodes = make(map[int]struct{}, len(defaultStatusCodes))
+	for _, v := range defaultStatusCodes {
+		defaultOptions.statusCodes[v] = struct{}{}
+	}
+}
+
+// defaultContentTypes is the default list of content types for which
+// a Handler considers gzip compression. This list originates from the
+// file compression.conf within the Apache configuration found at
+// https://html5boilerplate.com/
+var defaultContentTypes = []string{
+	"application/atom+xml",
+	"application/font-sfnt",
+	"application/javascript",
+	"application/json",
+	"application/ld+json",
+	"application/manifest+json",
+	"application/rdf+xml",
+	"application/rss+xml",
+	"application/schema+json",
+	"application/vnd.", // All custom vendor types
+	"application/x-font-ttf",
+	"application/x-javascript",
+	"application/x-web-app-manifest+json",
+	"application/xhtml+xml",
+	"application/xml",
+	"font/eot",
+	"font/opentype",
+	"image/bmp",
+	"image/svg+xml",
+	"image/vnd.microsoft.icon",
+	"image/x-icon",
+	"text/cache-manifest",
+	"text/css",
+	"text/html",
+	"text/javascript",
+	"text/plain",
+	"text/vcard",
+	"text/vnd.rim.location.xloc",
+	"text/vtt",
+	"text/x-component",
+	"text/x-cross-domain-policy",
+	"text/xml",
+}
+
+// defaultStatusCodes are the status codes that will be compressed.
+var defaultStatusCodes = []int{
+	http.StatusOK,
+	http.StatusCreated,
+	http.StatusAccepted,
+}
+
+// AddContentTypes allows to specify specific content types to encode.
+// Adds to previous content types.
+func AddContentTypes(types ...string) Option {
+	return func(c *options) error {
+		dst := make([]string, len(c.contentTypes)+len(types))
+		copy(dst, c.contentTypes)
+		copy(dst[len(c.contentTypes):], types)
+		c.contentTypes = dst
+		return nil
+	}
 }
 
-// handler struct contains the ServeHTTP method
-type handler struct {
-	pool sync.Pool
+// OnlyContentTypes allows to specify specific content types to encode.
+// Overrides previous content types.
+// no types = ignore content types (always compress).
+func OnlyContentTypes(types ...string) Option {
+	return func(c *options) error {
+		if len(types) == 0 {
+			c.contentTypes = nil
+			return nil
+		}
+		c.contentTypes = types
+		return nil
+	}
+}
+
+// AddStatusCodes allows to specify specific content types to encode.
+// All content types that has the supplied prefixes are compressed.
+func AddStatusCodes(codes ...int) Option {
+	return func(c *options) error {
+		dst := make(map[int]struct{}, len(c.statusCodes)+len(codes))
+		for code := range c.statusCodes {
+			dst[code] = struct{}{}
+		}
+		for _, code := range codes {
+			c.statusCodes[code] = struct{}{}
+		}
+		return nil
+	}
+}
+
+// OnlyStatusCodes allows to specify specific content types to encode.
+// All content types that has the supplied prefixes are compressed.
+// No codes = ignore content types (always compress).
+func OnlyStatusCodes(codes ...int) Option {
+	return func(c *options) error {
+		if len(codes) == 0 {
+			c.statusCodes = nil
+			return nil
+		}
+		c.statusCodes = make(map[int]struct{}, len(codes))
+		for _, code := range codes {
+			c.statusCodes[code] = struct{}{}
+		}
+		return nil
+	}
+}
+
+// MinSize will set a minimum size for compression.
+func MinSize(n int) Option {
+	return func(c *options) error {
+		if n <= 0 {
+			c.minSize = 0
+			return nil
+		}
+		c.minSize = n
+		return nil
+	}
 }
 
 // Middleware encodes the response using Gzip encoding and sets all the appropriate
 // headers. If the Content-Type is not set, it will be set by calling
 // http.DetectContentType on the data being written.
-func Middleware(level int) goa.Middleware {
+func Middleware(level int, o ...Option) goa.Middleware {
+	opts := defaultOptions
+	for _, opt := range o {
+		err := opt(&opts)
+		if err != nil {
+			panic(err)
+		}
+	}
 	gzipPool := sync.Pool{
 		New: func() interface{} {
 			gz, err := gzip.NewWriterLevel(ioutil.Discard, level)
@@ -71,23 +269,17 @@ func Middleware(level int) goa.Middleware {
 
 			// Set the appropriate gzip headers.
 			resp := goa.ContextResponse(ctx)
-			resp.Header().Set(headerContentEncoding, encodingGzip)
-			resp.Header().Set(headerVary, headerAcceptEncoding)
-
-			// Retrieve gzip writer from the pool. Reset it to use the ResponseWriter.
-			// This allows us to re-use an already allocated buffer rather than
-			// allocating a new buffer for every request.
-			gz := gzipPool.Get().(*gzip.Writer)
 
 			// Get the original http.ResponseWriter
 			w := resp.SwitchWriter(nil)
-			// Reset our gzip writer to use the http.ResponseWriter
-			gz.Reset(w)
 
 			// Wrap the original http.ResponseWriter with our gzipResponseWriter
-			grw := gzipResponseWriter{
+			grw := &gzipResponseWriter{
 				ResponseWriter: w,
-				gzw:            gz,
+				pool:           &gzipPool,
+				buf:            bytes.NewBuffer(nil),
+				statusCode:     http.StatusOK,
+				o:              opts,
 			}
 
 			// Set the new http.ResponseWriter
@@ -102,9 +294,49 @@ func Middleware(level int) goa.Middleware {
 
 			// Delete the content length after we know we have been written to.
 			grw.Header().Del(headerContentLength)
-			gz.Close()
-			gzipPool.Put(gz)
+			if grw.buf.Len() > 0 {
+				_, err = grw.ResponseWriter.Write(grw.buf.Bytes())
+				if err != nil {
+					return err
+				}
+			}
+
+			// Flush+recycle gzip writer
+			if grw.gzw != nil {
+				err := grw.gzw.Close()
+				if err != nil {
+					return err
+				}
+				gzipPool.Put(grw.gzw)
+			}
 			return
 		}
 	}
 }
+
+// returns true if we've been configured to compress the specific content type.
+func (o options) shouldCompress(contentType string, statusCode int) bool {
+	// If contentTypes is nil we handle all content types.
+	if len(o.contentTypes) > 0 {
+		ct := strings.ToLower(contentType)
+		ct = strings.Split(ct, ";")[0]
+		found := false
+		for _, v := range o.contentTypes {
+			if strings.HasPrefix(ct, v) {
+				found = true
+				break
+			}
+		}
+		if !found {
+			return false
+		}
+	}
+	if len(o.statusCodes) > 0 {
+		_, ok := o.statusCodes[statusCode]
+		if !ok {
+			return false
+		}
+	}
+
+	return true
+}