Auto add 'Vary' header after compression (#1585)

* Auto add 'Vary' header after compression

Add config `SetAddVaryHeaderForCompression` to enable
'Vary: Accept-Encoding' header when compression is used.

* feat: always set the Vary header

* create and use `ResponseHeader.AddVaryBytes`

* not export 'AddVaryBytes'

1 files changed, 141 insertions, 0 deletions

diff --git a/server_test.go b/server_test.go
index e22da4e..67ca8a6 100644
--- a/server_test.go
+++ b/server_test.go
@@ -2035,6 +2035,147 @@ func TestCompressHandler(t *testing.T) {
 	}
 }
 
+func TestCompressHandlerVary(t *testing.T) {
+	t.Parallel()
+
+	expectedBody := string(createFixedBody(2e4))
+
+	h := CompressHandlerBrotliLevel(func(ctx *RequestCtx) {
+		ctx.WriteString(expectedBody) //nolint:errcheck
+	}, CompressBrotliBestSpeed, CompressBestSpeed)
+
+	var ctx RequestCtx
+	var resp Response
+
+	// verify uncompressed response
+	h(&ctx)
+	s := ctx.Response.String()
+	br := bufio.NewReader(bytes.NewBufferString(s))
+	if err := resp.Read(br); err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+	ce := resp.Header.ContentEncoding()
+	if string(ce) != "" {
+		t.Fatalf("unexpected Content-Encoding: %q. Expecting %q", ce, "")
+	}
+	vary := resp.Header.Peek("Vary")
+	if string(vary) != "" {
+		t.Fatalf("unexpected Vary: %q. Expecting %q", vary, "")
+	}
+	body := resp.Body()
+	if string(body) != expectedBody {
+		t.Fatalf("unexpected body %q. Expecting %q", body, expectedBody)
+	}
+
+	// verify gzip-compressed response
+	ctx.Request.Reset()
+	ctx.Response.Reset()
+	ctx.Request.Header.Set("Accept-Encoding", "gzip, deflate, sdhc")
+
+	h(&ctx)
+	s = ctx.Response.String()
+	br = bufio.NewReader(bytes.NewBufferString(s))
+	if err := resp.Read(br); err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+	ce = resp.Header.ContentEncoding()
+	if string(ce) != "gzip" {
+		t.Fatalf("unexpected Content-Encoding: %q. Expecting %q", ce, "gzip")
+	}
+	vary = resp.Header.Peek("Vary")
+	if string(vary) != "Accept-Encoding" {
+		t.Fatalf("unexpected Vary: %q. Expecting %q", vary, "Accept-Encoding")
+	}
+	body, err := resp.BodyGunzip()
+	if err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+	if string(body) != expectedBody {
+		t.Fatalf("unexpected body %q. Expecting %q", body, expectedBody)
+	}
+
+	// an attempt to compress already compressed response
+	ctx.Request.Reset()
+	ctx.Response.Reset()
+	ctx.Request.Header.Set("Accept-Encoding", "gzip, deflate, sdhc")
+	hh := CompressHandler(h)
+	hh(&ctx)
+	s = ctx.Response.String()
+	br = bufio.NewReader(bytes.NewBufferString(s))
+	if err := resp.Read(br); err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+	ce = resp.Header.ContentEncoding()
+	if string(ce) != "gzip" {
+		t.Fatalf("unexpected Content-Encoding: %q. Expecting %q", ce, "gzip")
+	}
+	vary = resp.Header.Peek("Vary")
+	if string(vary) != "Accept-Encoding" {
+		t.Fatalf("unexpected Vary: %q. Expecting %q", vary, "Accept-Encoding")
+	}
+	body, err = resp.BodyGunzip()
+	if err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+	if string(body) != expectedBody {
+		t.Fatalf("unexpected body %q. Expecting %q", body, expectedBody)
+	}
+
+	// verify deflate-compressed response
+	ctx.Request.Reset()
+	ctx.Response.Reset()
+	ctx.Request.Header.Set(HeaderAcceptEncoding, "foobar, deflate, sdhc")
+
+	h(&ctx)
+	s = ctx.Response.String()
+	br = bufio.NewReader(bytes.NewBufferString(s))
+	if err := resp.Read(br); err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+	ce = resp.Header.ContentEncoding()
+	if string(ce) != "deflate" {
+		t.Fatalf("unexpected Content-Encoding: %q. Expecting %q", ce, "deflate")
+	}
+	vary = resp.Header.Peek("Vary")
+	if string(vary) != "Accept-Encoding" {
+		t.Fatalf("unexpected Vary: %q. Expecting %q", vary, "Accept-Encoding")
+	}
+	body, err = resp.BodyInflate()
+	if err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+	if string(body) != expectedBody {
+		t.Fatalf("unexpected body %q. Expecting %q", body, expectedBody)
+	}
+
+	// verify br-compressed response
+	ctx.Request.Reset()
+	ctx.Response.Reset()
+	ctx.Request.Header.Set(HeaderAcceptEncoding, "gzip, deflate, br")
+
+	h(&ctx)
+	s = ctx.Response.String()
+	br = bufio.NewReader(bytes.NewBufferString(s))
+	if err := resp.Read(br); err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+	ce = resp.Header.ContentEncoding()
+	if string(ce) != "br" {
+		t.Fatalf("unexpected Content-Encoding: %q. Expecting %q", ce, "br")
+	}
+	vary = resp.Header.Peek("Vary")
+	if string(vary) != "Accept-Encoding" {
+		t.Fatalf("unexpected Vary: %q. Expecting %q", vary, "Accept-Encoding")
+	}
+	body, err = resp.BodyUnbrotli()
+	if err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+	if string(body) != expectedBody {
+		t.Fatalf("unexpected body %q. Expecting %q", body, expectedBody)
+	}
+}
+
 func TestRequestCtxWriteString(t *testing.T) {
 	t.Parallel()