Improve encoding decision

Especially when content-encoding is already set

middleware.go

@@ -48,25 +48,25 @@ func newMiddleware(h http.Handler, options ...Option) *middleware {
 }
 
 func (m *middleware) ServeHTTP(w http.ResponseWriter, r *http.Request) {
+	// for cache validation
+	w.Header().Add("Vary", "Accept-Encoding")
+
 	encoding := nego.NegotiateContentEncoding(r, m.supportedEncoding...)
 	enc, ok := m.encoders[encoding]
-	if !ok {
-		m.Handler.ServeHTTP(w, r)
-		return
-	}
-
-	mw := &middlewareWriter{
-		ResponseWriter: w,
-		ctx:            r.Context(),
-		factory:        enc.factory,
-		encoding:       encoding,
-		m:              m,
-		status:         http.StatusOK,
+	if ok {
+		mw := &middlewareWriter{
+			ResponseWriter: w,
+			ctx:            r.Context(),
+			factory:        enc.factory,
+			encoding:       encoding,
+			m:              m,
+			status:         http.StatusOK,
+		}
+		defer mw.flush()
+		w = mw
 	}
-	defer mw.flush()
-
-	m.Handler.ServeHTTP(mw, r)
 
+	m.Handler.ServeHTTP(w, r)
 }
 
 func (m *middleware) populateSupportedEncoding() {

writer.go

@@ -45,30 +45,34 @@ func (mw *middlewareWriter) WriteHeader(status int) {
 	mw.status = status
 	mw.headerSent = true
 
-	if cenc := mw.Header().Get("content-encoding"); cenc != "" {
+	cenc := mw.Header().Get("content-encoding")
+	ctype := mw.Header().Get("content-type")
+
+	// if content encoding already defined
+	// or content type is not defined
+	// or content type is not in allowed list
+	// => just forward the body
+	if cenc != "" || ctype == "" || !matchRegexes(ctype, mw.m.allowedType) {
+		mw.dontEncode = true
 		mw.ResponseWriter.WriteHeader(status)
 		return
 	}
 
-	if ctype := mw.Header().Get("content-type"); ctype != "" {
-		if !matchRegexes(ctype, mw.m.allowedType) {
-			mw.dontEncode = true
-			mw.ResponseWriter.WriteHeader(status)
-			return
-		}
-	}
-
+	// if content length is big enough, start encoding now
 	if clen := mw.Header().Get("content-length"); clen != "" {
 		len, err := strconv.ParseUint(clen, 10, 64)
 		if err == nil {
-			if len > mw.m.minSize {
+			if len >= mw.m.minSize {
 				mw.startEncoding()
-				return
+			} else {
+				mw.dontEncode = true
+				mw.ResponseWriter.WriteHeader(status)
 			}
+			return
 		}
 	}
 
-	// buffer the content until
+	// the content length is unknown, buffer the response until it exceeds minSize
 	mw.buff = make([]byte, mw.m.minSize)
 }
 
@@ -119,7 +123,7 @@ func (mw *middlewareWriter) Write(chunk []byte) (int, error) {
 		}
 		return mw.ResponseWriter.Write(chunk)
 	}
-	n := copy(mw.buff[mw.buffLen:], chunk[:])
+	n := copy(mw.buff[mw.buffLen:], chunk)
 	mw.buffLen = newBufLen
 	return n, nil
 }