github.com/gocaveman/caveman@v0.0.0-20191211162744-0ddf99dbdf6e/webutil/response-gzip.go (about)

     1  package webutil
     2  
     3  import (
     4  	"bufio"
     5  	"compress/gzip"
     6  	"net"
     7  	"net/http"
     8  	"strings"
     9  )
    10  
    11  // TODO: make a GzipResponseWriter and the ChainHandler that creates it,
    12  // use WrapResponseWriter as a base (gzip is a great example of how it's used).
    13  
    14  type GzipHandler struct{}
    15  
    16  func NewGzipHandler() GzipHandler {
    17  	return GzipHandler{}
    18  }
    19  
    20  func (h GzipHandler) ServeHTTPChain(w http.ResponseWriter, r *http.Request) (http.ResponseWriter, *http.Request) {
    21  
    22  	// TODO: we really should properly parse the accept-encoding header, but its kind of a pain; go simple for now
    23  
    24  	// Verify client accepts gzip
    25  	if !strings.Contains(r.Header.Get("accept-encoding"), "gzip") {
    26  		return w, r
    27  	}
    28  
    29  	w.Header().Set("Content-Encoding", "gzip")
    30  
    31  	gzw := gzip.NewWriter(w)
    32  
    33  	retw := &GzipResponseWriter{
    34  		ResponseWriter: w,
    35  		gzw:            gzw,
    36  	}
    37  	return retw, r
    38  }
    39  
    40  type GzipResponseWriter struct {
    41  	http.ResponseWriter
    42  	gzw      *gzip.Writer
    43  	written  bool
    44  	hijacked bool
    45  }
    46  
    47  func (w *GzipResponseWriter) Write(p []byte) (int, error) {
    48  
    49  	// TODO: Review if there are some common cases where we do not want to apply gzipping, e.g.
    50  	// - Binary content types possibly don't need it
    51  	// - See if it interferes at all with things like server-sent events
    52  	// - Also websockets, will things work correctly?
    53  	// The logic could be improved here to automatically disable in cases we can detect.
    54  
    55  	if !w.written {
    56  
    57  		// if no content type, we autodetect it here - because if we don't then the underlying default
    58  		// ResponseWriter will try it on the gzipped bytes which will say that everthing is a gzip file
    59  		if w.Header().Get("content-type") == "" {
    60  			ct := http.DetectContentType(p)
    61  			w.Header().Set("content-type", ct)
    62  		}
    63  
    64  		w.WriteHeader(http.StatusOK)
    65  
    66  		// Send an empty write through to the underlying writer - to be sure the context cancelation works.
    67  		// w.gzw.Write() below is not guaranteed to reach the underlying response writer in this call.
    68  		w.ResponseWriter.Write(nil)
    69  
    70  	}
    71  
    72  	return w.gzw.Write(p)
    73  }
    74  
    75  func (w *GzipResponseWriter) WriteHeader(c int) {
    76  	w.written = true
    77  	w.ResponseWriter.WriteHeader(c)
    78  }
    79  
    80  // func (w *GzipResponseWriter) Close() (err error) {
    81  // 	w.gzw.Close()
    82  // 	if c, ok := w.ResponseWriter.(io.Closer); ok {
    83  // 		err = c.Close()
    84  // 	}
    85  // 	return err
    86  // }
    87  
    88  // Flush calls Flush on the gzip.Writer and then the underlying ResponseWriter
    89  func (w *GzipResponseWriter) Flush() {
    90  	// do nothing if hijacked
    91  	if w.hijacked {
    92  		return
    93  	}
    94  	w.gzw.Flush()
    95  	w.ResponseWriter.(http.Flusher).Flush()
    96  }
    97  
    98  func (w *GzipResponseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
    99  	w.hijacked = true
   100  	return w.ResponseWriter.(http.Hijacker).Hijack()
   101  }
   102  
   103  func (w *GzipResponseWriter) CloseNotify() <-chan bool {
   104  	return w.ResponseWriter.(http.CloseNotifier).CloseNotify()
   105  }
   106  
   107  func (w *GzipResponseWriter) Push(target string, opts *http.PushOptions) error {
   108  	return w.ResponseWriter.(http.Pusher).Push(target, opts)
   109  }