123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399 |
- package middleware
-
- import (
- "bufio"
- "compress/flate"
- "compress/gzip"
- "errors"
- "fmt"
- "io"
- "io/ioutil"
- "net"
- "net/http"
- "strings"
- "sync"
- )
-
- var defaultCompressibleContentTypes = []string{
- "text/html",
- "text/css",
- "text/plain",
- "text/javascript",
- "application/javascript",
- "application/x-javascript",
- "application/json",
- "application/atom+xml",
- "application/rss+xml",
- "image/svg+xml",
- }
-
- // Compress is a middleware that compresses response
- // body of a given content types to a data format based
- // on Accept-Encoding request header. It uses a given
- // compression level.
- //
- // NOTE: make sure to set the Content-Type header on your response
- // otherwise this middleware will not compress the response body. For ex, in
- // your handler you should set w.Header().Set("Content-Type", http.DetectContentType(yourBody))
- // or set it manually.
- //
- // Passing a compression level of 5 is sensible value
- func Compress(level int, types ...string) func(next http.Handler) http.Handler {
- compressor := NewCompressor(level, types...)
- return compressor.Handler
- }
-
- // Compressor represents a set of encoding configurations.
- type Compressor struct {
- // The mapping of encoder names to encoder functions.
- encoders map[string]EncoderFunc
- // The mapping of pooled encoders to pools.
- pooledEncoders map[string]*sync.Pool
- // The set of content types allowed to be compressed.
- allowedTypes map[string]struct{}
- allowedWildcards map[string]struct{}
- // The list of encoders in order of decreasing precedence.
- encodingPrecedence []string
- level int // The compression level.
- }
-
- // NewCompressor creates a new Compressor that will handle encoding responses.
- //
- // The level should be one of the ones defined in the flate package.
- // The types are the content types that are allowed to be compressed.
- func NewCompressor(level int, types ...string) *Compressor {
- // If types are provided, set those as the allowed types. If none are
- // provided, use the default list.
- allowedTypes := make(map[string]struct{})
- allowedWildcards := make(map[string]struct{})
- if len(types) > 0 {
- for _, t := range types {
- if strings.Contains(strings.TrimSuffix(t, "/*"), "*") {
- panic(fmt.Sprintf("middleware/compress: Unsupported content-type wildcard pattern '%s'. Only '/*' supported", t))
- }
- if strings.HasSuffix(t, "/*") {
- allowedWildcards[strings.TrimSuffix(t, "/*")] = struct{}{}
- } else {
- allowedTypes[t] = struct{}{}
- }
- }
- } else {
- for _, t := range defaultCompressibleContentTypes {
- allowedTypes[t] = struct{}{}
- }
- }
-
- c := &Compressor{
- level: level,
- encoders: make(map[string]EncoderFunc),
- pooledEncoders: make(map[string]*sync.Pool),
- allowedTypes: allowedTypes,
- allowedWildcards: allowedWildcards,
- }
-
- // Set the default encoders. The precedence order uses the reverse
- // ordering that the encoders were added. This means adding new encoders
- // will move them to the front of the order.
- //
- // TODO:
- // lzma: Opera.
- // sdch: Chrome, Android. Gzip output + dictionary header.
- // br: Brotli, see https://github.com/go-chi/chi/pull/326
-
- // HTTP 1.1 "deflate" (RFC 2616) stands for DEFLATE data (RFC 1951)
- // wrapped with zlib (RFC 1950). The zlib wrapper uses Adler-32
- // checksum compared to CRC-32 used in "gzip" and thus is faster.
- //
- // But.. some old browsers (MSIE, Safari 5.1) incorrectly expect
- // raw DEFLATE data only, without the mentioned zlib wrapper.
- // Because of this major confusion, most modern browsers try it
- // both ways, first looking for zlib headers.
- // Quote by Mark Adler: http://stackoverflow.com/a/9186091/385548
- //
- // The list of browsers having problems is quite big, see:
- // http://zoompf.com/blog/2012/02/lose-the-wait-http-compression
- // https://web.archive.org/web/20120321182910/http://www.vervestudios.co/projects/compression-tests/results
- //
- // That's why we prefer gzip over deflate. It's just more reliable
- // and not significantly slower than gzip.
- c.SetEncoder("deflate", encoderDeflate)
-
- // TODO: Exception for old MSIE browsers that can't handle non-HTML?
- // https://zoompf.com/blog/2012/02/lose-the-wait-http-compression
- c.SetEncoder("gzip", encoderGzip)
-
- // NOTE: Not implemented, intentionally:
- // case "compress": // LZW. Deprecated.
- // case "bzip2": // Too slow on-the-fly.
- // case "zopfli": // Too slow on-the-fly.
- // case "xz": // Too slow on-the-fly.
- return c
- }
-
- // SetEncoder can be used to set the implementation of a compression algorithm.
- //
- // The encoding should be a standardised identifier. See:
- // https://developer.mozilla.org/en-US/docs/Web/HTTP/Headers/Accept-Encoding
- //
- // For example, add the Brotli algortithm:
- //
- // import brotli_enc "gopkg.in/kothar/brotli-go.v0/enc"
- //
- // compressor := middleware.NewCompressor(5, "text/html")
- // compressor.SetEncoder("br", func(w http.ResponseWriter, level int) io.Writer {
- // params := brotli_enc.NewBrotliParams()
- // params.SetQuality(level)
- // return brotli_enc.NewBrotliWriter(params, w)
- // })
- func (c *Compressor) SetEncoder(encoding string, fn EncoderFunc) {
- encoding = strings.ToLower(encoding)
- if encoding == "" {
- panic("the encoding can not be empty")
- }
- if fn == nil {
- panic("attempted to set a nil encoder function")
- }
-
- // If we are adding a new encoder that is already registered, we have to
- // clear that one out first.
- if _, ok := c.pooledEncoders[encoding]; ok {
- delete(c.pooledEncoders, encoding)
- }
- if _, ok := c.encoders[encoding]; ok {
- delete(c.encoders, encoding)
- }
-
- // If the encoder supports Resetting (IoReseterWriter), then it can be pooled.
- encoder := fn(ioutil.Discard, c.level)
- if encoder != nil {
- if _, ok := encoder.(ioResetterWriter); ok {
- pool := &sync.Pool{
- New: func() interface{} {
- return fn(ioutil.Discard, c.level)
- },
- }
- c.pooledEncoders[encoding] = pool
- }
- }
- // If the encoder is not in the pooledEncoders, add it to the normal encoders.
- if _, ok := c.pooledEncoders[encoding]; !ok {
- c.encoders[encoding] = fn
- }
-
- for i, v := range c.encodingPrecedence {
- if v == encoding {
- c.encodingPrecedence = append(c.encodingPrecedence[:i], c.encodingPrecedence[i+1:]...)
- }
- }
-
- c.encodingPrecedence = append([]string{encoding}, c.encodingPrecedence...)
- }
-
- // Handler returns a new middleware that will compress the response based on the
- // current Compressor.
- func (c *Compressor) Handler(next http.Handler) http.Handler {
- return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
- encoder, encoding, cleanup := c.selectEncoder(r.Header, w)
-
- cw := &compressResponseWriter{
- ResponseWriter: w,
- w: w,
- contentTypes: c.allowedTypes,
- contentWildcards: c.allowedWildcards,
- encoding: encoding,
- compressable: false, // determined in post-handler
- }
- if encoder != nil {
- cw.w = encoder
- }
- // Re-add the encoder to the pool if applicable.
- defer cleanup()
- defer cw.Close()
-
- next.ServeHTTP(cw, r)
- })
- }
-
- // selectEncoder returns the encoder, the name of the encoder, and a closer function.
- func (c *Compressor) selectEncoder(h http.Header, w io.Writer) (io.Writer, string, func()) {
- header := h.Get("Accept-Encoding")
-
- // Parse the names of all accepted algorithms from the header.
- accepted := strings.Split(strings.ToLower(header), ",")
-
- // Find supported encoder by accepted list by precedence
- for _, name := range c.encodingPrecedence {
- if matchAcceptEncoding(accepted, name) {
- if pool, ok := c.pooledEncoders[name]; ok {
- encoder := pool.Get().(ioResetterWriter)
- cleanup := func() {
- pool.Put(encoder)
- }
- encoder.Reset(w)
- return encoder, name, cleanup
-
- }
- if fn, ok := c.encoders[name]; ok {
- return fn(w, c.level), name, func() {}
- }
- }
-
- }
-
- // No encoder found to match the accepted encoding
- return nil, "", func() {}
- }
-
- func matchAcceptEncoding(accepted []string, encoding string) bool {
- for _, v := range accepted {
- if strings.Contains(v, encoding) {
- return true
- }
- }
- return false
- }
-
- // An EncoderFunc is a function that wraps the provided io.Writer with a
- // streaming compression algorithm and returns it.
- //
- // In case of failure, the function should return nil.
- type EncoderFunc func(w io.Writer, level int) io.Writer
-
- // Interface for types that allow resetting io.Writers.
- type ioResetterWriter interface {
- io.Writer
- Reset(w io.Writer)
- }
-
- type compressResponseWriter struct {
- http.ResponseWriter
-
- // The streaming encoder writer to be used if there is one. Otherwise,
- // this is just the normal writer.
- w io.Writer
- contentTypes map[string]struct{}
- contentWildcards map[string]struct{}
- encoding string
- wroteHeader bool
- compressable bool
- }
-
- func (cw *compressResponseWriter) isCompressable() bool {
- // Parse the first part of the Content-Type response header.
- contentType := cw.Header().Get("Content-Type")
- if idx := strings.Index(contentType, ";"); idx >= 0 {
- contentType = contentType[0:idx]
- }
-
- // Is the content type compressable?
- if _, ok := cw.contentTypes[contentType]; ok {
- return true
- }
- if idx := strings.Index(contentType, "/"); idx > 0 {
- contentType = contentType[0:idx]
- _, ok := cw.contentWildcards[contentType]
- return ok
- }
- return false
- }
-
- func (cw *compressResponseWriter) WriteHeader(code int) {
- if cw.wroteHeader {
- cw.ResponseWriter.WriteHeader(code) // Allow multiple calls to propagate.
- return
- }
- cw.wroteHeader = true
- defer cw.ResponseWriter.WriteHeader(code)
-
- // Already compressed data?
- if cw.Header().Get("Content-Encoding") != "" {
- return
- }
-
- if !cw.isCompressable() {
- cw.compressable = false
- return
- }
-
- if cw.encoding != "" {
- cw.compressable = true
- cw.Header().Set("Content-Encoding", cw.encoding)
- cw.Header().Set("Vary", "Accept-Encoding")
-
- // The content-length after compression is unknown
- cw.Header().Del("Content-Length")
- }
- }
-
- func (cw *compressResponseWriter) Write(p []byte) (int, error) {
- if !cw.wroteHeader {
- cw.WriteHeader(http.StatusOK)
- }
-
- return cw.writer().Write(p)
- }
-
- func (cw *compressResponseWriter) writer() io.Writer {
- if cw.compressable {
- return cw.w
- } else {
- return cw.ResponseWriter
- }
- }
-
- type compressFlusher interface {
- Flush() error
- }
-
- func (cw *compressResponseWriter) Flush() {
- if f, ok := cw.writer().(http.Flusher); ok {
- f.Flush()
- }
- // If the underlying writer has a compression flush signature,
- // call this Flush() method instead
- if f, ok := cw.writer().(compressFlusher); ok {
- f.Flush()
-
- // Also flush the underlying response writer
- if f, ok := cw.ResponseWriter.(http.Flusher); ok {
- f.Flush()
- }
- }
- }
-
- func (cw *compressResponseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
- if hj, ok := cw.writer().(http.Hijacker); ok {
- return hj.Hijack()
- }
- return nil, nil, errors.New("chi/middleware: http.Hijacker is unavailable on the writer")
- }
-
- func (cw *compressResponseWriter) Push(target string, opts *http.PushOptions) error {
- if ps, ok := cw.writer().(http.Pusher); ok {
- return ps.Push(target, opts)
- }
- return errors.New("chi/middleware: http.Pusher is unavailable on the writer")
- }
-
- func (cw *compressResponseWriter) Close() error {
- if c, ok := cw.writer().(io.WriteCloser); ok {
- return c.Close()
- }
- return errors.New("chi/middleware: io.WriteCloser is unavailable on the writer")
- }
-
- func encoderGzip(w io.Writer, level int) io.Writer {
- gw, err := gzip.NewWriterLevel(w, level)
- if err != nil {
- return nil
- }
- return gw
- }
-
- func encoderDeflate(w io.Writer, level int) io.Writer {
- dw, err := flate.NewWriter(w, level)
- if err != nil {
- return nil
- }
- return dw
- }
|