Add Brotli support with pure go implementation

go.mod

@@ -3,6 +3,7 @@ module github.com/labstack/echo/v4
 go 1.14
 
 require (
+	github.com/andybalholm/brotli v1.0.0
 	github.com/dgrijalva/jwt-go v3.2.0+incompatible
 	github.com/labstack/gommon v0.3.0
 	github.com/mattn/go-colorable v0.1.6 // indirect

go.sum

@@ -1,3 +1,5 @@
+github.com/andybalholm/brotli v1.0.0 h1:7UCwP93aiSfvWpapti8g88vVVGp2qqtGyePsSuDafo4=
+github.com/andybalholm/brotli v1.0.0/go.mod h1:loMXtMfwqflxFJPmdbJO0a3KNoPuLBgiu3qAvBg8x/Y=
 github.com/davecgh/go-spew v1.1.0 h1:ZDRjVQ15GmhC3fiQ8ni8+OwkZQO4DARzQgrnXU1Liz8=
 github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/dgrijalva/jwt-go v1.0.2 h1:KPldsxuKGsS2FPWsNeg9ZO18aCrGKujPoWXn2yo+KQM=

middleware/compress_brotli.go

@@ -0,0 +1,121 @@
+package middleware
+
+import (
+	"bufio"
+	"io"
+	"io/ioutil"
+	"net"
+	"net/http"
+	"strings"
+
+	"github.com/andybalholm/brotli"
+
+	"github.com/labstack/echo/v4"
+)
+
+type (
+	// BrotliConfig defines the config for Brotli middleware.
+	BrotliConfig struct {
+		// Skipper defines a function to skip middleware.
+		Skipper Skipper
+
+		// Brotli compression level.
+		// Optional. Default value -1.
+		Level int `yaml:"level"`
+	}
+
+	brotliResponseWriter struct {
+		io.Writer
+		http.ResponseWriter
+	}
+)
+
+const (
+	brotliScheme = "br"
+)
+
+var (
+	// DefaultBrotliConfig is the default Brotli middleware config.
+	DefaultBrotliConfig = BrotliConfig{
+		Skipper: DefaultSkipper,
+		Level:   brotli.DefaultCompression,
+	}
+)
+
+// Brotli returns a middleware which compresses HTTP response using brotli compression
+// scheme.
+func Brotli() echo.MiddlewareFunc {
+	return BrotliWithConfig(DefaultBrotliConfig)
+}
+
+// BrotliWithConfig return Brotli middleware with config.
+// See: `Brotli()`.
+func BrotliWithConfig(config BrotliConfig) echo.MiddlewareFunc {
+	// Defaults
+	if config.Skipper == nil {
+		config.Skipper = DefaultBrotliConfig.Skipper
+	}
+	if config.Level == 0 {
+		config.Level = DefaultBrotliConfig.Level
+	}
+
+	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), brotliScheme) {
+				res.Header().Set(echo.HeaderContentEncoding, brotliScheme) // Issue #806
+				rw := res.Writer
+
+				w := brotli.NewWriterOptions(rw, brotli.WriterOptions{Quality: config.Level})
+
+				defer func() {
+					if res.Size == 0 {
+						if res.Header().Get(echo.HeaderContentEncoding) == brotliScheme {
+							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()
+				}()
+				grw := &brotliResponseWriter{Writer: w, ResponseWriter: rw}
+				res.Writer = grw
+			}
+			return next(c)
+		}
+	}
+}
+
+func (w *brotliResponseWriter) 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 *brotliResponseWriter) 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 *brotliResponseWriter) Flush() {
+	w.Writer.(*brotli.Writer).Flush()
+	if flusher, ok := w.ResponseWriter.(http.Flusher); ok {
+		flusher.Flush()
+	}
+}
+
+func (w *brotliResponseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
+	return w.ResponseWriter.(http.Hijacker).Hijack()
+}

middleware/compress_brotli_test.go

@@ -0,0 +1,144 @@
+package middleware
+
+import (
+	"bytes"
+	"io"
+	"io/ioutil"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+
+	"github.com/andybalholm/brotli"
+
+	"github.com/labstack/echo/v4"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestBrotli(t *testing.T) {
+	e := echo.New()
+	req := httptest.NewRequest(http.MethodGet, "/", nil)
+	rec := httptest.NewRecorder()
+	c := e.NewContext(req, rec)
+
+	// Skip if no Accept-Encoding header
+	h := Brotli()(func(c echo.Context) error {
+		c.Response().Write([]byte("test")) // For Content-Type sniffing
+		return nil
+	})
+	h(c)
+
+	assert := assert.New(t)
+
+	assert.Equal("test", rec.Body.String())
+
+	// Brotli
+	req = httptest.NewRequest(http.MethodGet, "/", nil)
+	req.Header.Set(echo.HeaderAcceptEncoding, brotliScheme)
+	rec = httptest.NewRecorder()
+	c = e.NewContext(req, rec)
+	h(c)
+	assert.Equal(brotliScheme, rec.Header().Get(echo.HeaderContentEncoding))
+	assert.Contains(rec.Header().Get(echo.HeaderContentType), echo.MIMETextPlain)
+
+	r := brotli.NewReader(rec.Body)
+	buf := new(bytes.Buffer)
+	buf.ReadFrom(r)
+	assert.Equal("test", buf.String())
+
+	chunkBuf := make([]byte, 5)
+
+	// Brotli chunked
+	req = httptest.NewRequest(http.MethodGet, "/", nil)
+	req.Header.Set(echo.HeaderAcceptEncoding, brotliScheme)
+	rec = httptest.NewRecorder()
+
+	c = e.NewContext(req, rec)
+	Brotli()(func(c echo.Context) error {
+		c.Response().Header().Set("Content-Type", "text/event-stream")
+		c.Response().Header().Set("Transfer-Encoding", "chunked")
+
+		// Write and flush the first part of the data
+		c.Response().Write([]byte("test\n"))
+		c.Response().Flush()
+
+		// Read the first part of the data
+		assert.True(rec.Flushed)
+		assert.Equal(brotliScheme, rec.Header().Get(echo.HeaderContentEncoding))
+		r.Reset(rec.Body)
+
+		_, err := io.ReadFull(r, chunkBuf)
+		assert.NoError(err)
+		assert.Equal("test\n", string(chunkBuf))
+
+		// Write and flush the second part of the data
+		c.Response().Write([]byte("test\n"))
+		c.Response().Flush()
+
+		_, err = io.ReadFull(r, chunkBuf)
+		assert.NoError(err)
+		assert.Equal("test\n", string(chunkBuf))
+
+		// Write the final part of the data and return
+		c.Response().Write([]byte("test"))
+		return nil
+	})(c)
+
+	buf = new(bytes.Buffer)
+	buf.ReadFrom(r)
+	assert.Equal("test", buf.String())
+}
+
+func TestBrotliNoContent(t *testing.T) {
+	e := echo.New()
+	req := httptest.NewRequest(http.MethodGet, "/", nil)
+	req.Header.Set(echo.HeaderAcceptEncoding, brotliScheme)
+	rec := httptest.NewRecorder()
+	c := e.NewContext(req, rec)
+	h := Brotli()(func(c echo.Context) error {
+		return c.NoContent(http.StatusNoContent)
+	})
+	if assert.NoError(t, h(c)) {
+		assert.Empty(t, rec.Header().Get(echo.HeaderContentEncoding))
+		assert.Empty(t, rec.Header().Get(echo.HeaderContentType))
+		assert.Equal(t, 0, len(rec.Body.Bytes()))
+	}
+}
+
+func TestBrotliErrorReturned(t *testing.T) {
+	e := echo.New()
+	e.Use(Brotli())
+	e.GET("/", func(c echo.Context) error {
+		return echo.ErrNotFound
+	})
+	req := httptest.NewRequest(http.MethodGet, "/", nil)
+	req.Header.Set(echo.HeaderAcceptEncoding, brotliScheme)
+	rec := httptest.NewRecorder()
+	e.ServeHTTP(rec, req)
+	assert.Equal(t, http.StatusNotFound, rec.Code)
+	assert.Empty(t, rec.Header().Get(echo.HeaderContentEncoding))
+}
+
+// Issue #806
+func TestBrotliWithStatic(t *testing.T) {
+	e := echo.New()
+	e.Use(Brotli())
+	e.Static("/test", "../_fixture/images")
+	req := httptest.NewRequest(http.MethodGet, "/test/walle.png", nil)
+	req.Header.Set(echo.HeaderAcceptEncoding, brotliScheme)
+	rec := httptest.NewRecorder()
+	e.ServeHTTP(rec, req)
+	assert.Equal(t, http.StatusOK, rec.Code)
+	// Data is written out in chunks when Content-Length == "", so only
+	// validate the content length if it's not set.
+	if cl := rec.Header().Get("Content-Length"); cl != "" {
+		assert.Equal(t, cl, rec.Body.Len())
+	}
+	r := brotli.NewReader(rec.Body)
+
+	want, err := ioutil.ReadFile("../_fixture/images/walle.png")
+	if assert.NoError(t, err) {
+		buf := new(bytes.Buffer)
+		buf.ReadFrom(r)
+		assert.Equal(t, want, buf.Bytes())
+	}
+}