github.com/gofunct/common@v0.0.0-20190131174352-fd058c7fbf22/pkg/transport/handlers/compress.go (about)

     1  package handlers
     2  
     3  import (
     4  	"compress/flate"
     5  	"compress/gzip"
     6  	"io"
     7  	"net/http"
     8  	"strings"
     9  )
    10  
    11  type compressResponseWriter struct {
    12  	io.Writer
    13  	http.ResponseWriter
    14  	http.Hijacker
    15  	http.Flusher
    16  	http.CloseNotifier
    17  }
    18  
    19  func (w *compressResponseWriter) WriteHeader(c int) {
    20  	w.ResponseWriter.Header().Del("Content-Length")
    21  	w.ResponseWriter.WriteHeader(c)
    22  }
    23  
    24  func (w *compressResponseWriter) Header() http.Header {
    25  	return w.ResponseWriter.Header()
    26  }
    27  
    28  func (w *compressResponseWriter) Write(b []byte) (int, error) {
    29  	h := w.ResponseWriter.Header()
    30  	if h.Get("Content-Type") == "" {
    31  		h.Set("Content-Type", http.DetectContentType(b))
    32  	}
    33  	h.Del("Content-Length")
    34  
    35  	return w.Writer.Write(b)
    36  }
    37  
    38  type flusher interface {
    39  	Flush() error
    40  }
    41  
    42  func (w *compressResponseWriter) Flush() {
    43  	// Flush compressed data if compressor supports it.
    44  	if f, ok := w.Writer.(flusher); ok {
    45  		f.Flush()
    46  	}
    47  	// Flush HTTP response.
    48  	if w.Flusher != nil {
    49  		w.Flusher.Flush()
    50  	}
    51  }
    52  
    53  // CompressHandler gzip compresses HTTP responses for clients that support it
    54  // via the 'Accept-Encoding' header.
    55  //
    56  // Compressing TLS traffic may leak the page contents to an attacker if the
    57  // page contains user input: http://security.stackexchange.com/a/102015/12208
    58  func CompressHandler(h http.Handler) http.Handler {
    59  	return CompressHandlerLevel(h, gzip.DefaultCompression)
    60  }
    61  
    62  // CompressHandlerLevel gzip compresses HTTP responses with specified compression level
    63  // for clients that support it via the 'Accept-Encoding' header.
    64  //
    65  // The compression level should be gzip.DefaultCompression, gzip.NoCompression,
    66  // or any integer value between gzip.BestSpeed and gzip.BestCompression inclusive.
    67  // gzip.DefaultCompression is used in case of invalid compression level.
    68  func CompressHandlerLevel(h http.Handler, level int) http.Handler {
    69  	if level < gzip.DefaultCompression || level > gzip.BestCompression {
    70  		level = gzip.DefaultCompression
    71  	}
    72  
    73  	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    74  	L:
    75  		for _, enc := range strings.Split(r.Header.Get("Accept-Encoding"), ",") {
    76  			switch strings.TrimSpace(enc) {
    77  			case "gzip":
    78  				w.Header().Set("Content-Encoding", "gzip")
    79  				w.Header().Add("Vary", "Accept-Encoding")
    80  
    81  				gw, _ := gzip.NewWriterLevel(w, level)
    82  				defer gw.Close()
    83  
    84  				h, hok := w.(http.Hijacker)
    85  				if !hok { /* w is not Hijacker... oh well... */
    86  					h = nil
    87  				}
    88  
    89  				f, fok := w.(http.Flusher)
    90  				if !fok {
    91  					f = nil
    92  				}
    93  
    94  				cn, cnok := w.(http.CloseNotifier)
    95  				if !cnok {
    96  					cn = nil
    97  				}
    98  
    99  				w = &compressResponseWriter{
   100  					Writer:         gw,
   101  					ResponseWriter: w,
   102  					Hijacker:       h,
   103  					Flusher:        f,
   104  					CloseNotifier:  cn,
   105  				}
   106  
   107  				break L
   108  			case "deflate":
   109  				w.Header().Set("Content-Encoding", "deflate")
   110  				w.Header().Add("Vary", "Accept-Encoding")
   111  
   112  				fw, _ := flate.NewWriter(w, level)
   113  				defer fw.Close()
   114  
   115  				h, hok := w.(http.Hijacker)
   116  				if !hok { /* w is not Hijacker... oh well... */
   117  					h = nil
   118  				}
   119  
   120  				f, fok := w.(http.Flusher)
   121  				if !fok {
   122  					f = nil
   123  				}
   124  
   125  				cn, cnok := w.(http.CloseNotifier)
   126  				if !cnok {
   127  					cn = nil
   128  				}
   129  
   130  				w = &compressResponseWriter{
   131  					Writer:         fw,
   132  					ResponseWriter: w,
   133  					Hijacker:       h,
   134  					Flusher:        f,
   135  					CloseNotifier:  cn,
   136  				}
   137  
   138  				break L
   139  			}
   140  		}
   141  
   142  		h.ServeHTTP(w, r)
   143  	})
   144  }