github.com/keysonZZZ/kmg@v0.0.0-20151121023212-05317bfd7d39/kmgNet/kmgHttp/Compress.go (about)

     1  package kmgHttp
     2  
     3  import (
     4  	"compress/flate"
     5  	"compress/gzip"
     6  	"github.com/bronze1man/kmg/kmgCompress"
     7  	"io"
     8  	"net/http"
     9  	"strconv"
    10  	"strings"
    11  )
    12  
    13  type ResponseWriterWraper struct {
    14  	io.Writer
    15  	http.ResponseWriter
    16  }
    17  
    18  // TODO 这个地方缺少猜类型功能,导致不能直接wrapper到任意Handler上去.(会导致 html Type 错误.)
    19  // TODO 有时候压缩会有负效果,此处应该可以自动判断出来.
    20  func (w ResponseWriterWraper) Write(b []byte) (int, error) {
    21  	return w.Writer.Write(b)
    22  }
    23  
    24  func CompressHandlerFunc(f func(w http.ResponseWriter, r *http.Request)) func(w http.ResponseWriter, r *http.Request) {
    25  	return func(w http.ResponseWriter, r *http.Request) {
    26  		acceptEncoding := r.Header.Get("Accept-Encoding")
    27  		switch {
    28  		case strings.Contains(acceptEncoding, "deflate"):
    29  			w.Header().Set("Content-Encoding", "deflate")
    30  			gzw, err := flate.NewWriter(w, -1)
    31  			if err != nil {
    32  				panic(err)
    33  			}
    34  			defer gzw.Close()
    35  			gzr := ResponseWriterWraper{Writer: gzw, ResponseWriter: w}
    36  			f(gzr, r)
    37  			return
    38  		case strings.Contains(acceptEncoding, "gzip"):
    39  			w.Header().Set("Content-Encoding", "gzip")
    40  			gzw := gzip.NewWriter(w)
    41  			defer gzw.Close()
    42  			gzr := ResponseWriterWraper{Writer: gzw, ResponseWriter: w}
    43  			f(gzr, r)
    44  			return
    45  		default:
    46  			f(w, r)
    47  		}
    48  	}
    49  }
    50  
    51  func CompressHandler(fn http.Handler) http.Handler {
    52  	return http.HandlerFunc(CompressHandlerFunc(fn.ServeHTTP))
    53  }
    54  
    55  // a flate(DEFLATE) compress wrap around http request and response,
    56  // !!not handle any http header!!
    57  func HttpHandleCompressFlateWrap(fn http.Handler) http.Handler {
    58  	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    59  		oldBody := r.Body
    60  		defer oldBody.Close()
    61  		r.Body = flate.NewReader(oldBody)
    62  		//w.Header().Set("Content-Encoding", "deflate")
    63  		gzw, err := flate.NewWriter(w, -1)
    64  		if err != nil {
    65  			panic(err)
    66  		}
    67  		defer gzw.Close()
    68  		gzr := ResponseWriterWraper{Writer: gzw, ResponseWriter: w}
    69  		fn.ServeHTTP(gzr, r)
    70  	})
    71  }
    72  
    73  // a flate(DEFLATE) compress wrap around http request and response,
    74  // !!not handle any http header!!
    75  func HttpHandleCompressGzipWrap(fn http.Handler) http.Handler {
    76  	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    77  		oldBody := r.Body
    78  		defer oldBody.Close()
    79  		var err error
    80  		r.Body, err = gzip.NewReader(oldBody)
    81  		if err != nil {
    82  			panic(err)
    83  		}
    84  		//w.Header().Set("Content-Encoding", "gzip")
    85  		gzw := gzip.NewWriter(w)
    86  		defer gzw.Close()
    87  		gzr := ResponseWriterWraper{Writer: gzw, ResponseWriter: w}
    88  		fn.ServeHTTP(gzr, r)
    89  	})
    90  }
    91  
    92  // 如果使用了这个,golang的猜测返回类型的东西会挂掉.请设置输出的内容的类型
    93  // 这个地方返回错误没有什么意义,(调用者无法处理)
    94  func CompressWriteByte(r *http.Request, w http.ResponseWriter, b []byte) {
    95  	acceptEncoding := r.Header.Get("Accept-Encoding")
    96  	switch {
    97  	case strings.Contains(acceptEncoding, "deflate"):
    98  		w.Header().Set("Content-Encoding", "deflate")
    99  		b = kmgCompress.FlateMustCompress(b)
   100  		w.Header().Set("Content-Length", strconv.Itoa(len(b)))
   101  		w.Write(b)
   102  		return
   103  	case strings.Contains(acceptEncoding, "gzip"):
   104  		w.Header().Set("Content-Encoding", "gzip")
   105  		b = kmgCompress.GzipMustCompress(b)
   106  		w.Header().Set("Content-Length", strconv.Itoa(len(b)))
   107  		w.Write(b)
   108  	default:
   109  		w.Header().Set("Content-Length", strconv.Itoa(len(b)))
   110  		w.Write(b)
   111  	}
   112  }
   113  
   114  func CompressWriteHtml(req *http.Request, w http.ResponseWriter, html []byte) {
   115  	w.Header().Set("Content-Type", "text/html; charset=utf-8")
   116  	CompressWriteByte(req, w, html)
   117  }