cmd/tsconnect,util/precompress: move precompression to its own package

We have very similar code in corp, moving it to util/precompress allows
it to be reused.

Updates #5133

Signed-off-by: Mihai Parparita <mihai@tailscale.com>

cmd/tsconnect/build.go

@@ -5,21 +5,16 @@
 package main
 
 import (
-	"bytes"
-	"compress/gzip"
 	"encoding/json"
 	"fmt"
-	"io"
-	"io/fs"
 	"io/ioutil"
 	"log"
 	"os"
 	"path"
 	"path/filepath"
 
-	"github.com/andybalholm/brotli"
 	esbuild "github.com/evanw/esbuild/pkg/api"
-	"golang.org/x/sync/errgroup"
+	"tailscale.com/util/precompress"
 )
 
 func runBuild() {
@@ -127,77 +122,10 @@ func cleanDist() error {
 
 func precompressDist(fastCompression bool) error {
 	log.Printf("Pre-compressing files in %s/...\n", *distDir)
-	var eg errgroup.Group
-	err := fs.WalkDir(os.DirFS(*distDir), ".", func(p string, d fs.DirEntry, err error) error {
-		if err != nil {
-			return err
-		}
-		if d.IsDir() {
-			return nil
-		}
-		if !compressibleExtensions[filepath.Ext(p)] {
-			return nil
-		}
-		p = path.Join(*distDir, p)
-		log.Printf("Pre-compressing %v\n", p)
-
-		eg.Go(func() error {
-			return precompress(p, fastCompression)
-		})
-		return nil
+	return precompress.PrecompressDir(*distDir, precompress.Options{
+		FastCompression: fastCompression,
+		ProgressFn: func(path string) {
+			log.Printf("Pre-compressing %v\n", path)
+		},
 	})
-	if err != nil {
-		return err
-	}
-	return eg.Wait()
-}
-
-var compressibleExtensions = map[string]bool{
-	".js":   true,
-	".css":  true,
-	".wasm": true,
-}
-
-func precompress(path string, fastCompression bool) error {
-	contents, err := os.ReadFile(path)
-	if err != nil {
-		return err
-	}
-	fi, err := os.Lstat(path)
-	if err != nil {
-		return err
-	}
-
-	gzipLevel := gzip.BestCompression
-	if fastCompression {
-		gzipLevel = gzip.BestSpeed
-	}
-	err = writeCompressed(contents, func(w io.Writer) (io.WriteCloser, error) {
-		return gzip.NewWriterLevel(w, gzipLevel)
-	}, path+".gz", fi.Mode())
-	if err != nil {
-		return err
-	}
-	brotliLevel := brotli.BestCompression
-	if fastCompression {
-		brotliLevel = brotli.BestSpeed
-	}
-	return writeCompressed(contents, func(w io.Writer) (io.WriteCloser, error) {
-		return brotli.NewWriterLevel(w, brotliLevel), nil
-	}, path+".br", fi.Mode())
-}
-
-func writeCompressed(contents []byte, compressedWriterCreator func(io.Writer) (io.WriteCloser, error), outputPath string, outputMode fs.FileMode) error {
-	var buf bytes.Buffer
-	compressedWriter, err := compressedWriterCreator(&buf)
-	if err != nil {
-		return err
-	}
-	if _, err := compressedWriter.Write(contents); err != nil {
-		return err
-	}
-	if err := compressedWriter.Close(); err != nil {
-		return err
-	}
-	return os.WriteFile(outputPath, buf.Bytes(), outputMode)
 }

cmd/tsconnect/serve.go

@@ -19,6 +19,7 @@ import (
 	"time"
 
 	"tailscale.com/tsweb"
+	"tailscale.com/util/precompress"
 )
 
 //go:embed index.html
@@ -120,28 +121,10 @@ var entryPointsToDefaultDistPaths = map[string]string{
 
 func handleServeDist(w http.ResponseWriter, r *http.Request, distFS fs.FS) {
 	path := r.URL.Path
-	var f fs.File
-	// Prefer pre-compressed versions generated during the build step.
-	if tsweb.AcceptsEncoding(r, "br") {
-		if brotliFile, err := distFS.Open(path + ".br"); err == nil {
-			f = brotliFile
-			w.Header().Set("Content-Encoding", "br")
-		}
-	}
-	if f == nil && tsweb.AcceptsEncoding(r, "gzip") {
-		if gzipFile, err := distFS.Open(path + ".gz"); err == nil {
-			f = gzipFile
-			w.Header().Set("Content-Encoding", "gzip")
-		}
-	}
-
-	if f == nil {
-		if rawFile, err := distFS.Open(path); err == nil {
-			f = rawFile
-		} else {
-			http.Error(w, err.Error(), http.StatusNotFound)
-			return
-		}
+	f, err := precompress.OpenPrecompressedFile(w, r, path, distFS)
+	if err != nil {
+		http.Error(w, err.Error(), http.StatusNotFound)
+		return
 	}
 	defer f.Close()
 

util/precompress/precompress.go

@@ -0,0 +1,131 @@
+// 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 precompress provides build- and serving-time support for
+// precompressed static resources, to avoid the cost of repeatedly compressing
+// unchanging resources.
+package precompress
+
+import (
+	"bytes"
+	"compress/gzip"
+	"io"
+	"io/fs"
+	"net/http"
+	"os"
+	"path"
+	"path/filepath"
+
+	"github.com/andybalholm/brotli"
+	"golang.org/x/sync/errgroup"
+	"tailscale.com/tsweb"
+)
+
+// PrecompressDir compresses static assets in dirPath using Gzip and Brotli, so
+// that they can be later served with OpenPrecompressedFile.
+func PrecompressDir(dirPath string, options Options) error {
+	var eg errgroup.Group
+	err := fs.WalkDir(os.DirFS(dirPath), ".", func(p string, d fs.DirEntry, err error) error {
+		if err != nil {
+			return err
+		}
+		if d.IsDir() {
+			return nil
+		}
+		if !compressibleExtensions[filepath.Ext(p)] {
+			return nil
+		}
+		p = path.Join(dirPath, p)
+		if options.ProgressFn != nil {
+			options.ProgressFn(p)
+		}
+
+		eg.Go(func() error {
+			return precompress(p, options)
+		})
+		return nil
+	})
+	if err != nil {
+		return err
+	}
+	return eg.Wait()
+}
+
+type Options struct {
+	// FastCompression controls whether compression should be optimized for
+	// speed rather than size.
+	FastCompression bool
+	// ProgressFn, if non-nil, is invoked when a file in the directory is about
+	// to be compressed.
+	ProgressFn func(path string)
+}
+
+// OpenPrecompressedFile opens a file from fs, preferring compressed versions
+// generated by PrecompressDir if possible.
+func OpenPrecompressedFile(w http.ResponseWriter, r *http.Request, path string, fs fs.FS) (fs.File, error) {
+	if tsweb.AcceptsEncoding(r, "br") {
+		if f, err := fs.Open(path + ".br"); err == nil {
+			w.Header().Set("Content-Encoding", "br")
+			return f, nil
+		}
+	}
+	if tsweb.AcceptsEncoding(r, "gzip") {
+		if f, err := fs.Open(path + ".gz"); err == nil {
+			w.Header().Set("Content-Encoding", "gzip")
+			return f, nil
+		}
+	}
+
+	return fs.Open(path)
+}
+
+var compressibleExtensions = map[string]bool{
+	".js":   true,
+	".css":  true,
+	".wasm": true,
+}
+
+func precompress(path string, options Options) error {
+	contents, err := os.ReadFile(path)
+	if err != nil {
+		return err
+	}
+	fi, err := os.Lstat(path)
+	if err != nil {
+		return err
+	}
+
+	gzipLevel := gzip.BestCompression
+	if options.FastCompression {
+		gzipLevel = gzip.BestSpeed
+	}
+	err = writeCompressed(contents, func(w io.Writer) (io.WriteCloser, error) {
+		return gzip.NewWriterLevel(w, gzipLevel)
+	}, path+".gz", fi.Mode())
+	if err != nil {
+		return err
+	}
+	brotliLevel := brotli.BestCompression
+	if options.FastCompression {
+		brotliLevel = brotli.BestSpeed
+	}
+	return writeCompressed(contents, func(w io.Writer) (io.WriteCloser, error) {
+		return brotli.NewWriterLevel(w, brotliLevel), nil
+	}, path+".br", fi.Mode())
+}
+
+func writeCompressed(contents []byte, compressedWriterCreator func(io.Writer) (io.WriteCloser, error), outputPath string, outputMode fs.FileMode) error {
+	var buf bytes.Buffer
+	compressedWriter, err := compressedWriterCreator(&buf)
+	if err != nil {
+		return err
+	}
+	if _, err := compressedWriter.Write(contents); err != nil {
+		return err
+	}
+	if err := compressedWriter.Close(); err != nil {
+		return err
+	}
+	return os.WriteFile(outputPath, buf.Bytes(), outputMode)
+}