Split the SSO branch of Login into its own method

2026-09-12 17:59:06 +02:00 · 2026-08-25 17:56:21 +02:00
parent e8d6aa6344
commit 384102d017
1 changed files with 76 additions and 53 deletions
@@ -675,59 +675,7 @@ func (s *Server) Login(callerCtx context.Context, msg *proto.LoginRequest) (*pro
 	}

 	if msg.SetupKey == "" {
-		hint := ""
-		if msg.Hint != nil {
-			hint = *msg.Hint
-		}
-		oAuthFlow, err := auth.NewOAuthFlow(ctx, config, msg.IsUnixDesktopClient, false, hint)
-		if err != nil {
-			state.Set(internal.StatusLoginFailed)
-			// enrolling a device is the one flow a setup key can replace. NotFound so the CLI
-			// stops its backoff loop and shows this instead of retrying a permanent condition.
-			if auth.IsSSOUnavailable(err) {
-				return nil, gstatus.Error(codes.NotFound, auth.WithSetupKeyAdvice(err).Error())
-			}
-			return nil, err
-		}
-
-		if s.oauthAuthFlow.flow != nil && s.oauthAuthFlow.flow.GetClientID(ctx) == oAuthFlow.GetClientID(ctx) {
-			if s.oauthAuthFlow.expiresAt.After(time.Now().Add(90 * time.Second)) {
-				log.Debugf("using previous oauth flow info")
-				state.Set(internal.StatusNeedsLogin)
-				return &proto.LoginResponse{
-					NeedsSSOLogin:           true,
-					VerificationURI:         s.oauthAuthFlow.info.VerificationURI,
-					VerificationURIComplete: s.oauthAuthFlow.info.VerificationURIComplete,
-					UserCode:                s.oauthAuthFlow.info.UserCode,
-				}, nil
-			} else {
-				log.Warnf("canceling previous waiting execution")
-				if s.oauthAuthFlow.waitCancel != nil {
-					s.oauthAuthFlow.waitCancel()
-				}
-			}
-		}
-
-		authInfo, err := oAuthFlow.RequestAuthInfo(ctx)
-		if err != nil {
-			log.Errorf("getting a request OAuth flow failed: %v", err)
-			return nil, err
-		}
-
-		s.mutex.Lock()
-		s.oauthAuthFlow.flow = oAuthFlow
-		s.oauthAuthFlow.info = authInfo
-		s.oauthAuthFlow.expiresAt = time.Now().Add(time.Duration(authInfo.ExpiresIn) * time.Second)
-		s.mutex.Unlock()
-
-		state.Set(internal.StatusNeedsLogin)
-
-		return &proto.LoginResponse{
-			NeedsSSOLogin:           true,
-			VerificationURI:         authInfo.VerificationURI,
-			VerificationURIComplete: authInfo.VerificationURIComplete,
-			UserCode:                authInfo.UserCode,
-		}, nil
+		return s.ssoLogin(ctx, config, msg)
 	}

 	// Setup-key path: we are about to dial Management with the key, so the
@@ -743,6 +691,81 @@ func (s *Server) Login(callerCtx context.Context, msg *proto.LoginRequest) (*pro
 	return &proto.LoginResponse{}, nil
 }

+// ssoLogin starts the browser leg of a login that carries no setup key: it builds the OAuth
+// flow, reuses the one already in flight when it targets the same client, and returns the
+// verification URI for the caller to open.
+func (s *Server) ssoLogin(ctx context.Context, config *profilemanager.Config, msg *proto.LoginRequest) (*proto.LoginResponse, error) {
+	state := internal.CtxGetState(s.rootCtx)
+
+	var hint string
+	if msg.Hint != nil {
+		hint = *msg.Hint
+	}
+
+	oAuthFlow, err := auth.NewOAuthFlow(ctx, config, msg.IsUnixDesktopClient, false, hint)
+	if err != nil {
+		state.Set(internal.StatusLoginFailed)
+		// enrolling a device is the one flow a setup key can replace. NotFound so the CLI
+		// stops its backoff loop and shows this instead of retrying a permanent condition.
+		if auth.IsSSOUnavailable(err) {
+			return nil, gstatus.Error(codes.NotFound, auth.WithSetupKeyAdvice(err).Error())
+		}
+		return nil, err
+	}
+
+	if resp := s.reusableOAuthFlow(ctx, oAuthFlow); resp != nil {
+		state.Set(internal.StatusNeedsLogin)
+		return resp, nil
+	}
+
+	authInfo, err := oAuthFlow.RequestAuthInfo(ctx)
+	if err != nil {
+		log.Errorf("getting a request OAuth flow failed: %v", err)
+		return nil, err
+	}
+
+	s.mutex.Lock()
+	s.oauthAuthFlow.flow = oAuthFlow
+	s.oauthAuthFlow.info = authInfo
+	s.oauthAuthFlow.expiresAt = time.Now().Add(time.Duration(authInfo.ExpiresIn) * time.Second)
+	s.mutex.Unlock()
+
+	state.Set(internal.StatusNeedsLogin)
+
+	return &proto.LoginResponse{
+		NeedsSSOLogin:           true,
+		VerificationURI:         authInfo.VerificationURI,
+		VerificationURIComplete: authInfo.VerificationURIComplete,
+		UserCode:                authInfo.UserCode,
+	}, nil
+}
+
+// reusableOAuthFlow returns the response for a flow that is still in flight for the same client,
+// or nil when there is none to reuse. A flow too close to expiry is dropped along with the
+// execution waiting on it, so the caller starts a fresh one.
+func (s *Server) reusableOAuthFlow(ctx context.Context, oAuthFlow auth.OAuthFlow) *proto.LoginResponse {
+	if s.oauthAuthFlow.flow == nil || s.oauthAuthFlow.flow.GetClientID(ctx) != oAuthFlow.GetClientID(ctx) {
+		return nil
+	}
+
+	if !s.oauthAuthFlow.expiresAt.After(time.Now().Add(90 * time.Second)) {
+		log.Warnf("canceling previous waiting execution")
+		if s.oauthAuthFlow.waitCancel != nil {
+			s.oauthAuthFlow.waitCancel()
+		}
+		return nil
+	}
+
+	log.Debugf("using previous oauth flow info")
+
+	return &proto.LoginResponse{
+		NeedsSSOLogin:           true,
+		VerificationURI:         s.oauthAuthFlow.info.VerificationURI,
+		VerificationURIComplete: s.oauthAuthFlow.info.VerificationURIComplete,
+		UserCode:                s.oauthAuthFlow.info.UserCode,
+	}
+}
+
 // WaitSSOLogin validates the supplied userCode against the in-flight OAuth
 // device/PKCE flow and blocks until the user finishes the browser leg.
 //