From 454c79cb3066091fd41b8266c21a04de3eab9e12 Mon Sep 17 00:00:00 2001 From: vicanso Date: Wed, 15 Apr 2020 22:14:02 +0800 Subject: [PATCH] feat: add compress middleware --- elton.go | 3 +- middleware/compress.go | 179 +++++++++++++++++++++ middleware/compress_test.go | 299 ++++++++++++++++++++++++++++++++++++ middleware/gzip.go | 98 ++++++++++++ middleware/gzip_test.go | 74 +++++++++ 5 files changed, 651 insertions(+), 2 deletions(-) create mode 100644 middleware/compress.go create mode 100644 middleware/compress_test.go create mode 100644 middleware/gzip.go create mode 100644 middleware/gzip_test.go diff --git a/elton.go b/elton.go index 5d5c76c..1e1a4ea 100644 --- a/elton.go +++ b/elton.go @@ -51,8 +51,6 @@ const ( type ( // Skipper check for skip middleware Skipper func(c *Context) bool - // Validator validate function for param - Validator func(value string) error // RouterInfo router's info RouterInfo struct { Method string `json:"method,omitempty"` @@ -281,6 +279,7 @@ func (e *Elton) Handle(method, path string, handlerList ...Handler) { index := -1 var traceInfos TraceInfos if e.EnableTrace { + // TODO 复用tracInfos traceInfos = make(TraceInfos, 0, maxNext) } c.Next = func() error { diff --git a/middleware/compress.go b/middleware/compress.go new file mode 100644 index 0000000..ffa7dfb --- /dev/null +++ b/middleware/compress.go @@ -0,0 +1,179 @@ +// MIT License + +// Copyright (c) 2020 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" + "regexp" + "strings" + + "github.com/vicanso/elton" +) + +var ( + DefaultCompressRegexp = regexp.MustCompile("text|javascript|json") +) + +const ( + DefaultCompressMinLength = 1024 +) + +type ( + // Compressor compressor interface + Compressor interface { + // Accept accept check function + Accept(c *elton.Context, bodySize int) (acceptable bool, encoding string) + // Compress compress function + Compress([]byte) (*bytes.Buffer, error) + // Pipe pipe function + Pipe(*elton.Context) error + } + // Config compress config + CompressConfig struct { + // Checker check the data is compressable + Checker *regexp.Regexp + // Compressors compressor list + Compressors []Compressor + // Skipper skipper function + Skipper elton.Skipper + } +) + +// AcceptEncoding check request accept encoding +func AcceptEncoding(c *elton.Context, encoding string) (bool, string) { + acceptEncoding := c.GetRequestHeader(elton.HeaderAcceptEncoding) + if strings.Contains(acceptEncoding, encoding) { + return true, encoding + } + return false, "" +} + +// AddCompressor add compressor +func (conf *CompressConfig) AddCompressor(compressor Compressor) { + if conf.Compressors == nil { + conf.Compressors = make([]Compressor, 0) + } + conf.Compressors = append(conf.Compressors, compressor) +} + +// NewCompressConfig create a new compress config +func NewCompressConfig(compressors ...Compressor) CompressConfig { + cfg := CompressConfig{} + cfg.AddCompressor(new(GzipCompressor)) + return cfg +} + +// NewDefaultCompress create a new default compress middleware(include gzip) +func NewDefaultCompress() elton.Handler { + cfg := NewCompressConfig(new(GzipCompressor)) + return NewCompress(cfg) +} + +// NewCompress create a new compress middleware +func NewCompress(config CompressConfig) elton.Handler { + skipper := config.Skipper + if skipper == nil { + skipper = elton.DefaultSkipper + } + checker := config.Checker + if checker == nil { + checker = DefaultCompressRegexp + } + compressorList := config.Compressors + return func(c *elton.Context) (err error) { + if skipper(c) || compressorList == nil { + return c.Next() + } + err = c.Next() + if err != nil { + return + } + isReaderBody := c.IsReaderBody() + // 如果数据为空，而且body不是reader，直接跳过 + if c.BodyBuffer == nil && !isReaderBody { + return + } + + // encoding 不为空，已做处理，无需要压缩 + if c.GetHeader(elton.HeaderContentEncoding) != "" { + return + } + contentType := c.GetHeader(elton.HeaderContentType) + // 数据类型为非可压缩，则返回 + if !checker.MatchString(contentType) { + return + } + + var body []byte + if c.BodyBuffer != nil { + body = c.BodyBuffer.Bytes() + } + // 对于reader类，无法判断长度，认为长度为-1 + bodySize := -1 + if !isReaderBody { + // 如果数据长度少于最小压缩长度 + bodySize = len(body) + } + + fillHeader := func(encoding string) { + c.SetHeader(elton.HeaderContentEncoding, encoding) + c.AddHeader("Vary", "Accept-Encoding") + etagValue := c.GetHeader(elton.HeaderETag) + // after compress, etag should be weak etag + if etagValue != "" && !strings.HasPrefix(etagValue, "W/") { + c.SetHeader(elton.HeaderETag, "W/"+etagValue) + } + } + + for _, compressor := range compressorList { + acceptable, encoding := compressor.Accept(c, bodySize) + if !acceptable { + continue + } + if isReaderBody { + fillHeader(encoding) + err = compressor.Pipe(c) + // 如果出错直接返回 + if err != nil { + return + } + // 成功跳出循环 + // pipe 将数据直接转至原有的Response，因此设置committed为true + c.Committed = true + // 清除 reader body + c.Body = nil + break + } + + newBuf, e := compressor.Compress(body) + // 如果压缩成功，则使用压缩数据 + // 失败则忽略 + if e == nil { + fillHeader(encoding) + c.BodyBuffer = newBuf + break + } + } + return + } +} diff --git a/middleware/compress_test.go b/middleware/compress_test.go new file mode 100644 index 0000000..959aeb6 --- /dev/null +++ b/middleware/compress_test.go @@ -0,0 +1,299 @@ +// MIT License + +// Copyright (c) 2020 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" + "errors" + "math/rand" + "net/http/httptest" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/vicanso/elton" +) + +var letterRunes = []rune("ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789-_") + +type testCompressor struct{} + +func (t *testCompressor) Accept(c *elton.Context, bodySize int) (acceptable bool, encoding string) { + return AcceptEncoding(c, "br") +} + +func (t *testCompressor) Compress(buf []byte) (*bytes.Buffer, error) { + return bytes.NewBufferString("abcd"), nil +} + +func (t *testCompressor) Pipe(c *elton.Context) error { + return nil +} + +// randomString get random string +func randomString(n int) string { + b := make([]rune, n) + rand.Seed(time.Now().UnixNano()) + for i := range b { + b[i] = letterRunes[rand.Intn(len(letterRunes))] + } + return string(b) +} + +func TestAcceptEncoding(t *testing.T) { + assert := assert.New(t) + req := httptest.NewRequest("GET", "/", nil) + c := elton.NewContext(nil, req) + acceptable, encoding := AcceptEncoding(c, elton.Gzip) + assert.False(acceptable) + assert.Empty(encoding) + + c.SetRequestHeader(elton.HeaderAcceptEncoding, elton.Gzip) + acceptable, encoding = AcceptEncoding(c, elton.Gzip) + assert.True(acceptable) + assert.Equal(elton.Gzip, encoding) +} + +func TestCompress(t *testing.T) { + t.Run("skip", func(t *testing.T) { + assert := assert.New(t) + c := elton.NewContext(nil, nil) + done := false + c.Next = func() error { + done = true + return nil + } + fn := NewCompress(CompressConfig{ + Skipper: func(c *elton.Context) bool { + return true + }, + }) + err := fn(c) + assert.Nil(err) + assert.True(done) + }) + + t.Run("nil body", func(t *testing.T) { + assert := assert.New(t) + c := elton.NewContext(nil, nil) + done := false + c.Next = func() error { + done = true + return nil + } + fn := NewDefaultCompress() + err := fn(c) + assert.Nil(err) + assert.True(done) + }) + + t.Run("return error", func(t *testing.T) { + assert := assert.New(t) + c := elton.NewContext(nil, nil) + customErr := errors.New("abccd") + c.Next = func() error { + return customErr + } + fn := NewDefaultCompress() + err := fn(c) + assert.Equal(err, customErr) + }) + + t.Run("normal", func(t *testing.T) { + assert := assert.New(t) + conf := NewCompressConfig(&GzipCompressor{ + MinLength: 1, + }) + fn := NewCompress(conf) + + req := httptest.NewRequest("GET", "/users/me", nil) + req.Header.Set(elton.HeaderAcceptEncoding, "gzip") + resp := httptest.NewRecorder() + c := elton.NewContext(resp, req) + c.SetHeader(elton.HeaderContentType, "text/html") + c.BodyBuffer = bytes.NewBuffer([]byte("" + randomString(8192) + "")) + originalSize := c.BodyBuffer.Len() + done := false + c.Next = func() error { + done = true + return nil + } + err := fn(c) + assert.Nil(err) + assert.True(done) + assert.True(c.BodyBuffer.Len() < originalSize) + assert.Equal(elton.Gzip, c.GetHeader(elton.HeaderContentEncoding)) + }) + + t.Run("encoding done", func(t *testing.T) { + assert := assert.New(t) + fn := NewDefaultCompress() + req := httptest.NewRequest("GET", "/users/me", nil) + resp := httptest.NewRecorder() + c := elton.NewContext(resp, req) + c.Next = func() error { + return nil + } + body := bytes.NewBufferString(randomString(4096)) + c.BodyBuffer = body + c.SetHeader(elton.HeaderContentEncoding, "gzip") + err := fn(c) + assert.Nil(err) + assert.Equal(body.Bytes(), c.BodyBuffer.Bytes()) + }) + + t.Run("body size is less than min length", func(t *testing.T) { + assert := assert.New(t) + fn := NewDefaultCompress() + + req := httptest.NewRequest("GET", "/users/me", nil) + req.Header.Set(elton.HeaderAcceptEncoding, "gzip") + resp := httptest.NewRecorder() + c := elton.NewContext(resp, req) + c.Next = func() error { + return nil + } + body := bytes.NewBufferString("abcd") + c.BodyBuffer = body + c.SetHeader(elton.HeaderContentType, "text/plain") + err := fn(c) + assert.Nil(err) + assert.Equal(body.Bytes(), c.BodyBuffer.Bytes()) + assert.Empty(c.GetHeader(elton.HeaderContentEncoding)) + }) + + t.Run("image should not be compress", func(t *testing.T) { + assert := assert.New(t) + + fn := NewDefaultCompress() + + req := httptest.NewRequest("GET", "/users/me", nil) + req.Header.Set(elton.HeaderAcceptEncoding, "gzip") + resp := httptest.NewRecorder() + c := elton.NewContext(resp, req) + c.SetHeader(elton.HeaderContentType, "image/jpeg") + c.Next = func() error { + return nil + } + body := bytes.NewBufferString(randomString(4096)) + c.BodyBuffer = body + err := fn(c) + assert.Nil(err) + assert.Equal(body.Bytes(), c.BodyBuffer.Bytes()) + assert.Empty(c.GetHeader(elton.HeaderContentEncoding)) + }) + + t.Run("not accept gzip should not compress", func(t *testing.T) { + assert := assert.New(t) + + fn := NewDefaultCompress() + + req := httptest.NewRequest("GET", "/users/me", nil) + resp := httptest.NewRecorder() + c := elton.NewContext(resp, req) + c.SetHeader(elton.HeaderContentType, "text/html") + c.Next = func() error { + return nil + } + body := bytes.NewBufferString(randomString(4096)) + c.BodyBuffer = body + err := fn(c) + assert.Nil(err) + assert.Equal(body.Bytes(), c.BodyBuffer.Bytes()) + assert.Empty(c.GetHeader(elton.HeaderContentEncoding)) + }) + + t.Run("custom compress", func(t *testing.T) { + assert := assert.New(t) + compressorList := make([]Compressor, 0) + compressorList = append(compressorList, new(testCompressor)) + fn := NewCompress(CompressConfig{ + Compressors: compressorList, + }) + + req := httptest.NewRequest("GET", "/users/me", nil) + req.Header.Set("Accept-Encoding", "gzip, deflate, br") + resp := httptest.NewRecorder() + c := elton.NewContext(resp, req) + c.SetHeader(elton.HeaderContentType, "text/html") + c.BodyBuffer = bytes.NewBufferString("" + randomString(8192) + "") + done := false + c.Next = func() error { + done = true + return nil + } + err := fn(c) + assert.Nil(err) + assert.True(done) + assert.Equal(4, c.BodyBuffer.Len()) + assert.Equal("br", c.GetHeader(elton.HeaderContentEncoding)) + }) + + t.Run("update etag", func(t *testing.T) { + assert := assert.New(t) + compressorList := make([]Compressor, 0) + compressorList = append(compressorList, new(GzipCompressor)) + fn := NewCompress(CompressConfig{ + Compressors: compressorList, + }) + + req := httptest.NewRequest("GET", "/users/me", nil) + req.Header.Set("Accept-Encoding", "gzip") + resp := httptest.NewRecorder() + c := elton.NewContext(resp, req) + c.SetHeader(elton.HeaderContentType, "text/html") + c.SetHeader(elton.HeaderETag, "123") + c.BodyBuffer = bytes.NewBufferString("" + randomString(8192) + "") + done := false + c.Next = func() error { + done = true + return nil + } + err := fn(c) + assert.Nil(err) + assert.True(done) + assert.Equal("W/123", c.GetHeader(elton.HeaderETag)) + }) + + t.Run("reader body", func(t *testing.T) { + assert := assert.New(t) + + fn := NewDefaultCompress() + + req := httptest.NewRequest("GET", "/users/me", nil) + req.Header.Set(elton.HeaderAcceptEncoding, "gzip") + resp := httptest.NewRecorder() + c := elton.NewContext(resp, req) + c.SetHeader(elton.HeaderContentType, "text/html") + c.Next = func() error { + return nil + } + body := bytes.NewBufferString(randomString(4096)) + c.Body = body + err := fn(c) + assert.True(c.Committed) + assert.Nil(err) + assert.NotEmpty(resp.Body.Bytes()) + assert.Equal(elton.Gzip, c.GetHeader(elton.HeaderContentEncoding)) + }) +} diff --git a/middleware/gzip.go b/middleware/gzip.go new file mode 100644 index 0000000..186b810 --- /dev/null +++ b/middleware/gzip.go @@ -0,0 +1,98 @@ +// MIT License + +// Copyright (c) 2020 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" + "compress/gzip" + "io" + + "github.com/vicanso/elton" +) + +const ( + // GzipEncoding gzip encoding + GzipEncoding = "gzip" +) + +type ( + // GzipCompressor gzip compress + GzipCompressor struct { + Level int + MinLength int + } +) + +// Accept accept gzip encoding +func (g *GzipCompressor) Accept(c *elton.Context, bodySize int) (acceptable bool, encoding string) { + // 如果数据少于最低压缩长度，则不压缩 + if bodySize >= 0 && bodySize < g.getMinLength() { + return + } + return AcceptEncoding(c, GzipEncoding) +} + +// Compress compress data by gzip +func (g *GzipCompressor) Compress(buf []byte) (*bytes.Buffer, error) { + level := g.getLevel() + buffer := new(bytes.Buffer) + + w, _ := gzip.NewWriterLevel(buffer, level) + defer w.Close() + _, err := w.Write(buf) + if err != nil { + return nil, err + } + return buffer, nil +} + +func (g *GzipCompressor) getLevel() int { + level := g.Level + if level <= 0 { + level = gzip.DefaultCompression + } + if level > gzip.BestCompression { + level = gzip.BestCompression + } + return level +} + +func (g *GzipCompressor) getMinLength() int { + if g.MinLength == 0 { + return DefaultCompressMinLength + } + return g.MinLength +} + +// Pipe compress by pipe +func (g *GzipCompressor) Pipe(c *elton.Context) (err error) { + r := c.Body.(io.Reader) + closer, ok := c.Body.(io.Closer) + if ok { + defer closer.Close() + } + w, _ := gzip.NewWriterLevel(c.Response, g.getLevel()) + defer w.Close() + _, err = io.Copy(w, r) + return +} diff --git a/middleware/gzip_test.go b/middleware/gzip_test.go new file mode 100644 index 0000000..ad32b39 --- /dev/null +++ b/middleware/gzip_test.go @@ -0,0 +1,74 @@ +// MIT License + +// Copyright (c) 2020 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" + "compress/gzip" + "io/ioutil" + "net/http/httptest" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/vicanso/elton" +) + +func TestGzipCompress(t *testing.T) { + assert := assert.New(t) + originalData := randomString(1024) + g := new(GzipCompressor) + req := httptest.NewRequest("GET", "/users/me", nil) + req.Header.Set("Accept-Encoding", "gzip, deflate, br") + c := elton.NewContext(nil, req) + acceptable, encoding := g.Accept(c, 0) + assert.False(acceptable) + assert.Empty(encoding) + acceptable, encoding = g.Accept(c, len(originalData)) + assert.True(acceptable) + assert.Equal(GzipEncoding, encoding) + buf, err := g.Compress([]byte(originalData)) + assert.Nil(err) + r, err := gzip.NewReader(bytes.NewReader(buf.Bytes())) + assert.Nil(err) + defer r.Close() + originlBuf, _ := ioutil.ReadAll(r) + assert.Equal(originalData, string(originlBuf)) +} + +func TestGzipPipe(t *testing.T) { + assert := assert.New(t) + resp := httptest.NewRecorder() + originalData := randomString(1024) + c := elton.NewContext(resp, nil) + + c.Body = bytes.NewReader([]byte(originalData)) + + g := new(GzipCompressor) + err := g.Pipe(c) + assert.Nil(err) + r, err := gzip.NewReader(resp.Body) + assert.Nil(err) + defer r.Close() + buf, _ := ioutil.ReadAll(r) + assert.Equal(originalData, string(buf)) +}