github.com/sohaha/zlsgo@v1.7.13-0.20240501141223-10dd1a906f76/znet/gzip/gzip.go (about)

     1  package gzip
     2  
     3  import (
     4  	"bytes"
     5  	"compress/gzip"
     6  	"strings"
     7  
     8  	"github.com/sohaha/zlsgo/znet"
     9  )
    10  
    11  func Default() znet.HandlerFunc {
    12  	return New(Config{
    13  		CompressionLevel: 7,
    14  		PoolMaxSize:      1024,
    15  		MinContentLength: 1024,
    16  	})
    17  }
    18  
    19  func New(conf Config) znet.HandlerFunc {
    20  	pool := &poolCap{
    21  		c: make(chan *gzip.Writer, conf.PoolMaxSize),
    22  		l: conf.CompressionLevel,
    23  	}
    24  	return func(c *znet.Context) {
    25  		if !strings.Contains(c.GetHeader("Accept-Encoding"), "gzip") {
    26  			c.Next()
    27  		} else {
    28  			c.Next()
    29  			p := c.PrevContent()
    30  			if len(p.Content) < conf.MinContentLength {
    31  				return
    32  			}
    33  
    34  			g, err := pool.Get()
    35  			if err != nil {
    36  				return
    37  			}
    38  			defer pool.Put(g)
    39  
    40  			be := &bytes.Buffer{}
    41  			g.Reset(be)
    42  			_, err = g.Write(p.Content)
    43  			if err != nil {
    44  				return
    45  			}
    46  			_ = g.Flush()
    47  
    48  			c.SetHeader("Content-Encoding", "gzip")
    49  			c.Byte(p.Code.Load(), be.Bytes())
    50  		}
    51  	}
    52  }