diff --git a/gzip.go b/gzip.go index d2f616f..0bd54b0 100644 --- a/gzip.go +++ b/gzip.go @@ -335,8 +335,14 @@ func (h *handler) pool() *sync.Pool { } func (h *handler) shouldGzip(r *http.Request) bool { - if h.forceGzip { - return true + if h.config.shouldGzip != nil { + switch h.config.shouldGzip(r) { + case NegotiateGzip: + case SkipGzip: + return false + case ForceGzip: + return true + } } match := httputils.Negotiate(r.Header, "Accept-Encoding", "gzip") @@ -419,7 +425,7 @@ type config struct { level int minSize int contentTypes []string - forceGzip bool + shouldGzip func(*http.Request) ShouldGzipType } // Option customizes the behaviour of the gzip handler. @@ -482,21 +488,40 @@ func ContentTypes(types []string) Option { } } -// ForceGzip makes the handler always return a gzipped -// response and not consult the request's Accept-Encoding +// ShouldGzip provides control over when the handler should +// return a gzipped response. It allows handlers to implement +// logic that doesn't consult the request's Accept-Encoding // header. // // By default, responses are only gzipped if the request's // Accept-Encoding header indicates gzip support. // -// Note: ForceGzip does not affect MinSize or ContentTypes, -// it simply ignores the Accept-Encoding header. -var ForceGzip Option = forceGzip - -func forceGzip(c *config) { - c.forceGzip = true +// Note: ShouldGzip does not affect MinSize or ContentTypes, +// it simply provides control over negotiating gzip support. +func ShouldGzip(fn func(*http.Request) ShouldGzipType) Option { + return func(c *config) { + c.shouldGzip = fn + } } +// ShouldGzipType controls how the handler determines gzip +// support. +type ShouldGzipType int + +const ( + // NegotiateGzip defers gzip negotiation to the + // request's Accept-Encoding header. + NegotiateGzip ShouldGzipType = iota + + // SkipGzip skips gzipping the response. + SkipGzip + + // ForceGzip ignores the request's Accept-Encoding + // header and always gzips the response. + // (See ShouldGzip note). + ForceGzip +) + type ( // Each of these structs is intentionally small (1 pointer wide) so // as to fit inside an interface{} without causing an allocaction. diff --git a/gzip_test.go b/gzip_test.go index 8a712c2..1895383 100644 --- a/gzip_test.go +++ b/gzip_test.go @@ -681,18 +681,43 @@ func TestWrapper(t *testing.T) { assert.Equal(t, Gzip(handler, MinSize(42)), Wrapper(MinSize(42))(handler)) } -func TestForceGzip(t *testing.T) { - handler := newTestHandler(testBody, ForceGzip) +func TestShouldGzip(t *testing.T) { + for _, tc := range []struct { + shouldGzip ShouldGzipType + advertise bool + expect bool + }{ + {NegotiateGzip, false, false}, + {NegotiateGzip, true, true}, + {SkipGzip, false, false}, + {SkipGzip, true, false}, + {ForceGzip, false, true}, + {ForceGzip, true, true}, + } { + handler := newTestHandler(testBody, ShouldGzip(func(*http.Request) ShouldGzipType { + return tc.shouldGzip + })) - req := httptest.NewRequest(http.MethodGet, "/whatever", nil) - resp := httptest.NewRecorder() - handler.ServeHTTP(resp, req) + req := httptest.NewRequest(http.MethodGet, "/whatever", nil) + if tc.advertise { + req.Header.Set("Accept-Encoding", "gzip") + } - res := resp.Result() - assert.Equal(t, http.StatusOK, res.StatusCode) - assert.Equal(t, "gzip", res.Header.Get("Content-Encoding")) - assert.Equal(t, "Accept-Encoding", res.Header.Get("Vary")) - assert.Equal(t, gzipStrLevel(testBody, DefaultCompression), resp.Body.Bytes()) + resp := httptest.NewRecorder() + handler.ServeHTTP(resp, req) + + res := resp.Result() + assert.Equal(t, http.StatusOK, res.StatusCode, "%+v", tc) + assert.Equal(t, "Accept-Encoding", res.Header.Get("Vary"), "%+v", tc) + + if tc.expect { + assert.Equal(t, "gzip", res.Header.Get("Content-Encoding"), "%+v", tc) + assert.Equal(t, gzipStrLevel(testBody, DefaultCompression), resp.Body.Bytes(), "%+v", tc) + } else { + assert.Equal(t, "", res.Header.Get("Content-Encoding"), "%+v", tc) + assert.Equal(t, testBody, resp.Body.String(), "%+v", tc) + } + } } // --------------------------------------------------------------------