From 64a8b16ca2ac3729cdedf6d1a9dff85ed6b66418 Mon Sep 17 00:00:00 2001 From: Snider Date: Fri, 20 Feb 2026 23:37:46 +0000 Subject: [PATCH] feat: add WithBrotli response compression middleware Co-Authored-By: Virgil --- brotli.go | 120 ++++++++++++++++++++++++++++++++++++++++++++ brotli_test.go | 132 +++++++++++++++++++++++++++++++++++++++++++++++++ go.mod | 1 + go.sum | 4 ++ options.go | 13 +++++ 5 files changed, 270 insertions(+) create mode 100644 brotli.go create mode 100644 brotli_test.go diff --git a/brotli.go b/brotli.go new file mode 100644 index 0000000..b203cf2 --- /dev/null +++ b/brotli.go @@ -0,0 +1,120 @@ +// SPDX-License-Identifier: EUPL-1.2 + +package api + +import ( + "io" + "net/http" + "strconv" + "strings" + "sync" + + "github.com/andybalholm/brotli" + "github.com/gin-gonic/gin" +) + +const ( + // BrotliBestSpeed is the lowest (fastest) Brotli compression level. + BrotliBestSpeed = brotli.BestSpeed + + // BrotliBestCompression is the highest (smallest output) Brotli level. + BrotliBestCompression = brotli.BestCompression + + // BrotliDefaultCompression is the default Brotli compression level. + BrotliDefaultCompression = brotli.DefaultCompression +) + +// brotliHandler manages a pool of brotli writers for reuse across requests. +type brotliHandler struct { + pool sync.Pool + level int +} + +// newBrotliHandler creates a handler that pools brotli writers at the given level. +func newBrotliHandler(level int) *brotliHandler { + if level < BrotliBestSpeed || level > BrotliBestCompression { + level = BrotliDefaultCompression + } + return &brotliHandler{ + level: level, + pool: sync.Pool{ + New: func() any { + return brotli.NewWriterLevel(io.Discard, level) + }, + }, + } +} + +// Handle is the Gin middleware function that compresses responses with Brotli. +func (h *brotliHandler) Handle(c *gin.Context) { + if !strings.Contains(c.Request.Header.Get("Accept-Encoding"), "br") { + c.Next() + return + } + + w := h.pool.Get().(*brotli.Writer) + w.Reset(c.Writer) + + c.Header("Content-Encoding", "br") + c.Writer.Header().Add("Vary", "Accept-Encoding") + + bw := &brotliWriter{ResponseWriter: c.Writer, writer: w} + c.Writer = bw + + defer func() { + if bw.status >= http.StatusBadRequest { + bw.Header().Del("Content-Encoding") + bw.Header().Del("Vary") + w.Reset(io.Discard) + } else if c.Writer.Size() < 0 { + w.Reset(io.Discard) + } + _ = w.Close() + if c.Writer.Size() > -1 { + c.Header("Content-Length", strconv.Itoa(c.Writer.Size())) + } + h.pool.Put(w) + }() + + c.Next() +} + +// brotliWriter wraps gin.ResponseWriter to intercept writes through brotli. +type brotliWriter struct { + gin.ResponseWriter + writer *brotli.Writer + statusWritten bool + status int +} + +func (b *brotliWriter) Write(data []byte) (int, error) { + b.Header().Del("Content-Length") + + if !b.statusWritten { + b.status = b.ResponseWriter.Status() + } + + if b.status >= http.StatusBadRequest { + b.Header().Del("Content-Encoding") + b.Header().Del("Vary") + return b.ResponseWriter.Write(data) + } + + return b.writer.Write(data) +} + +func (b *brotliWriter) WriteString(s string) (int, error) { + return b.Write([]byte(s)) +} + +func (b *brotliWriter) WriteHeader(code int) { + b.status = code + b.statusWritten = true + b.Header().Del("Content-Length") + b.ResponseWriter.WriteHeader(code) +} + +func (b *brotliWriter) Flush() { + _ = b.writer.Flush() + b.ResponseWriter.Flush() +} diff --git a/brotli_test.go b/brotli_test.go new file mode 100644 index 0000000..45e8071 --- /dev/null +++ b/brotli_test.go @@ -0,0 +1,132 @@ +// SPDX-License-Identifier: EUPL-1.2 + +package api_test + +import ( + "net/http" + "net/http/httptest" + "testing" + + "github.com/gin-gonic/gin" + + api "forge.lthn.ai/core/go-api" +) + +// ── WithBrotli ──────────────────────────────────────────────────────── + +func TestWithBrotli_Good_CompressesResponse(t *testing.T) { + gin.SetMode(gin.TestMode) + e, _ := api.New(api.WithBrotli()) + e.Register(&stubGroup{}) + + h := e.Handler() + w := httptest.NewRecorder() + req, _ := http.NewRequest(http.MethodGet, "/stub/ping", nil) + req.Header.Set("Accept-Encoding", "br") + h.ServeHTTP(w, req) + + if w.Code != http.StatusOK { + t.Fatalf("expected 200, got %d", w.Code) + } + + ce := w.Header().Get("Content-Encoding") + if ce != "br" { + t.Fatalf("expected Content-Encoding=%q, got %q", "br", ce) + } +} + +func TestWithBrotli_Good_NoCompressionWithoutAcceptHeader(t *testing.T) { + gin.SetMode(gin.TestMode) + e, _ := api.New(api.WithBrotli()) + e.Register(&stubGroup{}) + + h := e.Handler() + w := httptest.NewRecorder() + req, _ := http.NewRequest(http.MethodGet, "/stub/ping", nil) + // Deliberately not setting Accept-Encoding header. + h.ServeHTTP(w, req) + + if w.Code != http.StatusOK { + t.Fatalf("expected 200, got %d", w.Code) + } + + ce := w.Header().Get("Content-Encoding") + if ce == "br" { + t.Fatal("expected no br Content-Encoding when client does not request it") + } +} + +func TestWithBrotli_Good_DefaultLevel(t *testing.T) { + // Calling WithBrotli() with no arguments should use default compression + // and not panic. + gin.SetMode(gin.TestMode) + e, _ := api.New(api.WithBrotli()) + e.Register(&stubGroup{}) + + h := e.Handler() + w := httptest.NewRecorder() + req, _ := http.NewRequest(http.MethodGet, "/stub/ping", nil) + req.Header.Set("Accept-Encoding", "br") + h.ServeHTTP(w, req) + + if w.Code != http.StatusOK { + t.Fatalf("expected 200, got %d", w.Code) + } + + ce := w.Header().Get("Content-Encoding") + if ce != "br" { + t.Fatalf("expected Content-Encoding=%q with default level, got %q", "br", ce) + } +} + +func TestWithBrotli_Good_CustomLevel(t *testing.T) { + // WithBrotli(BrotliBestSpeed) should work without panicking and still compress. + gin.SetMode(gin.TestMode) + e, _ := api.New(api.WithBrotli(api.BrotliBestSpeed)) + e.Register(&stubGroup{}) + + h := e.Handler() + w := httptest.NewRecorder() + req, _ := http.NewRequest(http.MethodGet, "/stub/ping", nil) + req.Header.Set("Accept-Encoding", "br") + h.ServeHTTP(w, req) + + if w.Code != http.StatusOK { + t.Fatalf("expected 200, got %d", w.Code) + } + + ce := w.Header().Get("Content-Encoding") + if ce != "br" { + t.Fatalf("expected Content-Encoding=%q with BestSpeed, got %q", "br", ce) + } +} + +func TestWithBrotli_Good_CombinesWithOtherMiddleware(t *testing.T) { + gin.SetMode(gin.TestMode) + e, _ := api.New( + api.WithBrotli(), + api.WithRequestID(), + ) + e.Register(&stubGroup{}) + + h := e.Handler() + w := httptest.NewRecorder() + req, _ := http.NewRequest(http.MethodGet, "/stub/ping", nil) + req.Header.Set("Accept-Encoding", "br") + h.ServeHTTP(w, req) + + if w.Code != http.StatusOK { + t.Fatalf("expected 200, got %d", w.Code) + } + + // Both brotli compression and request ID should be present. + ce := w.Header().Get("Content-Encoding") + if ce != "br" { + t.Fatalf("expected Content-Encoding=%q, got %q", "br", ce) + } + + rid := w.Header().Get("X-Request-ID") + if rid == "" { + t.Fatal("expected X-Request-ID header from WithRequestID") + } +} diff --git a/go.mod b/go.mod index 185512a..25780ff 100644 --- a/go.mod +++ b/go.mod @@ -3,6 +3,7 @@ module forge.lthn.ai/core/go-api go 1.25.5 require ( + github.com/andybalholm/brotli v1.2.0 github.com/coreos/go-oidc/v3 v3.17.0 github.com/gin-contrib/cors v1.7.6 github.com/gin-contrib/gzip v1.2.5 diff --git a/go.sum b/go.sum index 9b9a819..de6d468 100644 --- a/go.sum +++ b/go.sum @@ -4,6 +4,8 @@ github.com/PuerkitoBio/purell v1.1.1 h1:WEQqlqaGbrPkxLJWfBwQmfEAE1Z7ONdDLqrN38tN github.com/PuerkitoBio/purell v1.1.1/go.mod h1:c11w/QuzBsJSee3cPx9rAFu61PvFxuPbtSwDGJws/X0= github.com/PuerkitoBio/urlesc v0.0.0-20170810143723-de5bf2ad4578 h1:d+Bc7a5rLufV/sSk/8dngufqelfh6jnri85riMAaF/M= github.com/PuerkitoBio/urlesc v0.0.0-20170810143723-de5bf2ad4578/go.mod h1:uGdkoq3SwY9Y+13GIhn11/XLaGBb4BfwItxLd5jeuXE= +github.com/andybalholm/brotli v1.2.0 h1:ukwgCxwYrmACq68yiUqwIWnGY0cTPox/M94sVwToPjQ= +github.com/andybalholm/brotli v1.2.0/go.mod h1:rzTDkvFWvIrjDXZHkuS16NPggd91W3kUSvPlQ1pLaKY= github.com/bytedance/gopkg v0.1.3 h1:TPBSwH8RsouGCBcMBktLt1AymVo2TVsBVCY4b6TnZ/M= github.com/bytedance/gopkg v0.1.3/go.mod h1:576VvJ+eJgyCzdjS+c4+77QF3p7ubbtiKARP3TxducM= github.com/bytedance/sonic v1.14.2 h1:k1twIoe97C1DtYUo+fZQy865IuHia4PR5RPiuGPPIIE= @@ -124,6 +126,8 @@ github.com/twitchyliquid64/golang-asm v0.15.1 h1:SU5vSMR7hnwNxj24w34ZyCi/FmDZTkS github.com/twitchyliquid64/golang-asm v0.15.1/go.mod h1:a1lVb/DtPvCB8fslRZhAngC2+aY1QWCk3Cedj/Gdt08= github.com/ugorji/go/codec v1.3.1 h1:waO7eEiFDwidsBN6agj1vJQ4AG7lh2yqXyOXqhgQuyY= github.com/ugorji/go/codec v1.3.1/go.mod h1:pRBVtBSKl77K30Bv8R2P+cLSGaTtex6fsA2Wjqmfxj4= +github.com/xyproto/randomstring v1.0.5 h1:YtlWPoRdgMu3NZtP45drfy1GKoojuR7hmRcnhZqKjWU= +github.com/xyproto/randomstring v1.0.5/go.mod h1:rgmS5DeNXLivK7YprL0pY+lTuhNQW3iGxZ18UQApw/E= github.com/yuin/goldmark v1.4.13/go.mod h1:6yULJ656Px+3vBD8DxQVa3kxgyrAnzto9xy5taEt/CY= go.uber.org/mock v0.6.0 h1:hyF9dfmbgIX5EfOdasqLsWD6xqpNZlXblLB/Dbnwv3Y= go.uber.org/mock v0.6.0/go.mod h1:KiVJ4BqZJaMj4svdfmHM0AUx4NJYO8ZNpPnZn1Z+BBU= diff --git a/options.go b/options.go index 64a1957..634107d 100644 --- a/options.go +++ b/options.go @@ -145,6 +145,19 @@ func WithGzip(level ...int) Option { } } +// WithBrotli adds Brotli response compression middleware using andybalholm/brotli. +// An optional compression level may be supplied (e.g. BrotliBestSpeed, +// BrotliBestCompression). If omitted, BrotliDefaultCompression is used. +func WithBrotli(level ...int) Option { + return func(e *Engine) { + l := BrotliDefaultCompression + if len(level) > 0 { + l = level[0] + } + e.middlewares = append(e.middlewares, newBrotliHandler(l).Handle) + } +} + // WithSlog adds structured request logging middleware via gin-contrib/slog. // Each request is logged with method, path, status code, latency, and client IP. // If logger is nil, slog.Default() is used.