mirror of
https://github.com/labstack/echo.git
synced 2024-12-18 16:20:53 +02:00
14e020bc07
Fixing a http.Request.Body leak on the decompress middleware that were not properly Close Removing the defer on the call to gzip.Reader, because that reader is already exausted after the call to io.Copy
147 lines
3.4 KiB
Go
147 lines
3.4 KiB
Go
package middleware
|
|
|
|
import (
|
|
"bufio"
|
|
"compress/gzip"
|
|
"io"
|
|
"io/ioutil"
|
|
"net"
|
|
"net/http"
|
|
"strings"
|
|
"sync"
|
|
|
|
"github.com/labstack/echo/v4"
|
|
)
|
|
|
|
type (
|
|
// GzipConfig defines the config for Gzip middleware.
|
|
GzipConfig struct {
|
|
// Skipper defines a function to skip middleware.
|
|
Skipper Skipper
|
|
|
|
// Gzip compression level.
|
|
// Optional. Default value -1.
|
|
Level int `yaml:"level"`
|
|
}
|
|
|
|
gzipResponseWriter struct {
|
|
io.Writer
|
|
http.ResponseWriter
|
|
}
|
|
)
|
|
|
|
const (
|
|
gzipScheme = "gzip"
|
|
)
|
|
|
|
var (
|
|
// DefaultGzipConfig is the default Gzip middleware config.
|
|
DefaultGzipConfig = GzipConfig{
|
|
Skipper: DefaultSkipper,
|
|
Level: -1,
|
|
}
|
|
)
|
|
|
|
// Gzip returns a middleware which compresses HTTP response using gzip compression
|
|
// scheme.
|
|
func Gzip() echo.MiddlewareFunc {
|
|
return GzipWithConfig(DefaultGzipConfig)
|
|
}
|
|
|
|
// GzipWithConfig return Gzip middleware with config.
|
|
// See: `Gzip()`.
|
|
func GzipWithConfig(config GzipConfig) echo.MiddlewareFunc {
|
|
// Defaults
|
|
if config.Skipper == nil {
|
|
config.Skipper = DefaultGzipConfig.Skipper
|
|
}
|
|
if config.Level == 0 {
|
|
config.Level = DefaultGzipConfig.Level
|
|
}
|
|
|
|
pool := gzipCompressPool(config)
|
|
|
|
return func(next echo.HandlerFunc) echo.HandlerFunc {
|
|
return func(c echo.Context) error {
|
|
if config.Skipper(c) {
|
|
return next(c)
|
|
}
|
|
|
|
res := c.Response()
|
|
res.Header().Add(echo.HeaderVary, echo.HeaderAcceptEncoding)
|
|
if strings.Contains(c.Request().Header.Get(echo.HeaderAcceptEncoding), gzipScheme) {
|
|
res.Header().Set(echo.HeaderContentEncoding, gzipScheme) // Issue #806
|
|
i := pool.Get()
|
|
w, ok := i.(*gzip.Writer)
|
|
if !ok {
|
|
return echo.NewHTTPError(http.StatusInternalServerError, i.(error).Error())
|
|
}
|
|
rw := res.Writer
|
|
w.Reset(rw)
|
|
defer func() {
|
|
if res.Size == 0 {
|
|
if res.Header().Get(echo.HeaderContentEncoding) == gzipScheme {
|
|
res.Header().Del(echo.HeaderContentEncoding)
|
|
}
|
|
// We have to reset response to it's pristine state when
|
|
// nothing is written to body or error is returned.
|
|
// See issue #424, #407.
|
|
res.Writer = rw
|
|
w.Reset(ioutil.Discard)
|
|
}
|
|
w.Close()
|
|
pool.Put(w)
|
|
}()
|
|
grw := &gzipResponseWriter{Writer: w, ResponseWriter: rw}
|
|
res.Writer = grw
|
|
}
|
|
return next(c)
|
|
}
|
|
}
|
|
}
|
|
|
|
func (w *gzipResponseWriter) WriteHeader(code int) {
|
|
if code == http.StatusNoContent { // Issue #489
|
|
w.ResponseWriter.Header().Del(echo.HeaderContentEncoding)
|
|
}
|
|
w.Header().Del(echo.HeaderContentLength) // Issue #444
|
|
w.ResponseWriter.WriteHeader(code)
|
|
}
|
|
|
|
func (w *gzipResponseWriter) Write(b []byte) (int, error) {
|
|
if w.Header().Get(echo.HeaderContentType) == "" {
|
|
w.Header().Set(echo.HeaderContentType, http.DetectContentType(b))
|
|
}
|
|
return w.Writer.Write(b)
|
|
}
|
|
|
|
func (w *gzipResponseWriter) Flush() {
|
|
w.Writer.(*gzip.Writer).Flush()
|
|
if flusher, ok := w.ResponseWriter.(http.Flusher); ok {
|
|
flusher.Flush()
|
|
}
|
|
}
|
|
|
|
func (w *gzipResponseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
|
|
return w.ResponseWriter.(http.Hijacker).Hijack()
|
|
}
|
|
|
|
func (w *gzipResponseWriter) Push(target string, opts *http.PushOptions) error {
|
|
if p, ok := w.ResponseWriter.(http.Pusher); ok {
|
|
return p.Push(target, opts)
|
|
}
|
|
return http.ErrNotSupported
|
|
}
|
|
|
|
func gzipCompressPool(config GzipConfig) sync.Pool {
|
|
return sync.Pool{
|
|
New: func() interface{} {
|
|
w, err := gzip.NewWriterLevel(ioutil.Discard, config.Level)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
return w
|
|
},
|
|
}
|
|
}
|