Skip to content

Commit

Permalink
Auto add 'Vary' header after compression (#1585)
Browse files Browse the repository at this point in the history
* 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'
  • Loading branch information
AutumnSun1996 authored Jul 2, 2023
1 parent d229959 commit 0d0bbfe
Show file tree
Hide file tree
Showing 5 changed files with 219 additions and 0 deletions.
12 changes: 12 additions & 0 deletions header.go
Original file line number Diff line number Diff line change
Expand Up @@ -344,6 +344,18 @@ func (h *ResponseHeader) SetContentEncodingBytes(contentEncoding []byte) {
h.contentEncoding = append(h.contentEncoding[:0], contentEncoding...)
}

// addVaryBytes add value to the 'Vary' header if it's not included
func (h *ResponseHeader) addVaryBytes(value []byte) {
v := h.peek(strVary)
if len(v) == 0 {
// 'Vary' is not set
h.SetBytesV(HeaderVary, value)
} else if !bytes.Contains(v, value) {
// 'Vary' is set and not contains target value
h.SetBytesV(HeaderVary, append(append(v, ','), value...))
} // else: 'Vary' is set and contains target value
}

// Server returns Server header value.
func (h *ResponseHeader) Server() []byte {
return h.server
Expand Down
62 changes: 62 additions & 0 deletions header_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -3007,3 +3007,65 @@ func TestResponseHeader_Keys(t *testing.T) {
t.Fatalf("Unexpected value %q. Expected %q", actualTrailerKeys, expectedTrailerKeys)
}
}

func TestAddVaryHeader(t *testing.T) {
t.Parallel()

var h ResponseHeader

h.addVaryBytes([]byte("Accept-Encoding"))
got := string(h.Peek("Vary"))
expected := "Accept-Encoding"
if got != expected {
t.Errorf("expected %q got %q", expected, got)
}

var buf bytes.Buffer
h.WriteTo(&buf) //nolint:errcheck

if n := strings.Count(buf.String(), "Vary: "); n != 1 {
t.Errorf("Vary occurred %d times", n)
}
}

func TestAddVaryHeaderExisting(t *testing.T) {
t.Parallel()

var h ResponseHeader

h.Set("Vary", "Accept")
h.addVaryBytes([]byte("Accept-Encoding"))
got := string(h.Peek("Vary"))
expected := "Accept,Accept-Encoding"
if got != expected {
t.Errorf("expected %q got %q", expected, got)
}

var buf bytes.Buffer
h.WriteTo(&buf) //nolint:errcheck

if n := strings.Count(buf.String(), "Vary: "); n != 1 {
t.Errorf("Vary occurred %d times", n)
}
}

func TestAddVaryHeaderExistingAcceptEncoding(t *testing.T) {
t.Parallel()

var h ResponseHeader

h.Set("Vary", "Accept-Encoding")
h.addVaryBytes([]byte("Accept-Encoding"))
got := string(h.Peek("Vary"))
expected := "Accept-Encoding"
if got != expected {
t.Errorf("expected %q got %q", expected, got)
}

var buf bytes.Buffer
h.WriteTo(&buf) //nolint:errcheck

if n := strings.Count(buf.String(), "Vary: "); n != 1 {
t.Errorf("Vary occurred %d times", n)
}
}
3 changes: 3 additions & 0 deletions http.go
Original file line number Diff line number Diff line change
Expand Up @@ -1723,6 +1723,7 @@ func (resp *Response) brotliBody(level int) error {
resp.bodyRaw = nil
}
resp.Header.SetContentEncodingBytes(strBr)
resp.Header.addVaryBytes(strAcceptEncoding)
return nil
}

Expand Down Expand Up @@ -1778,6 +1779,7 @@ func (resp *Response) gzipBody(level int) error {
resp.bodyRaw = nil
}
resp.Header.SetContentEncodingBytes(strGzip)
resp.Header.addVaryBytes(strAcceptEncoding)
return nil
}

Expand Down Expand Up @@ -1833,6 +1835,7 @@ func (resp *Response) deflateBody(level int) error {
resp.bodyRaw = nil
}
resp.Header.SetContentEncodingBytes(strDeflate)
resp.Header.addVaryBytes(strAcceptEncoding)
return nil
}

Expand Down
141 changes: 141 additions & 0 deletions server_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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()

Expand Down
1 change: 1 addition & 0 deletions strings.go
Original file line number Diff line number Diff line change
Expand Up @@ -57,6 +57,7 @@ var (
strProxyAuthenticate = []byte(HeaderProxyAuthenticate)
strProxyAuthorization = []byte(HeaderProxyAuthorization)
strWWWAuthenticate = []byte(HeaderWWWAuthenticate)
strVary = []byte(HeaderVary)

strCookieExpires = []byte("expires")
strCookieDomain = []byte("domain")
Expand Down

0 comments on commit 0d0bbfe

Please sign in to comment.