client/web: move more session logic to auth.go

Updates tailscale/corp#14335

Signed-off-by: Sonia Appasamy <sonia@tailscale.com>

client/web/auth.go

@@ -77,8 +77,8 @@ var (
 	errNotOwner           = errors.New("not-owner")
 )
 
-// getTailscaleBrowserSession retrieves the browser session associated with
-// the request, if one exists.
+// getSession retrieves the browser session associated with the request,
+// if one exists.
 //
 // An error is returned in any of the following cases:
 //
@@ -104,7 +104,7 @@ var (
 //
 // The WhoIsResponse is always populated, with a non-nil Node and UserProfile,
 // unless getTailscaleBrowserSession reports errNotUsingTailscale.
-func (s *Server) getTailscaleBrowserSession(r *http.Request) (*browserSession, *apitype.WhoIsResponse, error) {
+func (s *Server) getSession(r *http.Request) (*browserSession, *apitype.WhoIsResponse, error) {
 	whoIs, whoIsErr := s.lc.WhoIs(r.Context(), r.RemoteAddr)
 	status, statusErr := s.lc.StatusWithoutPeers(r.Context())
 	switch {
@@ -148,6 +148,52 @@ func (s *Server) getTailscaleBrowserSession(r *http.Request) (*browserSession, *
 	return session, whoIs, nil
 }
 
+// newSession creates a new session associated with the given source user/node,
+// and stores it back to the session cache. Creating of a new session includes
+// generating a new auth URL from the control server.
+func (s *Server) newSession(ctx context.Context, src *apitype.WhoIsResponse) (*browserSession, error) {
+	d, err := s.getOrAwaitAuth(ctx, "", src.Node.ID)
+	if err != nil {
+		return nil, err
+	}
+	sid, err := s.newSessionID()
+	if err != nil {
+		return nil, err
+	}
+	session := &browserSession{
+		ID:      sid,
+		SrcNode: src.Node.ID,
+		SrcUser: src.UserProfile.ID,
+		AuthID:  d.ID,
+		AuthURL: d.URL,
+		Created: s.timeNow(),
+	}
+	s.browserSessions.Store(sid, session)
+	return session, nil
+}
+
+// awaitUserAuth blocks until the given session auth has been completed
+// by the user on the control server, then updates the session cache upon
+// completion. An error is returned if control auth failed for any reason.
+func (s *Server) awaitUserAuth(ctx context.Context, session *browserSession) error {
+	if session.isAuthorized(s.timeNow()) {
+		return nil // already authorized
+	}
+	d, err := s.getOrAwaitAuth(ctx, session.AuthID, session.SrcNode)
+	if err != nil {
+		// Clean up the session. Doing this on any error from control
+		// server to avoid the user getting stuck with a bad session
+		// cookie.
+		s.browserSessions.Delete(session.ID)
+		return err
+	}
+	if d.Complete {
+		session.Authenticated = d.Complete
+		s.browserSessions.Store(session.ID, session)
+	}
+	return nil
+}
+
 // getOrAwaitAuth connects to the control server for user auth,
 // with the following behavior:
 //

client/web/web.go

@@ -214,10 +214,10 @@ func (s *Server) authorizeRequest(w http.ResponseWriter, r *http.Request) (ok bo
 		case strings.HasPrefix(r.URL.Path, "/api/"):
 			// All other /api/ endpoints require a valid browser session.
 			//
-			// TODO(sonia): s.getTailscaleBrowserSession calls whois again,
+			// TODO(sonia): s.getSession calls whois again,
 			// should try and use the above call instead of running another
 			// localapi request.
-			session, _, err := s.getTailscaleBrowserSession(r)
+			session, _, err := s.getSession(r)
 			if err != nil || !session.isAuthorized(s.timeNow()) {
 				http.Error(w, "no valid session", http.StatusUnauthorized)
 				return false
@@ -289,57 +289,34 @@ func (s *Server) serveTailscaleAuth(w http.ResponseWriter, r *http.Request) {
 	}
 	var resp authResponse
 
-	session, whois, err := s.getTailscaleBrowserSession(r)
+	session, whois, err := s.getSession(r)
 	switch {
 	case err != nil && !errors.Is(err, errNoSession):
 		http.Error(w, err.Error(), http.StatusUnauthorized)
 		return
 	case session == nil:
 		// Create a new session.
-		d, err := s.getOrAwaitAuth(r.Context(), "", whois.Node.ID)
+		session, err := s.newSession(r.Context(), whois)
 		if err != nil {
 			http.Error(w, err.Error(), http.StatusInternalServerError)
 			return
 		}
-		sid, err := s.newSessionID()
-		if err != nil {
-			http.Error(w, err.Error(), http.StatusInternalServerError)
-			return
-		}
-		session := &browserSession{
-			ID:      sid,
-			SrcNode: whois.Node.ID,
-			SrcUser: whois.UserProfile.ID,
-			AuthID:  d.ID,
-			AuthURL: d.URL,
-			Created: s.timeNow(),
-		}
-		s.browserSessions.Store(sid, session)
 		// Set the cookie on browser.
 		http.SetCookie(w, &http.Cookie{
 			Name:    sessionCookieName,
-			Value:   sid,
-			Raw:     sid,
+			Value:   session.ID,
+			Raw:     session.ID,
 			Path:    "/",
 			Expires: session.expires(),
 		})
-		resp = authResponse{OK: false, AuthURL: d.URL}
+		resp = authResponse{OK: false, AuthURL: session.AuthURL}
 	case !session.isAuthorized(s.timeNow()):
 		if r.URL.Query().Get("wait") == "true" {
 			// Client requested we block until user completes auth.
-			d, err := s.getOrAwaitAuth(r.Context(), session.AuthID, whois.Node.ID)
-			if err != nil {
+			if err := s.awaitUserAuth(r.Context(), session); err != nil {
 				http.Error(w, err.Error(), http.StatusUnauthorized)
-				// Clean up the session. Doing this on any error from control
-				// server to avoid the user getting stuck with a bad session
-				// cookie.
-				s.browserSessions.Delete(session.ID)
 				return
 			}
-			if d.Complete {
-				session.Authenticated = d.Complete
-				s.browserSessions.Store(session.ID, session)
-			}
 		}
 		if session.isAuthorized(s.timeNow()) {
 			resp = authResponse{OK: true}

client/web/web_test.go

@@ -302,7 +302,7 @@ func TestGetTailscaleBrowserSession(t *testing.T) {
 			if tt.cookie != "" {
 				r.AddCookie(&http.Cookie{Name: sessionCookieName, Value: tt.cookie})
 			}
-			session, _, err := s.getTailscaleBrowserSession(r)
+			session, _, err := s.getSession(r)
 			if !errors.Is(err, tt.wantError) {
 				t.Errorf("wrong error; want=%v, got=%v", tt.wantError, err)
 			}