tsweb: add transparent compression for StdHandlerraggi/tsweb-compression

Implements inline compression for both gzip and brotli via the brotli
library. The library requires that Content-Type is set. The
implementation here explicitly avoids wrapping the ResponseWriter in
cases where Accept-Encoding is not set so as to maximally attempt to get
out of the way of hijack and upgrade concerns. It also avoids any
attempt at compression if Content-Encoding is already set so that
handlers that already perform compression are unaffected.

Signed-off-by: James Tucker <james@tailscale.com>

4 files changed, 246 insertions, 1 deletions

diff --git a/cmd/tailscaled/depaware.txt b/cmd/tailscaled/depaware.txt
index 051bdc0bd..8e8bfd6d1 100644
--- a/cmd/tailscaled/depaware.txt
+++ b/cmd/tailscaled/depaware.txt
@@ -8,6 +8,7 @@ tailscale.com/cmd/tailscaled dependencies: (generated by github.com/tailscale/de
    W 💣 github.com/alexbrainman/sspi                                 from github.com/alexbrainman/sspi/internal/common+
    W    github.com/alexbrainman/sspi/internal/common                 from github.com/alexbrainman/sspi/negotiate
    W 💣 github.com/alexbrainman/sspi/negotiate                       from tailscale.com/net/tshttpproxy
+        github.com/andybalholm/brotli                                from tailscale.com/tsweb
   LD    github.com/anmitsu/go-shlex                                  from tailscale.com/tempfork/gliderlabs/ssh
    L    github.com/aws/aws-sdk-go-v2                                 from github.com/aws/aws-sdk-go-v2/internal/ini
    L    github.com/aws/aws-sdk-go-v2/aws                             from github.com/aws/aws-sdk-go-v2/aws/middleware+
diff --git a/tsweb/compress.go b/tsweb/compress.go
new file mode 100644
index 000000000..ee6f4564e
--- /dev/null
+++ b/tsweb/compress.go
@@ -0,0 +1,104 @@
+// 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 (
+	"bufio"
+	"errors"
+	"io"
+	"net"
+	"net/http"
+
+	"github.com/andybalholm/brotli"
+)
+
+type compressingHandler struct {
+	h http.Handler
+}
+
+func (h compressingHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
+	if !AcceptsEncoding(r, "br") && !AcceptsEncoding(r, "gzip") {
+		h.h.ServeHTTP(w, r)
+		return
+	}
+
+	cw := &compressingResponseWriter{
+		ResponseWriter: w,
+		r:              r,
+	}
+	defer cw.Close()
+
+	h.h.ServeHTTP(cw, r)
+}
+
+type compressingResponseWriter struct {
+	http.ResponseWriter
+	r *http.Request
+	w io.Writer
+}
+
+// WriteHeader implements http.ResponseWriter.
+func (w *compressingResponseWriter) WriteHeader(code int) {
+	// If a handler has already set a Content-Encoding, such as for precompressed
+	// assets, skip the compressing writer. This must be recorded before
+	// WriteHeader call as "The header map is cleared when 2xx-5xx headers are
+	// sent".
+	if w.w == nil {
+		if w.ResponseWriter.Header().Get("Content-Encoding") == "" {
+			w.w = brotli.HTTPCompressor(w.ResponseWriter, w.r)
+		} else {
+			w.w = w.ResponseWriter
+		}
+	}
+	w.ResponseWriter.WriteHeader(code)
+}
+
+// Write implements http.ResponseWriter.
+func (w *compressingResponseWriter) Write(b []byte) (int, error) {
+	if w.w == nil {
+		w.WriteHeader(http.StatusOK)
+	}
+	return w.w.Write(b)
+}
+
+// Close implements io.Closer.
+func (w *compressingResponseWriter) Close() error {
+	if w.w == nil {
+		return nil
+	}
+	if c, ok := w.w.(io.Closer); ok {
+		return c.Close()
+	}
+	return nil
+}
+
+// flusher is an interface that is implemented by gzip.Writer and other writers
+// that differs from http.Flusher in that it may return an error.
+type flusher interface {
+	Flush() error
+}
+
+// Flush implements http.Flusher.
+func (w *compressingResponseWriter) Flush() {
+	// the writer may implement either of the flusher interfaces, so try both.
+	if f, ok := w.w.(flusher); ok {
+		_ = f.Flush()
+	}
+	if f, ok := w.w.(http.Flusher); ok {
+		f.Flush()
+	}
+
+	if f, ok := w.ResponseWriter.(http.Flusher); ok {
+		f.Flush()
+	}
+}
+
+// Hijack implements http.Hijacker.
+func (w *compressingResponseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
+	if hj, ok := w.ResponseWriter.(http.Hijacker); ok {
+		return hj.Hijack()
+	}
+	return nil, nil, errors.New("ResponseWriter is not a Hijacker")
+}
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")
+	})
+
+}
diff --git a/tsweb/tsweb.go b/tsweb/tsweb.go
index 8b52e656d..ac6a14332 100644
--- a/tsweb/tsweb.go
+++ b/tsweb/tsweb.go
@@ -239,7 +239,7 @@ func StdHandler(h ReturnHandler, opts HandlerOptions) http.Handler {
 	if opts.Logf == nil {
 		opts.Logf = logger.Discard
 	}
-	return retHandler{h, opts}
+	return compressingHandler{retHandler{h, opts}}
 }
 
 // retHandler is an http.Handler that wraps a Handler and handles errors.