recover from panics such as out-of-bounds array access & nil pointer access, print a stack trace and return 5xx error instead of the service crashing and relying on Execution framework to handle crashes
146 lines
3.4 KiB
Go
146 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
|
|
},
|
|
}
|
|
}
|