feat: Add support for extract and serve of zstd bins

coder · mafredri · Jun 21, 2022 · Jun 20, 2022 · Jun 20, 2022 · Jun 20, 2022
commit fb24c029549aa1ba011522eb21064fcb2f17d534
diff --git a/cli/server.go b/cli/server.go
@@ -254,6 +254,7 @@ func server() *cobra.Command {
 				Logger:               logger.Named("coderd"),
 				Database:             databasefake.New(),
 				Pubsub:               database.NewPubsubInMemory(),
+				CacheDir:             cacheDir,
 				GoogleTokenValidator: googleTokenValidator,
 				SecureAuthCookie:     secureAuthCookie,
 				SSHKeygenAlgorithm:   sshKeygenAlgorithm,

diff --git a/coderd/coderd.go b/coderd/coderd.go
@@ -7,6 +7,7 @@ import (
 	"io"
 	"net/http"
 	"net/url"
+	"path/filepath"
 	"sync"
 	"time"
 
@@ -42,6 +43,9 @@ type Options struct {
 	Database  database.Store
 	Pubsub    database.Pubsub
 
+	// CacheDir is used for caching files served by the API.
+	CacheDir string
+
 	AgentConnectionUpdateFrequency time.Duration
 	// APIRateLimit is the minutely throughput rate limit per user or ip.
 	// Setting a rate limit <0 will disable the rate limiter across the entire
@@ -78,11 +82,20 @@ func New(options *Options) *API {
 		}
 	}
 
+	siteCacheDir := options.CacheDir
+	if siteCacheDir != "" {
+		siteCacheDir = filepath.Join(siteCacheDir, "site")
+	}
+	binFS, err := site.ExtractOrReadBinFS(siteCacheDir, site.FS())
+	if err != nil {
+		panic(xerrors.Errorf("read site bin failed: %w", err))
+	}
+
 	r := chi.NewRouter()
 	api := &API{
 		Options:     options,
 		Handler:     r,
-		siteHandler: site.Handler(site.FS()),
+		siteHandler: site.Handler(site.FS(), binFS),
 	}
 	api.workspaceAgentCache = wsconncache.New(api.dialWorkspaceAgent, 0)
 

@@ -1,21 +1,25 @@
 package site
 
 import (
+	"archive/tar"
 	"bytes"
 	"context"
-
+	"errors"
 	"fmt"
 	"io"
 	"io/fs"
 	"net/http"
+	"os"
 	"path"
 	"path/filepath"
 	"strings"
 	"text/template" // html/template escapes some nonces
 	"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,134 @@ 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 destionation on fs, embedd fs is the only option.
+		binFS, err := fs.Sub(siteFS, "bin")
+		if err != nil {
+			return nil, xerrors.Errorf("cache path is empty and embedd 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) (n 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 give us a slight performance
+		// increase, we assume no corruption due to embedding.
+		zstd.IgnoreChecksum(true),
+		// 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)
+	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++
+	}
+}