gzhttp: Fix missing content type on Close (#883)

If compression had not yet been triggered in Write, be sure to detect content type on Close. Fixes #882
klauspost · Nov 10, 2023 · 80ba129 · 80ba129
1 parent 6a59dbb
commit 80ba129
Show file tree

Hide file tree

Showing 2 changed files with 76 additions and 1 deletion.
diff --git a/gzhttp/compress.go b/gzhttp/compress.go
@@ -335,7 +335,16 @@ func (w *GzipResponseWriter) Close() error {
 			ce = w.Header().Get(contentEncoding)
 			cr = w.Header().Get(contentRange)
 		)
-		// fmt.Println(len(w.buf) == 0, len(w.buf) < w.minSize, len(w.Header()[HeaderNoCompression]) != 0, ce != "", cr != "", !w.contentTypeFilter(ct))
+		if ct == "" {
+			ct = http.DetectContentType(w.buf)
+
+			// Handles the intended case of setting a nil Content-Type (as for http/server or http/fs)
+			// Set the header only if the key does not exist
+			if _, ok := w.Header()[contentType]; w.setContentType && !ok {
+				w.Header().Set(contentType, ct)
+			}
+		}
+
 		if len(w.buf) == 0 || len(w.buf) < w.minSize || len(w.Header()[HeaderNoCompression]) != 0 || ce != "" || cr != "" || !w.contentTypeFilter(ct) {
 			// GZIP not triggered, write out regular response.
 			return w.startPlain()

diff --git a/gzhttp/compress_test.go b/gzhttp/compress_test.go
@@ -14,6 +14,7 @@ import (
 	"net/url"
 	"os"
 	"strconv"
+	"strings"
 	"testing"
 
 	"github.com/klauspost/compress/gzip"
@@ -1883,3 +1884,68 @@ func Test1xxResponses(t *testing.T) {
 	body, _ := io.ReadAll(res.Body)
 	assertEqual(t, gzipStrLevel(testBody, gzip.DefaultCompression), body)
 }
+
+func TestContentTypeDetectWithJitter(t *testing.T) {
+	t.Parallel()
+	handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		content := `<!DOCTYPE html>` + strings.Repeat("foo", 400)
+		w.Write([]byte(content))
+	})
+
+	for _, tc := range []struct {
+		name    string
+		wrapper func(http.Handler) (http.Handler, error)
+	}{
+		{
+			name: "no wrapping",
+			wrapper: func(h http.Handler) (http.Handler, error) {
+				return h, nil
+			},
+		},
+		{
+			name: "default",
+			wrapper: func(h http.Handler) (http.Handler, error) {
+				wrapper, err := NewWrapper()
+				if err != nil {
+					return nil, err
+				}
+				return wrapper(h), nil
+			},
+		},
+		{
+			name: "jitter, default buffer",
+			wrapper: func(h http.Handler) (http.Handler, error) {
+				wrapper, err := NewWrapper(RandomJitter(32, 0, false))
+				if err != nil {
+					return nil, err
+				}
+				return wrapper(h), nil
+			},
+		},
+		{
+			name: "jitter, small buffer",
+			wrapper: func(h http.Handler) (http.Handler, error) {
+				wrapper, err := NewWrapper(RandomJitter(32, DefaultMinSize, false))
+				if err != nil {
+					return nil, err
+				}
+				return wrapper(h), nil
+			},
+		},
+	} {
+		tc := tc
+		t.Run(tc.name, func(t *testing.T) {
+			t.Parallel()
+
+			handler, err := tc.wrapper(handler)
+			assertNil(t, err)
+
+			req, resp := httptest.NewRequest(http.MethodGet, "/", nil), httptest.NewRecorder()
+			req.Header.Add("Accept-Encoding", "gzip")
+
+			handler.ServeHTTP(resp, req)
+
+			assertEqual(t, "text/html; charset=utf-8", resp.Header().Get("Content-Type"))
+		})
+	}
+}