tailscale/tsweb/compress_test.go

141 lines
3.7 KiB
Go

// Copyright (c) 2022 Tailscale Inc & AUTHORS All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
package tsweb
import (
"compress/gzip"
"io"
"io/ioutil"
"net/http"
"net/http/httptest"
"testing"
"github.com/andybalholm/brotli"
)
func TestCompressingHandler(t *testing.T) {
h := compressingHandler{nil}
var _ http.Handler = h
w := &compressingResponseWriter{}
var (
_ http.ResponseWriter = w
_ http.Flusher = w
_ http.Hijacker = w
)
// testRequest constructs a response recorder and a compressing handler that
// wraps the given handler h, it calls the handler with r, and returns the
// response recorder. If r is nil, then a GET request is made to "/" with no
// additional headers.
testRequest := func(r *http.Request, h http.HandlerFunc) *httptest.ResponseRecorder {
t.Helper()
w := httptest.NewRecorder()
if r == nil {
r = httptest.NewRequest("GET", "/", nil)
}
compressingHandler{h}.ServeHTTP(w, r)
return w
}
checkBody := func(r io.Reader, want string) {
t.Helper()
body, err := ioutil.ReadAll(r)
if err != nil {
t.Fatal(err)
}
if string(body) != want {
t.Errorf("got body %q, want %q", body, want)
}
}
checkHeader := func(h http.Header, key, want string) {
t.Helper()
if got := h.Get(key); got != want {
t.Errorf("got header %q=%q, want %q", key, got, want)
}
}
t.Run("transparently compresses content with brotli", func(t *testing.T) {
r := httptest.NewRequest("GET", "/", nil)
r.Header.Set("Accept-Encoding", "br")
w := testRequest(r, func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "text/plain")
w.Write([]byte("hello world"))
})
checkHeader(w.Header(), "Content-Encoding", "br")
checkBody(brotli.NewReader(w.Body), "hello world")
})
t.Run("transparently compresses content with gzip", func(t *testing.T) {
r := httptest.NewRequest("GET", "/", nil)
r.Header.Set("Accept-Encoding", "gzip")
w := testRequest(r, func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "text/plain")
w.Write([]byte("hello world"))
})
checkHeader(w.Header(), "Content-Encoding", "gzip")
br, err := gzip.NewReader(w.Body)
if err != nil {
t.Fatal(err)
}
checkBody(br, "hello world")
})
t.Run("does not compress content if client does not accept compressed content", func(t *testing.T) {
w := testRequest(nil, func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "text/plain")
w.Write([]byte("hello world"))
})
checkHeader(w.Header(), "Content-Encoding", "")
checkBody(w.Body, "hello world")
})
t.Run("does not recompress content if client accepts compressed content but content is already compressed", func(t *testing.T) {
r := httptest.NewRequest("GET", "/", nil)
r.Header.Set("Accept-Encoding", "br")
w := testRequest(r, func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "text/plain")
w.Header().Set("Content-Encoding", "magic")
w.Write([]byte("hello world"))
})
checkHeader(w.Header(), "Content-Encoding", "magic")
checkBody(w.Body, "hello world")
})
t.Run("integration", func(t *testing.T) {
s := httptest.NewServer(compressingHandler{http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "text/plain")
w.Write([]byte("hello world"))
})})
defer s.Close()
r, err := http.NewRequest("GET", s.URL, nil)
if err != nil {
t.Fatal(err)
}
r.Header.Set("Accept-Encoding", "gzip")
res, err := s.Client().Do(r)
if err != nil {
t.Fatal(err)
}
checkHeader(res.Header, "Content-Encoding", "gzip")
br, err := gzip.NewReader(res.Body)
if err != nil {
t.Fatal(err)
}
checkBody(br, "hello world")
})
}