1 files changed, 140 insertions, 0 deletions
diff --git a/tsweb/compress_test.go b/tsweb/compress_test.go
new file mode 100644
index 000000000..f1c8b6e89
--- /dev/null
+++ b/tsweb/compress_test.go
@@ -0,0 +1,140 @@
+// 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")
+	})
+
+}