server: push serveUIAssets into package ui

Better separate concerns by making package ui responsible for providing an HTTP handler that can serve its assets. As a side effect of the refactor, fix rendering of the "no UI installed" page in short binaries, which was broken in #25195. This is not pure code movement, so ui.Handler should be reviewed as if it were new code.
cockroachdb · Sep 7, 2018 · b498b9c · b498b9c
1 parent 5e6181f
commit b498b9c
Show file tree

Hide file tree

Showing 3 changed files with 88 additions and 74 deletions.
diff --git a/pkg/build/info.go b/pkg/build/info.go
@@ -92,6 +92,15 @@ func (b Info) Short() string {
 		b.Distribution, b.Tag, plat, b.Time, b.GoVersion)
 }
 
+// GoTime parses the utcTime string and returns a time.Time.
+func (b Info) GoTime() time.Time {
+	val, err := time.Parse(TimeFormat, b.Time)
+	if err != nil {
+		return time.Time{}
+	}
+	return val
+}
+
 // Timestamp parses the utcTime string and returns the number of seconds since epoch.
 func (b Info) Timestamp() (int64, error) {
 	val, err := time.Parse(TimeFormat, b.Time)

diff --git a/pkg/server/server.go b/pkg/server/server.go
@@ -18,9 +18,7 @@ import (
 	"compress/gzip"
 	"context"
 	"crypto/tls"
-	"encoding/json"
 	"fmt"
-	"html/template"
 	"io"
 	"io/ioutil"
 	"math"
@@ -33,7 +31,6 @@ import (
 	"sync/atomic"
 	"time"
 
-	assetfs "github.com/elazarl/go-bindata-assetfs"
 	raven "github.com/getsentry/raven-go"
 	gwruntime "github.com/grpc-ecosystem/grpc-gateway/runtime"
 	opentracing "github.com/opentracing/opentracing-go"
@@ -42,7 +39,6 @@ import (
 
 	"github.com/cockroachdb/cmux"
 	"github.com/cockroachdb/cockroach/pkg/base"
-	"github.com/cockroachdb/cockroach/pkg/build"
 	"github.com/cockroachdb/cockroach/pkg/gossip"
 	"github.com/cockroachdb/cockroach/pkg/internal/client"
 	"github.com/cockroachdb/cockroach/pkg/jobs"
@@ -1234,17 +1230,19 @@ func (s *Server) Start(ctx context.Context) error {
 	// endpoints.
 	s.mux.Handle(debug.Endpoint, debug.NewServer(s.st))
 
-	fileServer := http.FileServer(&assetfs.AssetFS{
-		Asset:     ui.Asset,
-		AssetDir:  ui.AssetDir,
-		AssetInfo: ui.AssetInfo,
-	})
-
 	// Serve UI assets. This needs to be before the gRPC handlers are registered, otherwise
 	// the `s.mux.Handle("/", ...)` would cover all URLs, allowing anonymous access.
 	maybeAuthMux := newAuthenticationMuxAllowAnonymous(
-		s.authentication, serveUIAssets(fileServer, s.cfg),
-	)
+		s.authentication, ui.Handler(ui.Config{
+			ExperimentalUseLogin: s.cfg.EnableWebSessionAuthentication,
+			LoginEnabled:         s.cfg.RequireWebSession(),
+			GetUser: func(ctx context.Context) *string {
+				if u, ok := ctx.Value(webSessionUserKey{}).(string); ok {
+					return &u
+				}
+				return nil
+			},
+		}))
 	s.mux.Handle("/", maybeAuthMux)
 
 	// Initialize grpc-gateway mux and context in order to get the /health
@@ -1948,40 +1946,3 @@ func (w *gzipResponseWriter) Close() error {
 	gzipResponseWriterPool.Put(w)
 	return err
 }
-
-func serveUIAssets(fileServer http.Handler, cfg Config) http.Handler {
-	return http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
-		if request.URL.Path != "/" {
-			fileServer.ServeHTTP(writer, request)
-			return
-		}
-
-		// Construct arguments for template.
-		tmplArgs := ui.IndexHTMLArgs{
-			ExperimentalUseLogin: cfg.EnableWebSessionAuthentication,
-			LoginEnabled:         cfg.RequireWebSession(),
-			Tag:                  build.GetInfo().Tag,
-			Version:              build.VersionPrefix(),
-		}
-		loggedInUser, ok := request.Context().Value(webSessionUserKey{}).(string)
-		if ok && loggedInUser != "" {
-			tmplArgs.LoggedInUser = &loggedInUser
-		}
-
-		argsJSON, err := json.Marshal(tmplArgs)
-		if err != nil {
-			http.Error(writer, err.Error(), 500)
-			return
-		}
-
-		// Execute the template.
-		writer.Header().Add("Content-Type", "text/html")
-		if err := ui.IndexHTMLTemplate.Execute(writer, map[string]template.JS{
-			"DataFromServer": template.JS(string(argsJSON)),
-		}); err != nil {
-			wrappedErr := errors.Wrap(err, "templating index.html")
-			http.Error(writer, wrappedErr.Error(), 500)
-			log.Error(request.Context(), wrappedErr)
-		}
-	})
-}
diff --git a/pkg/ui/ui.go b/pkg/ui/ui.go
@@ -21,27 +21,22 @@
 package ui
 
 import (
+	"bytes"
+	"context"
 	"fmt"
 	"html/template"
+	"net/http"
 	"os"
 
 	"github.com/cockroachdb/cockroach/pkg/build"
+	"github.com/cockroachdb/cockroach/pkg/util/log"
+	assetfs "github.com/elazarl/go-bindata-assetfs"
+	"github.com/pkg/errors"
 )
 
-var indexHTML = []byte(fmt.Sprintf(`<!DOCTYPE html>
-<title>CockroachDB</title>
-Binary built without web UI.
-<hr>
-<em>%s</em>`, build.GetInfo().Short()))
-
 // Asset loads and returns the asset for the given name. It returns an error if
 // the asset could not be found or could not be loaded.
-var Asset = func(name string) ([]byte, error) {
-	if name == "index.html" {
-		return indexHTML, nil
-	}
-	return nil, os.ErrNotExist
-}
+var Asset func(name string) ([]byte, error)
 
 // AssetDir returns the file names below a certain directory in the embedded
 // filesystem.
@@ -58,21 +53,22 @@ var Asset = func(name string) ([]byte, error) {
 // AssetDir("data") returns []string{"foo.txt", "img"}
 // AssetDir("data/img") returns []string{"a.png", "b.png"}
 // AssetDir("foo.txt") and AssetDir("notexist") return errors
-var AssetDir = func(name string) ([]string, error) {
-	if name == "" {
-		return []string{"index.html"}, nil
-	}
-	return nil, os.ErrNotExist
-}
+var AssetDir func(name string) ([]string, error)
 
 // AssetInfo loads and returns metadata for the asset with the given name. It
 // returns an error if the asset could not be found or could not be loaded.
 var AssetInfo func(name string) (os.FileInfo, error)
 
-// IndexHTMLTemplate takes arguments about the current session and returns HTML which
-// includes the UI JavaScript bundles, plus a script tag which sets the currently logged in user
-// so that the UI JavaScript can decide whether to show a login page.
-var IndexHTMLTemplate = template.Must(template.New("index").Parse(`<!DOCTYPE html>
+// haveUI returns whether the admin UI has been linked into the binary.
+func haveUI() bool {
+	return Asset != nil && AssetDir != nil && AssetInfo != nil
+}
+
+// indexTemplate takes arguments about the current session and returns HTML
+// which includes the UI JavaScript bundles, plus a script tag which sets the
+// currently logged in user so that the UI JavaScript can decide whether to show
+// a login page.
+var indexHTMLTemplate = template.Must(template.New("index").Parse(`<!DOCTYPE html>
 <html>
 	<head>
 		<title>Cockroach Console</title>
@@ -83,7 +79,7 @@ var IndexHTMLTemplate = template.Must(template.New("index").Parse(`<!DOCTYPE htm
 		<div id="react-layout"></div>
 
 		<script>
-			window.dataFromServer = {{ .DataFromServer }};
+			window.dataFromServer = {{.}};
 		</script>
 
 		<script src="protos.dll.js" type="text/javascript"></script>
@@ -93,11 +89,59 @@ var IndexHTMLTemplate = template.Must(template.New("index").Parse(`<!DOCTYPE htm
 </html>
 `))
 
-// IndexHTMLArgs are the arguments to IndexHTMLTemplate.
-type IndexHTMLArgs struct {
+type indexHTMLArgs struct {
 	ExperimentalUseLogin bool
 	LoginEnabled         bool
 	LoggedInUser         *string
 	Tag                  string
 	Version              string
 }
+
+// bareIndexHTML is used in place of indexHTMLTemplate when the binary is built
+// without the web UI.
+var bareIndexHTML = []byte(fmt.Sprintf(`<!DOCTYPE html>
+<title>CockroachDB</title>
+Binary built without web UI.
+<hr>
+<em>%s</em>`, build.GetInfo().Short()))
+
+// Config contains the configuration parameters for Handler.
+type Config struct {
+	ExperimentalUseLogin bool
+	LoginEnabled         bool
+	GetUser              func(ctx context.Context) *string
+}
+
+// Handler returns an http.Handler that serves the UI.
+func Handler(cfg Config) http.Handler {
+	fileServer := http.FileServer(&assetfs.AssetFS{
+		Asset:     Asset,
+		AssetDir:  AssetDir,
+		AssetInfo: AssetInfo,
+	})
+	buildInfo := build.GetInfo()
+
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if !haveUI() {
+			http.ServeContent(w, r, "index.html", buildInfo.GoTime(), bytes.NewReader(bareIndexHTML))
+			return
+		}
+
+		if r.URL.Path != "/" {
+			fileServer.ServeHTTP(w, r)
+			return
+		}
+
+		if err := indexHTMLTemplate.Execute(w, indexHTMLArgs{
+			ExperimentalUseLogin: cfg.ExperimentalUseLogin,
+			LoginEnabled:         cfg.LoginEnabled,
+			LoggedInUser:         cfg.GetUser(r.Context()),
+			Tag:                  buildInfo.Tag,
+			Version:              build.VersionPrefix(),
+		}); err != nil {
+			err = errors.Wrap(err, "templating index.html")
+			http.Error(w, err.Error(), 500)
+			log.Error(r.Context(), err)
+		}
+	})
+}