From be7253c084416b9598b17a050a68466c6bdf55ac Mon Sep 17 00:00:00 2001
From: Aliaksandr Valialkin <valyala@gmail.com>
Date: Fri, 22 May 2020 16:44:09 +0300
Subject: [PATCH] lib/httpserver: do not recompress already compressed response

This shoud help with vmauth issue - https://github.com/VictoriaMetrics/VictoriaMetrics/issues/514
---
 lib/httpserver/httpserver.go | 29 +++++++++++++++++++----------
 1 file changed, 19 insertions(+), 10 deletions(-)

diff --git a/lib/httpserver/httpserver.go b/lib/httpserver/httpserver.go
index 6d7a9f397..3b5accbff 100644
--- a/lib/httpserver/httpserver.go
+++ b/lib/httpserver/httpserver.go
@@ -232,10 +232,10 @@ func maybeGzipResponseWriter(w http.ResponseWriter, r *http.Request) http.Respon
 	ae = strings.ToLower(ae)
 	n := strings.Index(ae, "gzip")
 	if n < 0 {
+		// Do not apply gzip encoding to the response.
 		return w
 	}
-	h := w.Header()
-	h.Set("Content-Encoding", "gzip")
+	// Apply gzip encoding to the response.
 	zw := getGzipWriter(w)
 	bw := getBufioWriter(zw)
 	zrw := &gzipResponseWriter{
@@ -250,7 +250,14 @@ func maybeGzipResponseWriter(w http.ResponseWriter, r *http.Request) http.Respon
 //
 // The function must be called before the first w.Write* call.
 func DisableResponseCompression(w http.ResponseWriter) {
-	w.Header().Del("Content-Encoding")
+	zrw, ok := w.(*gzipResponseWriter)
+	if !ok {
+		return
+	}
+	if zrw.firstWriteDone {
+		logger.Panicf("BUG: DisableResponseCompression must be called before sending the response")
+	}
+	zrw.disableCompression = true
 }
 
 // EnableCORS enables https://developer.mozilla.org/en-US/docs/Web/HTTP/CORS
@@ -292,14 +299,16 @@ type gzipResponseWriter struct {
 func (zrw *gzipResponseWriter) Write(p []byte) (int, error) {
 	if !zrw.firstWriteDone {
 		h := zrw.Header()
-		if h.Get("Content-Encoding") != "gzip" {
-			// The request handler disabled gzip encoding.
-			// Send uncompressed response body.
+		if h.Get("Content-Encoding") != "" {
 			zrw.disableCompression = true
-		} else if h.Get("Content-Type") == "" {
-			// Disable auto-detection of content-type, since it
-			// is incorrectly detected after the compression.
-			h.Set("Content-Type", "text/html")
+		}
+		if !zrw.disableCompression {
+			h.Set("Content-Encoding", "gzip")
+			if h.Get("Content-Type") == "" {
+				// Disable auto-detection of content-type, since it
+				// is incorrectly detected after the compression.
+				h.Set("Content-Type", "text/html")
+			}
 		}
 		zrw.firstWriteDone = true
 	}