compress.go

Documentation: github.com/gorilla/handlers

     1  // Copyright 2013 The Gorilla Authors. All rights reserved.
     2  // Use of this source code is governed by a BSD-style
     3  // license that can be found in the LICENSE file.
     4  
     5  package handlers
     6  
     7  import (
     8  	"compress/flate"
     9  	"compress/gzip"
    10  	"io"
    11  	"net/http"
    12  	"strings"
    13  
    14  	"github.com/felixge/httpsnoop"
    15  )
    16  
    17  const acceptEncoding string = "Accept-Encoding"
    18  
    19  type compressResponseWriter struct {
    20  	compressor io.Writer
    21  	w          http.ResponseWriter
    22  }
    23  
    24  func (cw *compressResponseWriter) WriteHeader(c int) {
    25  	cw.w.Header().Del("Content-Length")
    26  	cw.w.WriteHeader(c)
    27  }
    28  
    29  func (cw *compressResponseWriter) Write(b []byte) (int, error) {
    30  	h := cw.w.Header()
    31  	if h.Get("Content-Type") == "" {
    32  		h.Set("Content-Type", http.DetectContentType(b))
    33  	}
    34  	h.Del("Content-Length")
    35  
    36  	return cw.compressor.Write(b)
    37  }
    38  
    39  func (cw *compressResponseWriter) ReadFrom(r io.Reader) (int64, error) {
    40  	return io.Copy(cw.compressor, r)
    41  }
    42  
    43  type flusher interface {
    44  	Flush() error
    45  }
    46  
    47  func (w *compressResponseWriter) Flush() {
    48  	// Flush compressed data if compressor supports it.
    49  	if f, ok := w.compressor.(flusher); ok {
    50  		f.Flush()
    51  	}
    52  	// Flush HTTP response.
    53  	if f, ok := w.w.(http.Flusher); ok {
    54  		f.Flush()
    55  	}
    56  }
    57  
    58  // CompressHandler gzip compresses HTTP responses for clients that support it
    59  // via the 'Accept-Encoding' header.
    60  //
    61  // Compressing TLS traffic may leak the page contents to an attacker if the
    62  // page contains user input: http://security.stackexchange.com/a/102015/12208
    63  func CompressHandler(h http.Handler) http.Handler {
    64  	return CompressHandlerLevel(h, gzip.DefaultCompression)
    65  }
    66  
    67  // CompressHandlerLevel gzip compresses HTTP responses with specified compression level
    68  // for clients that support it via the 'Accept-Encoding' header.
    69  //
    70  // The compression level should be gzip.DefaultCompression, gzip.NoCompression,
    71  // or any integer value between gzip.BestSpeed and gzip.BestCompression inclusive.
    72  // gzip.DefaultCompression is used in case of invalid compression level.
    73  func CompressHandlerLevel(h http.Handler, level int) http.Handler {
    74  	if level < gzip.DefaultCompression || level > gzip.BestCompression {
    75  		level = gzip.DefaultCompression
    76  	}
    77  
    78  	const (
    79  		gzipEncoding  = "gzip"
    80  		flateEncoding = "deflate"
    81  	)
    82  
    83  	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    84  		// detect what encoding to use
    85  		var encoding string
    86  		for _, curEnc := range strings.Split(r.Header.Get(acceptEncoding), ",") {
    87  			curEnc = strings.TrimSpace(curEnc)
    88  			if curEnc == gzipEncoding || curEnc == flateEncoding {
    89  				encoding = curEnc
    90  				break
    91  			}
    92  		}
    93  
    94  		// always add Accept-Encoding to Vary to prevent intermediate caches corruption
    95  		w.Header().Add("Vary", acceptEncoding)
    96  
    97  		// if we weren't able to identify an encoding we're familiar with, pass on the
    98  		// request to the handler and return
    99  		if encoding == "" {
   100  			h.ServeHTTP(w, r)
   101  			return
   102  		}
   103  
   104  		if r.Header.Get("Upgrade") != "" {
   105  			h.ServeHTTP(w, r)
   106  			return
   107  		}
   108  
   109  		// wrap the ResponseWriter with the writer for the chosen encoding
   110  		var encWriter io.WriteCloser
   111  		if encoding == gzipEncoding {
   112  			encWriter, _ = gzip.NewWriterLevel(w, level)
   113  		} else if encoding == flateEncoding {
   114  			encWriter, _ = flate.NewWriter(w, level)
   115  		}
   116  		defer encWriter.Close()
   117  
   118  		w.Header().Set("Content-Encoding", encoding)
   119  		r.Header.Del(acceptEncoding)
   120  
   121  		cw := &compressResponseWriter{
   122  			w:          w,
   123  			compressor: encWriter,
   124  		}
   125  
   126  		w = httpsnoop.Wrap(w, httpsnoop.Hooks{
   127  			Write: func(httpsnoop.WriteFunc) httpsnoop.WriteFunc {
   128  				return cw.Write
   129  			},
   130  			WriteHeader: func(httpsnoop.WriteHeaderFunc) httpsnoop.WriteHeaderFunc {
   131  				return cw.WriteHeader
   132  			},
   133  			Flush: func(httpsnoop.FlushFunc) httpsnoop.FlushFunc {
   134  				return cw.Flush
   135  			},
   136  			ReadFrom: func(rff httpsnoop.ReadFromFunc) httpsnoop.ReadFromFunc {
   137  				return cw.ReadFrom
   138  			},
   139  		})
   140  
   141  		h.ServeHTTP(w, r)
   142  	})
   143  }
   144
View as plain text