github.com/bytom/bytom@v1.1.2-0.20221014091027-bbcba3df6075/net/http/gzip/gzip.go (about)

     1  package gzip
     2  
     3  import (
     4  	"bufio"
     5  	"compress/gzip"
     6  	"errors"
     7  	"io"
     8  	"net"
     9  	"net/http"
    10  	"strings"
    11  	"sync"
    12  )
    13  
    14  var pool = sync.Pool{
    15  	New: func() interface{} {
    16  		w, _ := gzip.NewWriterLevel(nil, gzip.BestSpeed) // #nosec
    17  		return w
    18  	},
    19  }
    20  
    21  func getWriter(w io.Writer) *gzip.Writer {
    22  	gz := pool.Get().(*gzip.Writer)
    23  	gz.Reset(w)
    24  	return gz
    25  }
    26  
    27  type Handler struct {
    28  	Handler http.Handler
    29  }
    30  
    31  func (h Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
    32  	w.Header().Add("Vary", "Accept-Encoding")
    33  	if !strings.Contains(r.Header.Get("Accept-Encoding"), "gzip") {
    34  		h.Handler.ServeHTTP(w, r)
    35  		return
    36  	}
    37  	w.Header().Set("Content-Encoding", "gzip")
    38  	gz := getWriter(w)
    39  	w = &responseWriter{gz, w}
    40  	h.Handler.ServeHTTP(w, r)
    41  	gz.Close()
    42  	pool.Put(gz)
    43  }
    44  
    45  type responseWriter struct {
    46  	w                   io.Writer // w wraps only method Write
    47  	http.ResponseWriter           // embedded for the other methods
    48  }
    49  
    50  var _ http.ResponseWriter = (*responseWriter)(nil)
    51  var _ http.Hijacker = (*responseWriter)(nil)
    52  
    53  func (w *responseWriter) Write(p []byte) (int, error) { return w.w.Write(p) }
    54  
    55  func (w *responseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
    56  	h, ok := w.ResponseWriter.(http.Hijacker)
    57  	if !ok {
    58  		return nil, nil, errors.New("not a hijacker")
    59  	}
    60  	return h.Hijack()
    61  }