feat: Compress and extract slim binaries with zstd (#2533)

Fixes #2202

Co-authored-by: Dean Sheather <dean@deansheather.com>

site/site.go

@@ -1,13 +1,15 @@
 package site
 
 import (
+	"archive/tar"
 	"bytes"
 	"context"
-
+	"errors"
 	"fmt"
 	"io"
 	"io/fs"
 	"net/http"
+	"os"
 	"path"
 	"path/filepath"
 	"strings"
@@ -15,7 +17,9 @@ import (
 	"time"
 
 	"github.com/justinas/nosurf"
+	"github.com/klauspost/compress/zstd"
 	"github.com/unrolled/secure"
+	"golang.org/x/exp/slices"
 	"golang.org/x/xerrors"
 )
 
@@ -29,22 +33,25 @@ func WithAPIResponse(ctx context.Context, apiResponse APIResponse) context.Conte
 }
 
 // Handler returns an HTTP handler for serving the static site.
-func Handler(fileSystem fs.FS) http.Handler {
+func Handler(siteFS fs.FS, binFS http.FileSystem) http.Handler {
 	// html files are handled by a text/template. Non-html files
 	// are served by the default file server.
 	//
 	// REMARK: text/template is needed to inject values on each request like
 	//         CSRF.
-	files, err := htmlFiles(fileSystem)
-
+	files, err := htmlFiles(siteFS)
 	if err != nil {
 		panic(xerrors.Errorf("Failed to return handler for static files. Html files failed to load: %w", err))
 	}
 
+	mux := http.NewServeMux()
+	mux.Handle("/bin/", http.StripPrefix("/bin", http.FileServer(binFS)))
+	mux.Handle("/", http.FileServer(http.FS(siteFS))) // All other non-html static files.
+
 	return secureHeaders(&handler{
-		fs:        fileSystem,
+		fs:        siteFS,
 		htmlFiles: files,
-		h:         http.FileServer(http.FS(fileSystem)), // All other non-html static files
+		h:         mux,
 	})
 }
 
@@ -146,8 +153,13 @@ func (h *handler) ServeHTTP(resp http.ResponseWriter, req *http.Request) {
 		return
 	}
 
+	switch {
+	// If requesting binaries, serve straight up.
+	case reqFile == "bin" || strings.HasPrefix(reqFile, "bin/"):
+		h.h.ServeHTTP(resp, req)
+		return
 	// If the original file path exists we serve it.
-	if h.exists(reqFile) {
+	case h.exists(reqFile):
 		if ShouldCacheFile(reqFile) {
 			resp.Header().Add("Cache-Control", "public, max-age=31536000, immutable")
 		}
@@ -357,7 +369,6 @@ func htmlFiles(files fs.FS) (*htmlTemplates, error) {
 
 		return nil
 	})
-
 	if err != nil {
 		return nil, err
 	}
@@ -366,3 +377,135 @@ func htmlFiles(files fs.FS) (*htmlTemplates, error) {
 		tpls: root,
 	}, nil
 }
+
+// ExtractOrReadBinFS checks the provided fs for compressed coder
+// binaries and extracts them into dest/bin if found. As a fallback,
+// the provided FS is checked for a /bin directory, if it is non-empty
+// it is returned. Finally dest/bin is returned as a fallback allowing
+// binaries to be manually placed in dest (usually
+// ${CODER_CACHE_DIRECTORY}/site/bin).
+func ExtractOrReadBinFS(dest string, siteFS fs.FS) (http.FileSystem, error) {
+	if dest == "" {
+		// No destination on fs, embedded fs is the only option.
+		binFS, err := fs.Sub(siteFS, "bin")
+		if err != nil {
+			return nil, xerrors.Errorf("cache path is empty and embedded fs does not have /bin: %w", err)
+		}
+		return http.FS(binFS), nil
+	}
+
+	dest = filepath.Join(dest, "bin")
+	mkdest := func() (http.FileSystem, error) {
+		err := os.MkdirAll(dest, 0o700)
+		if err != nil {
+			return nil, xerrors.Errorf("mkdir failed: %w", err)
+		}
+		return http.Dir(dest), nil
+	}
+
+	archive, err := siteFS.Open("bin/coder.tar.zst")
+	if err != nil {
+		if xerrors.Is(err, fs.ErrNotExist) {
+			files, err := fs.ReadDir(siteFS, "bin")
+			if err != nil {
+				if xerrors.Is(err, fs.ErrNotExist) {
+					// Given fs does not have a bin directory,
+					// serve from cache directory.
+					return mkdest()
+				}
+				return nil, xerrors.Errorf("site fs read dir failed: %w", err)
+			}
+
+			if len(filterFiles(files, "GITKEEP")) > 0 {
+				// If there are other files than bin/GITKEEP,
+				// serve the files.
+				binFS, err := fs.Sub(siteFS, "bin")
+				if err != nil {
+					return nil, xerrors.Errorf("site fs sub dir failed: %w", err)
+				}
+				return http.FS(binFS), nil
+			}
+
+			// Nothing we can do, serve the cache directory,
+			// thus allowing binaries to be places there.
+			return mkdest()
+		}
+		return nil, xerrors.Errorf("open coder binary archive failed: %w", err)
+	}
+	defer archive.Close()
+
+	dir, err := mkdest()
+	if err != nil {
+		return nil, err
+	}
+
+	n, err := extractBin(dest, archive)
+	if err != nil {
+		return nil, xerrors.Errorf("extract coder binaries failed: %w", err)
+	}
+	if n == 0 {
+		return nil, xerrors.New("no files were extracted from coder binaries archive")
+	}
+
+	return dir, nil
+}
+
+func filterFiles(files []fs.DirEntry, names ...string) []fs.DirEntry {
+	var filtered []fs.DirEntry
+	for _, f := range files {
+		if slices.Contains(names, f.Name()) {
+			continue
+		}
+		filtered = append(filtered, f)
+	}
+	return filtered
+}
+
+func extractBin(dest string, r io.Reader) (numExtraced int, err error) {
+	opts := []zstd.DOption{
+		// Concurrency doesn't help us when decoding the tar and
+		// can actually slow us down.
+		zstd.WithDecoderConcurrency(1),
+		// Ignoring checksums can give a slight performance
+		// boost but it's probalby not worth the reduced safety.
+		zstd.IgnoreChecksum(false),
+		// Allow the decoder to use more memory giving us a 2-3x
+		// performance boost.
+		zstd.WithDecoderLowmem(false),
+	}
+	zr, err := zstd.NewReader(r, opts...)
+	if err != nil {
+		return 0, xerrors.Errorf("open zstd archive failed: %w", err)
+	}
+	defer zr.Close()
+
+	tr := tar.NewReader(zr)
+	n := 0
+	for {
+		h, err := tr.Next()
+		if err != nil {
+			if errors.Is(err, io.EOF) {
+				return n, nil
+			}
+			return n, xerrors.Errorf("read tar archive failed: %w", err)
+		}
+
+		name := filepath.Join(dest, filepath.Base(h.Name))
+		f, err := os.Create(name)
+		if err != nil {
+			return n, xerrors.Errorf("create file failed: %w", err)
+		}
+		//#nosec // We created this tar, no risk of decompression bomb.
+		_, err = io.Copy(f, tr)
+		if err != nil {
+			_ = f.Close()
+			return n, xerrors.Errorf("write file contents failed: %w", err)
+		}
+		err = f.Close()
+		if err != nil {
+			return n, xerrors.Errorf("close file failed: %w", err)
+		}
+
+		n++
+	}
+}

site/site_test.go

@@ -7,10 +7,12 @@ import (
 	"io"
 	"net/http"
 	"net/http/httptest"
+	"strings"
 	"testing"
 	"testing/fstest"
 	"time"
 
+	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 
 	"github.com/coder/coder/site"
@@ -36,8 +38,9 @@ func TestCaching(t *testing.T) {
 			Data: []byte("folderFile"),
 		},
 	}
+	binFS := http.FS(fstest.MapFS{})
 
-	srv := httptest.NewServer(site.Handler(rootFS))
+	srv := httptest.NewServer(site.Handler(rootFS, binFS))
 	defer srv.Close()
 
 	// Create a context
@@ -95,15 +98,16 @@ func TestServingFiles(t *testing.T) {
 			Data: []byte("dashboard-css-bytes"),
 		},
 	}
+	binFS := http.FS(fstest.MapFS{})
 
-	srv := httptest.NewServer(site.Handler(rootFS))
+	srv := httptest.NewServer(site.Handler(rootFS, binFS))
 	defer srv.Close()
 
 	// Create a context
 	ctx, cancelFunc := context.WithTimeout(context.Background(), 1*time.Second)
 	defer cancelFunc()
 
-	var testCases = []struct {
+	testCases := []struct {
 		path     string
 		expected string
 	}{
@@ -150,7 +154,7 @@ func TestServingFiles(t *testing.T) {
 func TestShouldCacheFile(t *testing.T) {
 	t.Parallel()
 
-	var testCases = []struct {
+	testCases := []struct {
 		reqFile  string
 		expected bool
 	}{
@@ -171,6 +175,167 @@ func TestShouldCacheFile(t *testing.T) {
 	}
 }
 
+func TestServingBin(t *testing.T) {
+	t.Parallel()
+
+	// Create a misc rootfs for realistic test.
+	rootFS := fstest.MapFS{
+		"index.html": &fstest.MapFile{
+			Data: []byte("index-bytes"),
+		},
+		"favicon.ico": &fstest.MapFile{
+			Data: []byte("favicon-bytes"),
+		},
+		"dashboard.js": &fstest.MapFile{
+			Data: []byte("dashboard-js-bytes"),
+		},
+		"dashboard.css": &fstest.MapFile{
+			Data: []byte("dashboard-css-bytes"),
+		},
+	}
+
+	type req struct {
+		url        string
+		wantStatus int
+		wantBody   []byte
+	}
+	tests := []struct {
+		name    string
+		fs      fstest.MapFS
+		reqs    []req
+		wantErr bool
+	}{
+		{
+			name: "Extract and serve bin",
+			fs: fstest.MapFS{
+				"bin/coder.tar.zst": &fstest.MapFile{
+					// echo iamcoder >coder-linux-amd64
+					// tar cf coder.tar coder-linux-amd64
+					// zstd --long --ultra -22 coder.tar
+					Data: []byte{
+						0x28, 0xb5, 0x2f, 0xfd, 0x64, 0x00, 0x27, 0xf5, 0x02, 0x00, 0x12, 0xc4,
+						0x0e, 0x16, 0xa0, 0xb5, 0x39, 0x00, 0xe8, 0x67, 0x59, 0xaf, 0xe3, 0xdd,
+						0x8d, 0xfe, 0x47, 0xe8, 0x9d, 0x9c, 0x44, 0x0b, 0x75, 0x70, 0x61, 0x52,
+						0x0d, 0x56, 0xaa, 0x16, 0xb9, 0x5a, 0x0a, 0x4b, 0x40, 0xd2, 0x7a, 0x05,
+						0xd1, 0xd7, 0xe3, 0xf9, 0xf9, 0x07, 0xef, 0xda, 0x77, 0x04, 0xff, 0xe8,
+						0x7a, 0x94, 0x56, 0x9a, 0x40, 0x3b, 0x94, 0x61, 0x18, 0x91, 0x90, 0x21,
+						0x0c, 0x00, 0xf3, 0xc5, 0xe5, 0xd8, 0x80, 0x10, 0x06, 0x0a, 0x08, 0x86,
+						0xb2, 0x00, 0x60, 0x12, 0x70, 0xd3, 0x51, 0x05, 0x04, 0x20, 0x16, 0x2c,
+						0x79, 0xad, 0x01, 0xc0, 0xf5, 0x28, 0x08, 0x03, 0x1c, 0x4c, 0x84, 0xf4,
+					},
+				},
+			},
+			reqs: []req{
+				{url: "/bin/coder-linux-amd64", wantStatus: http.StatusOK, wantBody: []byte("iamcoder\n")},
+				{url: "/bin/GITKEEP", wantStatus: http.StatusNotFound},
+			},
+		},
+		{
+			name: "Error on invalid archive",
+			fs: fstest.MapFS{
+				"bin/coder.tar.zst": &fstest.MapFile{
+					Data: []byte{
+						0x28, 0xb5, 0x2f, 0xfd, 0x64, 0x00, 0x27, 0xf5, 0x02, 0x00, 0x12, 0xc4,
+						0x0e, 0x16, 0xa0, 0xb5, 0x39, 0x00, 0xe8, 0x67, 0x59, 0xaf, 0xe3, 0xdd,
+						0x8d, 0xfe, 0x47, 0xe8, 0x9d, 0x9c, 0x44, 0x0b, 0x75, 0x70, 0x61, 0x52,
+						0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, // Zeroed from above test.
+						0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, // Zeroed from above test.
+						0x7a, 0x94, 0x56, 0x9a, 0x40, 0x3b, 0x94, 0x61, 0x18, 0x91, 0x90, 0x21,
+						0x0c, 0x00, 0xf3, 0xc5, 0xe5, 0xd8, 0x80, 0x10, 0x06, 0x0a, 0x08, 0x86,
+						0xb2, 0x00, 0x60, 0x12, 0x70, 0xd3, 0x51, 0x05, 0x04, 0x20, 0x16, 0x2c,
+						0x79, 0xad, 0x01, 0xc0, 0xf5, 0x28, 0x08, 0x03, 0x1c, 0x4c, 0x84, 0xf4,
+					},
+				},
+			},
+			wantErr: true,
+		},
+		{
+			name: "Error on empty archive",
+			fs: fstest.MapFS{
+				"bin/coder.tar.zst": &fstest.MapFile{Data: []byte{}},
+			},
+			wantErr: true,
+		},
+		{
+			name: "Serve local fs",
+			fs: fstest.MapFS{
+				// Only GITKEEP file on embedded fs, won't be served.
+				"bin/GITKEEP": &fstest.MapFile{},
+			},
+			reqs: []req{
+				{url: "/bin/coder-linux-amd64", wantStatus: http.StatusNotFound},
+				{url: "/bin/GITKEEP", wantStatus: http.StatusNotFound},
+			},
+		},
+		{
+			name: "Serve local fs when embedd fs empty",
+			fs:   fstest.MapFS{},
+			reqs: []req{
+				{url: "/bin/coder-linux-amd64", wantStatus: http.StatusNotFound},
+				{url: "/bin/GITKEEP", wantStatus: http.StatusNotFound},
+			},
+		},
+		{
+			name: "Serve embedd fs",
+			fs: fstest.MapFS{
+				"bin/GITKEEP": &fstest.MapFile{
+					Data: []byte(""),
+				},
+				"bin/coder-linux-amd64": &fstest.MapFile{
+					Data: []byte("embedd"),
+				},
+			},
+			reqs: []req{
+				{url: "/bin/coder-linux-amd64", wantStatus: http.StatusOK, wantBody: []byte("embedd")},
+				{url: "/bin/GITKEEP", wantStatus: http.StatusOK, wantBody: []byte("")},
+			},
+		},
+	}
+	//nolint // Parallel test detection issue.
+	for _, tt := range tests {
+		tt := tt
+		t.Run(tt.name, func(t *testing.T) {
+			t.Parallel()
+
+			dest := t.TempDir()
+			binFS, err := site.ExtractOrReadBinFS(dest, tt.fs)
+			if !tt.wantErr && err != nil {
+				require.NoError(t, err, "extract or read failed")
+			} else if tt.wantErr {
+				require.Error(t, err, "extraction or read did not fail")
+			}
+
+			srv := httptest.NewServer(site.Handler(rootFS, binFS))
+			defer srv.Close()
+
+			// Create a context
+			ctx, cancelFunc := context.WithTimeout(context.Background(), 1*time.Second)
+			defer cancelFunc()
+
+			for _, tr := range tt.reqs {
+				t.Run(strings.TrimPrefix(tr.url, "/"), func(t *testing.T) {
+					req, err := http.NewRequestWithContext(ctx, "GET", srv.URL+tr.url, nil)
+					require.NoError(t, err, "http request failed")
+
+					resp, err := http.DefaultClient.Do(req)
+					require.NoError(t, err, "http do failed")
+					defer resp.Body.Close()
+
+					gotStatus := resp.StatusCode
+					gotBody, _ := io.ReadAll(resp.Body)
+
+					if tr.wantStatus > 0 {
+						assert.Equal(t, tr.wantStatus, gotStatus, "status did not match")
+					}
+					if tr.wantBody != nil {
+						assert.Equal(t, string(tr.wantBody), string(gotBody), "body did not match")
+					}
+				})
+			}
+		})
+	}
+}
+
 func TestServeAPIResponse(t *testing.T) {
 	t.Parallel()
 
@@ -180,6 +345,7 @@ func TestServeAPIResponse(t *testing.T) {
 			Data: []byte(`{"code":{{ .APIResponse.StatusCode }},"message":"{{ .APIResponse.Message }}"}`),
 		},
 	}
+	binFS := http.FS(fstest.MapFS{})
 
 	apiResponse := site.APIResponse{
 		StatusCode: http.StatusBadGateway,
@@ -187,7 +353,7 @@ func TestServeAPIResponse(t *testing.T) {
 	}
 	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		r = r.WithContext(site.WithAPIResponse(r.Context(), apiResponse))
-		site.Handler(rootFS).ServeHTTP(w, r)
+		site.Handler(rootFS, binFS).ServeHTTP(w, r)
 	}))
 	defer srv.Close()