diff --git a/df.go b/df.go index 9bca8aa..f6566c9 100644 --- a/df.go +++ b/df.go @@ -128,6 +128,8 @@ const ( Gzip = "gzip" // Br brotli compress Br = "br" + // Zstd zstd compress + Zstd = "zstd" ) var ( diff --git a/go.mod b/go.mod index d5cb099..b12916b 100644 --- a/go.mod +++ b/go.mod @@ -5,6 +5,7 @@ go 1.18 require ( github.com/andybalholm/brotli v1.0.6 github.com/hashicorp/golang-lru/v2 v2.0.7 + github.com/klauspost/compress v1.17.2 github.com/stretchr/testify v1.8.4 github.com/tidwall/gjson v1.17.0 github.com/vicanso/hes v0.6.2 diff --git a/go.sum b/go.sum index 766365a..2a2f3d8 100644 --- a/go.sum +++ b/go.sum @@ -5,6 +5,8 @@ github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= github.com/hashicorp/golang-lru/v2 v2.0.7 h1:a+bsQ5rvGLjzHuww6tVxozPZFVghXaHOwFs4luLUK2k= github.com/hashicorp/golang-lru/v2 v2.0.7/go.mod h1:QeFd9opnmA6QUJc5vARoKUSoFhyfM2/ZepoAG6RGpeM= +github.com/klauspost/compress v1.17.2 h1:RlWWUY/Dr4fL8qk9YG7DTZ7PDgME2V4csBXA8L/ixi4= +github.com/klauspost/compress v1.17.2/go.mod h1:ntbaceVETuRiXiv4DpjP66DpAtAGkEQskQzEyD//IeE= github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME= diff --git a/middleware/brotli.go b/middleware/brotli.go index c0f4efd..9eff577 100644 --- a/middleware/brotli.go +++ b/middleware/brotli.go @@ -32,7 +32,6 @@ import ( const ( // BrEncoding br encoding - BrEncoding = "br" maxBrQuality = 11 defaultBrQuality = 6 ) @@ -69,7 +68,7 @@ func (b *BrCompressor) Accept(c *elton.Context, bodySize int) (acceptable bool, if bodySize >= 0 && bodySize < b.getMinLength() { return } - return AcceptEncoding(c, BrEncoding) + return AcceptEncoding(c, elton.Br) } // BrotliCompress compress data by brotli diff --git a/middleware/cache_compressor.go b/middleware/cache_compressor.go index bfa0698..00e8976 100644 --- a/middleware/cache_compressor.go +++ b/middleware/cache_compressor.go @@ -26,6 +26,9 @@ import ( "bytes" "compress/gzip" "regexp" + + "github.com/klauspost/compress/zstd" + "github.com/vicanso/elton" ) type CompressionType uint8 @@ -37,6 +40,8 @@ const ( CompressionGzip // br compress CompressionBr + // zstd compress + CompressionZstd ) type CacheCompressor interface { @@ -81,18 +86,17 @@ func (br *CacheBrCompressor) Decompress(data *bytes.Buffer) (*bytes.Buffer, erro return BrotliDecompress(data.Bytes()) } func (br *CacheBrCompressor) GetEncoding() string { - return BrEncoding + return elton.Br } func (br *CacheBrCompressor) IsValid(contentType string, length int) bool { return isValidForCompress(br.ContentRegexp, br.MinLength, contentType, length) } - func (br *CacheBrCompressor) Compress(buffer *bytes.Buffer) (*bytes.Buffer, CompressionType, error) { data, err := BrotliCompress(buffer.Bytes(), br.Level) if err != nil { return nil, CompressionNone, err } - return data, CompressionBr, nil + return data, br.GetCompression(), nil } func (br *CacheBrCompressor) GetCompression() CompressionType { return CompressionBr @@ -114,7 +118,7 @@ func (g *CacheGzipCompressor) Decompress(data *bytes.Buffer) (*bytes.Buffer, err return GzipDecompress(data.Bytes()) } func (g *CacheGzipCompressor) GetEncoding() string { - return GzipEncoding + return elton.Gzip } func (g *CacheGzipCompressor) IsValid(contentType string, length int) bool { return isValidForCompress(g.ContentRegexp, g.MinLength, contentType, length) @@ -124,8 +128,39 @@ func (g *CacheGzipCompressor) Compress(buffer *bytes.Buffer) (*bytes.Buffer, Com if err != nil { return nil, CompressionNone, err } - return data, CompressionGzip, nil + return data, g.GetCompression(), nil } func (g *CacheGzipCompressor) GetCompression() CompressionType { return CompressionGzip } + +type CacheZstdCompressor struct { + Level int + MinLength int + ContentRegexp *regexp.Regexp +} + +func NewCacheZstdCompressor() *CacheZstdCompressor { + return &CacheZstdCompressor{ + Level: int(zstd.SpeedBetterCompression), + } +} +func (z *CacheZstdCompressor) Decompress(data *bytes.Buffer) (*bytes.Buffer, error) { + return ZstdDecompress(data.Bytes()) +} +func (z *CacheZstdCompressor) GetEncoding() string { + return elton.Zstd +} +func (z *CacheZstdCompressor) IsValid(contentType string, length int) bool { + return isValidForCompress(z.ContentRegexp, z.MinLength, contentType, length) +} +func (z *CacheZstdCompressor) Compress(buffer *bytes.Buffer) (*bytes.Buffer, CompressionType, error) { + data, err := ZstdCompress(buffer.Bytes(), z.Level) + if err != nil { + return nil, CompressionNone, err + } + return data, z.GetCompression(), nil +} +func (z *CacheZstdCompressor) GetCompression() CompressionType { + return CompressionZstd +} diff --git a/middleware/cache_test.go b/middleware/cache_test.go index 688257e..106be6d 100644 --- a/middleware/cache_test.go +++ b/middleware/cache_test.go @@ -177,8 +177,8 @@ func TestCacheResponseGetBody(t *testing.T) { } }, compressor: NewCacheBrCompressor(), - acceptEncoding: BrEncoding, - encoding: BrEncoding, + acceptEncoding: elton.Br, + encoding: elton.Br, body: brData, }, // 数据br, 客户端不支持br @@ -203,8 +203,8 @@ func TestCacheResponseGetBody(t *testing.T) { } }, compressor: NewCacheGzipCompressor(), - acceptEncoding: GzipEncoding, - encoding: GzipEncoding, + acceptEncoding: elton.Gzip, + encoding: elton.Gzip, body: gzipData, }, // 数据gzip,客户端不支持gzip diff --git a/middleware/compressor_test.go b/middleware/compressor_test.go index eba0885..adeeb05 100644 --- a/middleware/compressor_test.go +++ b/middleware/compressor_test.go @@ -42,7 +42,7 @@ func TestCompressor(t *testing.T) { }{ { compressor: new(GzipCompressor), - encoding: GzipEncoding, + encoding: elton.Gzip, uncompress: func(b []byte) ([]byte, error) { buffer, err := GzipDecompress(b) if err != nil { @@ -53,7 +53,7 @@ func TestCompressor(t *testing.T) { }, { compressor: new(BrCompressor), - encoding: BrEncoding, + encoding: elton.Br, uncompress: func(b []byte) ([]byte, error) { buffer, err := BrotliDecompress(b) if err != nil { diff --git a/middleware/gzip.go b/middleware/gzip.go index 8026f55..36b7ab4 100644 --- a/middleware/gzip.go +++ b/middleware/gzip.go @@ -30,11 +30,6 @@ import ( "github.com/vicanso/elton" ) -const ( - // GzipEncoding gzip encoding - GzipEncoding = "gzip" -) - type ( // GzipCompressor gzip compress GzipCompressor struct { @@ -49,7 +44,7 @@ func (g *GzipCompressor) Accept(c *elton.Context, bodySize int) (bool, string) { if bodySize >= 0 && bodySize < g.getMinLength() { return false, "" } - return AcceptEncoding(c, GzipEncoding) + return AcceptEncoding(c, elton.Gzip) } // GzipCompress compress data by gzip diff --git a/middleware/zstd.go b/middleware/zstd.go new file mode 100644 index 0000000..796d398 --- /dev/null +++ b/middleware/zstd.go @@ -0,0 +1,120 @@ +// MIT License + +// Copyright (c) 2023 Tree Xie + +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: + +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. + +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. + +package middleware + +import ( + "bytes" + "io" + + "github.com/klauspost/compress/zstd" + "github.com/vicanso/elton" +) + +type ( + // ZstdCompressor zstd compress + ZstdCompressor struct { + Level int + MinLength int + } +) + +// Accept accept zstd encoding +func (z *ZstdCompressor) Accept(c *elton.Context, bodySize int) (bool, string) { + + // 如果数据少于最低压缩长度,则不压缩(因为reader中的bodySize会设置为1,因此需要判断>=0) + if bodySize >= 0 && bodySize < z.getMinLength() { + return false, "" + } + return AcceptEncoding(c, elton.Zstd) +} + +func (z *ZstdCompressor) getLevel() int { + level := z.Level + if level <= 0 { + level = int(zstd.SpeedBetterCompression) + } + if level > int(zstd.SpeedBestCompression) { + level = int(zstd.SpeedBestCompression) + } + return level +} + +func (z *ZstdCompressor) getMinLength() int { + if z.MinLength == 0 { + return DefaultCompressMinLength + } + return z.MinLength +} + +// Compress compress data by zstd +func (z *ZstdCompressor) Compress(buf []byte, levels ...int) (*bytes.Buffer, error) { + level := z.getLevel() + if len(levels) != 0 && levels[0] != IgnoreCompression { + level = levels[0] + } + return GzipCompress(buf, level) +} + +// Pipe compress by pipe +func (z *ZstdCompressor) Pipe(c *elton.Context) error { + r := c.Body.(io.Reader) + closer, ok := c.Body.(io.Closer) + if ok { + defer closer.Close() + } + enc, err := zstd.NewWriter(c.Response, zstd.WithEncoderLevel(zstd.EncoderLevel(z.getLevel()))) + if err != nil { + return err + } + + _, err = io.Copy(enc, r) + if err != nil { + enc.Close() + return err + } + return enc.Close() +} + +// ZstdCompressor compress data by zstd +func ZstdCompress(buf []byte, level int) (*bytes.Buffer, error) { + encoder, err := zstd.NewWriter(nil, zstd.WithEncoderLevel(zstd.EncoderLevel(level))) + if err != nil { + return nil, err + } + dst := encoder.EncodeAll(buf, make([]byte, 0, len(buf))) + return bytes.NewBuffer(dst), nil +} + +// ZstdDecompress decompress data of zstd +func ZstdDecompress(buf []byte) (*bytes.Buffer, error) { + r, err := zstd.NewReader(bytes.NewBuffer(buf)) + if err != nil { + return nil, err + } + defer r.Close() + data, err := io.ReadAll(r) + if err != nil { + return nil, err + } + return bytes.NewBuffer(data), nil +}