github.com/gocaveman/caveman@v0.0.0-20191211162744-0ddf99dbdf6e/router/httpgzip/httpgzip.go (about)

     1  package httpgzip
     2  
     3  import (
     4  	"bufio"
     5  	"compress/gzip"
     6  	"net"
     7  	"net/http"
     8  	"strings"
     9  
    10  	"github.com/gocaveman/caveman/router"
    11  )
    12  
    13  // New will return a new GZIPChainHandler ready to compress responses using gzip.
    14  // By default, all responses matching pathPrefix (empty to mean all) will be
    15  // compressed.  Individual responses can opt out of compression by setting the response
    16  // "Content-Encoding" header to something other than "gzip".
    17  func New(pathPrefix string) *GZIPChainHandler {
    18  	return &GZIPChainHandler{
    19  		PathPrefix: pathPrefix,
    20  		Sequence:   router.RouteSequenceSetup,
    21  	}
    22  }
    23  
    24  type GZIPChainHandler struct {
    25  	PathPrefix string
    26  	Sequence   float64
    27  }
    28  
    29  func (h *GZIPChainHandler) RoutePathPrefix() string {
    30  	return h.PathPrefix
    31  }
    32  
    33  func (h *GZIPChainHandler) RouteSequence() float64 {
    34  	return h.Sequence
    35  }
    36  
    37  func (h *GZIPChainHandler) ServeHTTPChain(w http.ResponseWriter, r *http.Request) (wnext http.ResponseWriter, rnext *http.Request) {
    38  
    39  	// imperfect but simple and works for all practical cases
    40  	if !strings.Contains(r.Header.Get("Accept-Encoding"), "gzip") {
    41  		return w, r
    42  	}
    43  
    44  	wwrap := router.NewWrapResponseWriter(w, r)
    45  
    46  	// indicate gzip response
    47  	wwrap.Header().Set("Content-Encoding", "gzip")
    48  	// log.Printf("wwrap.Header(): %v", w.Header())
    49  
    50  	gw := gzip.NewWriter(w)
    51  	gwWritten := false
    52  
    53  	wwrap.SetWriteFunc(func(b []byte) (int, error) {
    54  		if gwWritten && gw != nil { // if we've been gzipping, keep going
    55  			return gw.Write(b)
    56  		}
    57  		// see if header is set up for it and gw hasn't been nuked
    58  		if gw != nil && wwrap.Header().Get("Content-Encoding") == "gzip" {
    59  			// start gzip writing
    60  			gwWritten = true
    61  			return gw.Write(b)
    62  		}
    63  		// nop, looks like they opted out, do a regular write
    64  		return wwrap.Parent().Write(b)
    65  	})
    66  
    67  	wwrap.SetFlushFunc(func() {
    68  		if gwWritten && gw != nil { // gzip flush if gw hasn't been nuked and some data was written
    69  			gw.Flush()
    70  		}
    71  		wwrap.Parent().(http.Flusher).Flush()
    72  	})
    73  
    74  	wwrap.SetHijackFunc(func() (net.Conn, *bufio.ReadWriter, error) {
    75  		gw = nil // disable any further gzip stuff
    76  		return wwrap.Parent().(http.Hijacker).Hijack()
    77  	})
    78  
    79  	wnext, rnext = wwrap, r
    80  
    81  	// set deferred handler so we can flush the gzip stream if any data was written
    82  	rnext = rnext.WithContext(router.DeferChainHandler(rnext.Context(), router.ChainHandlerFunc(func(w http.ResponseWriter, r *http.Request) (wret http.ResponseWriter, rret *http.Request) {
    83  		wret, rret = w, r
    84  		if gwWritten && gw != nil {
    85  			// log.Printf("flushing...")
    86  			w.(http.Flusher).Flush() // calling Flush will trigger the gzip flush
    87  			gw.Close()               // then closer the gzip writer
    88  			// gw.Reset(w)              // note totally sure on this, but this would mean that further writes will start a new gzip stream...
    89  		}
    90  		return
    91  	})))
    92  
    93  	return
    94  }