Negotiate gzip on the response content type

README.md

@@ -206,6 +206,18 @@ Sets headers (passed as key:value) to requests. I.e. `rest.Headers("Server:MySer
 Compresses response with gzip. Adds `Vary: Accept-Encoding` to every response it handles, compressed or not,
 so shared caches key on the encoding rather than serving gzip bytes to a client that never asked for them.
 
+The decision is made on the **response** content type, either the one the handler set or, when it set none,
+the type sniffed from the first chunk of the body. By default the common textual types are compressed
+(`text/html`, `text/plain`, `text/css`, `text/xml`, `text/javascript`, `application/javascript`,
+`application/x-javascript`, `application/json`); pass your own list to override, i.e. `rest.Gzip("text/html")`.
+
+`Accept-Encoding` is parsed rather than substring-matched, so `gzip;q=0` is honoured as a refusal and a named
+`gzip` entry outranks a `*` wildcard. Compression is skipped for responses that carry no body (204 and 304),
+for responses the handler already encoded (`Content-Encoding` set), and for partial responses (206 or a
+`Content-Range`), whose offsets describe the uncompressed representation. `Content-Length` is dropped when the
+body is compressed, interim 1xx responses pass through without becoming the final status, and `Flush` and
+`Hijack` pass through so streaming responses and protocol upgrades keep working.
+
 ### RealIP middleware
 
 RealIP is a middleware that sets a http.Request's RemoteAddr to the results of parsing various headers that contain the client's real IP address. It checks headers in the following priority order:

gzip.go

@@ -1,9 +1,13 @@
 package rest
 
 import (
+	"bufio"
 	"compress/gzip"
+	"fmt"
 	"io"
+	"net"
 	"net/http"
+	"strconv"
 	"strings"
 	"sync"
 )
@@ -23,69 +27,220 @@ var gzPool = sync.Pool{
 	New: func() any { return gzip.NewWriter(io.Discard) },
 }
 
+// gzipResponseWriter defers the compression decision until the response content type is known,
+// either from the header the handler set or sniffed from the first chunk of the body.
 type gzipResponseWriter struct {
-	io.Writer
 	http.ResponseWriter
+
+	gzCts       []string
+	gz          *gzip.Writer
+	status      int
+	statusSet   bool
+	decided     bool
+	wroteHeader bool
+	hijacked    bool
 }
 
 func (w *gzipResponseWriter) WriteHeader(status int) {
-	w.Header().Del("Content-Length")
-	w.ResponseWriter.WriteHeader(status)
+	// 1xx are interim responses, they pass straight through and the final status still follows.
+	// 101 is the exception, it hands the connection to another protocol and is final.
+	if status >= 100 && status < 200 && status != http.StatusSwitchingProtocols {
+		w.ResponseWriter.WriteHeader(status)
+		return
+	}
+	if w.statusSet {
+		return // net/http keeps the first final status, so later calls are ignored here too
+	}
+	w.statusSet = true
+	w.status = status
+
+	// with a content type in hand the decision can be made right away, otherwise it waits for the
+	// first Write so the body can be sniffed
+	if ctype := w.Header().Get("Content-Type"); ctype != "" {
+		w.decide(ctype)
+		w.commit()
+	}
 }
 
 func (w *gzipResponseWriter) Write(b []byte) (int, error) {
-	return w.Writer.Write(b)
+	if !w.decided {
+		ctype := w.Header().Get("Content-Type")
+		// net/http suppresses sniffing for an already encoded body, guessing a type from
+		// compressed bytes would only mislabel it
+		if ctype == "" && w.Header().Get("Content-Encoding") == "" {
+			ctype = http.DetectContentType(b)
+			w.Header().Set("Content-Type", ctype)
+		}
+		w.decide(ctype)
+	}
+	if !w.wroteHeader {
+		w.commit()
+	}
+	if w.gz != nil {
+		return w.gz.Write(b)
+	}
+	return w.ResponseWriter.Write(b)
+}
+
+// decide turns compression on if the response content type is one of the configured types
+func (w *gzipResponseWriter) decide(ctype string) {
+	w.decided = true
+
+	switch w.status {
+	case http.StatusSwitchingProtocols, http.StatusNoContent, http.StatusResetContent, http.StatusNotModified:
+		return // these carry no body to compress
+	}
+	if w.Header().Get("Content-Encoding") != "" {
+		return // the handler encoded the body itself, wrapping it again would mislabel the result
+	}
+	if w.status == http.StatusPartialContent || w.Header().Get("Content-Range") != "" {
+		return // the range metadata describes the identity representation
+	}
+
+	for _, c := range w.gzCts {
+		if !strings.HasPrefix(strings.ToLower(ctype), strings.ToLower(c)) {
+			continue
+		}
+		gz := gzPool.Get().(*gzip.Writer)
+		gz.Reset(w.ResponseWriter)
+		w.gz = gz
+		w.Header().Set("Content-Encoding", "gzip")
+		w.Header().Del("Content-Length") // the handler's length describes the uncompressed body
+		return
+	}
+}
+
+func (w *gzipResponseWriter) commit() {
+	w.wroteHeader = true
+	if w.status == 0 {
+		w.status = http.StatusOK
+	}
+	w.ResponseWriter.WriteHeader(w.status)
+}
+
+// close finishes the gzip stream and makes sure the status reaches the client even if the handler
+// wrote no body at all. finished reports whether the handler returned normally: when it panicked
+// instead, an uncommitted response is left alone so a recoverer upstream can still make it a 500.
+func (w *gzipResponseWriter) close(finished bool) {
+	if w.hijacked {
+		return // the handler owns the connection now, nothing may be written to it
+	}
+	if w.wroteHeader || finished {
+		if !w.decided {
+			w.decide(w.Header().Get("Content-Type"))
+		}
+		if !w.wroteHeader {
+			w.commit()
+		}
+	}
+	if w.gz == nil {
+		return
+	}
+	_ = w.gz.Close()
+	gzPool.Put(w.gz)
+	w.gz = nil
+}
+
+// Flush pushes buffered data out, keeping streaming responses working through the compressor
+func (w *gzipResponseWriter) Flush() {
+	if w.hijacked {
+		return
+	}
+	// decide before the headers leave, otherwise a later write could be compressed after the client
+	// was already told the body is identity
+	if !w.decided {
+		w.decide(w.Header().Get("Content-Type"))
+	}
+	if !w.wroteHeader {
+		w.commit()
+	}
+	if w.gz != nil {
+		_ = w.gz.Flush()
+	}
+	if f, ok := w.ResponseWriter.(http.Flusher); ok {
+		f.Flush()
+	}
 }
 
-// Gzip is a middleware compressing response
+// Hijack passes through to the underlying writer for protocol upgrades
+func (w *gzipResponseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
+	h, ok := w.ResponseWriter.(http.Hijacker)
+	if !ok {
+		return nil, nil, fmt.Errorf("http.Hijacker not supported")
+	}
+	conn, rw, err := h.Hijack()
+	if err == nil {
+		w.hijacked = true
+	}
+	return conn, rw, err
+}
+
+// Unwrap exposes the underlying writer to http.ResponseController
+func (w *gzipResponseWriter) Unwrap() http.ResponseWriter {
+	return w.ResponseWriter
+}
+
+// Gzip is a middleware compressing response. The decision is made on the response content type,
+// so it applies to what the handler actually produced. Content types default to the common textual
+// ones and can be overridden by the caller.
 func Gzip(contentTypes ...string) func(http.Handler) http.Handler {
 
 	gzCts := gzDefaultContentTypes
 	if len(contentTypes) > 0 {
 		gzCts = contentTypes
 	}
 
-	contentType := func(r *http.Request) string {
-		result := r.Header.Get("Content-type")
-		if result == "" {
-			return "application/octet-stream"
-		}
-		return result
-	}
-
-	f := func(next http.Handler) http.Handler {
+	return func(next http.Handler) http.Handler {
 		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 			// the representation depends on Accept-Encoding, caches must key on it even when not compressing
 			w.Header().Add("Vary", "Accept-Encoding")
 
-			if !strings.Contains(r.Header.Get("Accept-Encoding"), "gzip") {
+			if !acceptsGzip(r.Header.Values("Accept-Encoding")) {
 				next.ServeHTTP(w, r)
 				return
 			}
 
-			var gzOk bool
-			ctype := contentType(r)
-			for _, c := range gzCts {
-				if strings.HasPrefix(strings.ToLower(ctype), strings.ToLower(c)) {
-					gzOk = true
-					break
-				}
-			}
+			gw := &gzipResponseWriter{ResponseWriter: w, gzCts: gzCts}
+			finished := false
+			defer func() { gw.close(finished) }()
 
-			if !gzOk {
-				next.ServeHTTP(w, r)
-				return
-			}
-
-			w.Header().Set("Content-Encoding", "gzip")
-			gz := gzPool.Get().(*gzip.Writer)
-			defer gzPool.Put(gz)
+			next.ServeHTTP(gw, r)
+			finished = true
+		})
+	}
+}
 
-			gz.Reset(w)
-			defer gz.Close()
+// acceptsGzip reports whether the client accepts gzip, honoring an explicit q=0 rejection.
+// A named gzip entry decides the answer on its own, as it takes precedence over the "*" wildcard.
+// Repeated Accept-Encoding fields form a single list, so every field is examined.
+func acceptsGzip(headers []string) bool {
+	var wildcard, wildcardSeen bool
+
+	for _, header := range headers {
+		for enc := range strings.SplitSeq(header, ",") {
+			name, params, _ := strings.Cut(strings.TrimSpace(enc), ";")
+			n := strings.ToLower(strings.TrimSpace(name))
+			if n != "gzip" && n != "*" {
+				continue
+			}
+			acceptable := !rejectedByQuality(params)
+			if n == "gzip" {
+				return acceptable
+			}
+			if !wildcardSeen {
+				wildcard, wildcardSeen = acceptable, true
+			}
+		}
+	}
+	return wildcard
+}
 
-			next.ServeHTTP(&gzipResponseWriter{ResponseWriter: w, Writer: gz}, r)
-		})
+// rejectedByQuality reports whether the parameters of an Accept-Encoding entry carry q=0
+func rejectedByQuality(params string) bool {
+	q, ok := strings.CutPrefix(strings.ToLower(strings.TrimSpace(params)), "q=")
+	if !ok {
+		return false
 	}
-	return f
+	v, err := strconv.ParseFloat(strings.TrimSpace(q), 64)
+	return err == nil && v == 0
 }