server/ui: fix crash caused by race in server initialization

pkg/cmd/roachtest/cluster_init.go

@@ -23,6 +23,8 @@ import (
 	"strings"
 	"time"
 
+	"github.com/cockroachdb/cockroach/pkg/server"
+	"github.com/cockroachdb/cockroach/pkg/server/serverpb"
 	"github.com/cockroachdb/cockroach/pkg/util/retry"
 	"golang.org/x/sync/errgroup"
 )
@@ -143,16 +145,37 @@ func runClusterInit(ctx context.Context, t *test, c *cluster) {
 				{"/_status/nodes", http.StatusNotFound},
 			}
 			for _, tc := range httpTests {
-				resp, err := http.Get(urlMap[1] + tc.endpoint)
-				if err != nil {
-					t.Fatalf("unexpected error hitting %s endpoint: %v", tc.endpoint, err)
-				}
-				defer resp.Body.Close()
-				if resp.StatusCode != tc.expectedStatus {
-					bodyBytes, _ := ioutil.ReadAll(resp.Body)
-					t.Fatalf("unexpected response code %d (expected %d) hitting %s endpoint: %v",
-						resp.StatusCode, tc.expectedStatus, tc.endpoint, string(bodyBytes))
+				for _, withCookie := range []bool{false, true} {
+					req, err := http.NewRequest("GET", urlMap[1]+tc.endpoint, nil /* body */)
+					if err != nil {
+						t.Fatalf("unexpected error while constructing request for %s: %s", tc.endpoint, err)
+					}
+					if withCookie {
+						// Prevent regression of #25771 by also sending authenticated
+						// requests, like would be sent if an admin UI were open against
+						// this node while it booted.
+						cookie, err := server.EncodeSessionCookie(&serverpb.SessionCookie{
+							// The actual contents of the cookie don't matter; the presence of
+							// a valid encoded cookie is enough to trigger the authentication
+							// code paths.
+						})
+						if err != nil {
+							t.Fatal(err)
+						}
+						req.AddCookie(cookie)
+					}
+					resp, err := http.DefaultClient.Do(req)
+					if err != nil {
+						t.Fatalf("unexpected error hitting %s endpoint: %v", tc.endpoint, err)
+					}
+					defer resp.Body.Close()
+					if resp.StatusCode != tc.expectedStatus {
+						bodyBytes, _ := ioutil.ReadAll(resp.Body)
+						t.Fatalf("unexpected response code %d (expected %d) hitting %s endpoint: %v",
+							resp.StatusCode, tc.expectedStatus, tc.endpoint, string(bodyBytes))
+					}
 				}
+
 			}
 
 			c.Run(ctx, c.Node(initNode),

pkg/server/authentication.go

@@ -139,7 +139,7 @@ func (s *authenticationServer) UserLogin(
 		ID:     id,
 		Secret: secret,
 	}
-	cookie, err := encodeSessionCookie(cookieValue)
+	cookie, err := EncodeSessionCookie(cookieValue)
 	if err != nil {
 		return nil, apiInternalError(ctx, err)
 	}
@@ -327,6 +327,13 @@ type authenticationMux struct {
 	server *authenticationServer
 	inner  http.Handler
 
+	// allowAnonymous, if true, indicates that the authentication mux should
+	// call its inner HTTP handler even if the request doesn't have a valid
+	// session. If there is a valid session, the mux calls its inner handler
+	// with a context containing the username and session ID.
+	//
+	// If allowAnonymous is false, the mux returns an error if there is no
+	// valid session.
 	allowAnonymous bool
 }
 
@@ -369,7 +376,8 @@ func (am *authenticationMux) ServeHTTP(w http.ResponseWriter, req *http.Request)
 	am.inner.ServeHTTP(w, req)
 }
 
-func encodeSessionCookie(sessionCookie *serverpb.SessionCookie) (*http.Cookie, error) {
+// EncodeSessionCookie encodes a SessionCookie proto into an http.Cookie.
+func EncodeSessionCookie(sessionCookie *serverpb.SessionCookie) (*http.Cookie, error) {
 	cookieValueBytes, err := protoutil.Marshal(sessionCookie)
 	if err != nil {
 		return nil, errors.Wrap(err, "session cookie could not be encoded")

pkg/server/authentication_test.go

@@ -551,7 +551,7 @@ func TestLogout(t *testing.T) {
 	if err != nil {
 		t.Fatal(err)
 	}
-	encodedCookie, err := encodeSessionCookie(cookie)
+	encodedCookie, err := EncodeSessionCookie(cookie)
 	if err != nil {
 		t.Fatal(err)
 	}

pkg/server/server.go

@@ -1234,21 +1234,6 @@ func (s *Server) Start(ctx context.Context) error {
 	// endpoints.
 	s.mux.Handle(debug.Endpoint, debug.NewServer(s.st))
 
-	// 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, 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
 	// endpoint working even before the node has fully initialized.
 	jsonpb := &protoutil.JSONPb{
@@ -1269,11 +1254,6 @@ func (s *Server) Start(ctx context.Context) error {
 	gwCtx, gwCancel := context.WithCancel(s.AnnotateCtx(context.Background()))
 	s.stopper.AddCloser(stop.CloserFn(gwCancel))
 
-	var authHandler http.Handler = gwMux
-	if s.cfg.RequireWebSession() {
-		authHandler = newAuthenticationMux(s.authentication, authHandler)
-	}
-
 	// Setup HTTP<->gRPC handlers.
 	c1, c2 := net.Pipe()
 
@@ -1537,16 +1517,6 @@ func (s *Server) Start(ctx context.Context) error {
 
 	s.serveMode.set(modeOperational)
 
-	s.mux.Handle(adminPrefix, authHandler)
-	// Exempt the health check endpoint from authentication.
-	s.mux.Handle("/_admin/v1/health", gwMux)
-	s.mux.Handle(ts.URLPrefix, authHandler)
-	s.mux.Handle(statusPrefix, authHandler)
-	s.mux.Handle(loginPath, gwMux)
-	s.mux.Handle(logoutPath, authHandler)
-	s.mux.Handle(statusVars, http.HandlerFunc(s.status.handleVars))
-	log.Event(ctx, "added http endpoints")
-
 	log.Infof(ctx, "starting %s server at %s (use: %s)",
 		s.cfg.HTTPRequestScheme(), s.cfg.HTTPAddr, s.cfg.HTTPAdvertiseAddr)
 	log.Infof(ctx, "starting grpc/postgres server at %s", s.cfg.Addr)
@@ -1609,6 +1579,43 @@ func (s *Server) Start(ctx context.Context) error {
 	log.Info(ctx, "serving sql connections")
 	// Start servicing SQL connections.
 
+	// Serve UI assets.
+	//
+	// The authentication mux used here is created in "allow anonymous" mode so that the UI
+	// assets are served up whether or not there is a session. If there is a session, the mux
+	// adds it to the context, and it is templated into index.html so that the UI can show
+	// the username of the currently-logged-in user.
+	authenticatedUIHandler := newAuthenticationMuxAllowAnonymous(
+		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("/", authenticatedUIHandler)
+
+	// Register gRPC-gateway endpoints used by the admin UI.
+	var authHandler http.Handler = gwMux
+	if s.cfg.RequireWebSession() {
+		authHandler = newAuthenticationMux(s.authentication, authHandler)
+	}
+
+	s.mux.Handle(adminPrefix, authHandler)
+	// Exempt the health check endpoint from authentication.
+	s.mux.Handle("/_admin/v1/health", gwMux)
+	s.mux.Handle(ts.URLPrefix, authHandler)
+	s.mux.Handle(statusPrefix, authHandler)
+	s.mux.Handle(loginPath, gwMux)
+	s.mux.Handle(logoutPath, authHandler)
+	s.mux.Handle(statusVars, http.HandlerFunc(s.status.handleVars))
+	log.Event(ctx, "added http endpoints")
+
 	// Attempt to upgrade cluster version.
 	s.startAttemptUpgrade(ctx)
 

pkg/server/testserver.go

@@ -563,7 +563,7 @@ func (ts *TestServer) getAuthenticatedHTTPClientAndCookie() (
 				Secret: secret,
 			}
 			// Encode a session cookie and store it in a cookie jar.
-			cookie, err := encodeSessionCookie(rawCookie)
+			cookie, err := EncodeSessionCookie(rawCookie)
 			if err != nil {
 				return err
 			}

pkg/ui/ui.go

@@ -112,7 +112,9 @@ type Config struct {
 	GetUser              func(ctx context.Context) *string
 }
 
-// Handler returns an http.Handler that serves the UI.
+// Handler returns an http.Handler that serves the UI,
+// including index.html, which has some login-related variables
+// templated into it, as well as static assets.
 func Handler(cfg Config) http.Handler {
 	fileServer := http.FileServer(&assetfs.AssetFS{
 		Asset:     Asset,