Implement MFA token generation and update login response handling

- Removed the previous MFA check from the LoginThirdParty function and integrated MFA token generation within the LoginByUserID function.
- Enhanced the LoginResponse structure to include MFAToken and MFATokenExpiresIn fields for better MFA handling.
- Updated the SendLoginCookies function to manage MFA token cookies appropriately.
- Adjusted the authback function to respond with MFA status and token when MFA is required, improving the login flow for users with MFA enabled.

openapi/user/login.go

@@ -135,16 +135,6 @@ func LoginThirdParty(providerID string, userinfo *oauthtypes.OIDCUserInfo, ip st
 		return nil, err
 	}
 
-	// If MFA Enabled, should return MFA required response
-	mfaEnabled, err := userProvider.IsMFAEnabled(ctx, userID)
-	if err != nil {
-		return nil, err
-	}
-
-	if mfaEnabled {
-		return nil, response.ErrMFARequired
-	}
-
 	return LoginByUserID(userID, ip)
 }
 
@@ -166,6 +156,28 @@ func LoginByUserID(userid string, ip string) (*LoginResponse, error) {
 		return nil, err
 	}
 
+	// Get MFA enabled status from user data
+	mfaEnabled := toBool(user["mfa_enabled"])
+
+	// If MFA enabled, generate MFA token
+	if mfaEnabled {
+
+		// Sign temporary access token for MFA
+		var mfaExpire int = 10 * 60 // 10 minutes
+		mfaToken, err := oauth.OAuth.MakeAccessToken(yaoClientConfig.ClientID, ScopeMFAVerification, userid, mfaExpire)
+		if err != nil {
+			return nil, err
+		}
+
+		return &LoginResponse{
+			UserID:            userid,
+			MFAToken:          mfaToken,
+			MFATokenExpiresIn: mfaExpire,
+			MFAEnabled:        mfaEnabled,
+			Status:            LoginStatusMFA,
+		}, nil
+	}
+
 	// Update Last Login
 	err = userProvider.UpdateUserLastLogin(ctx, userid, ip)
 	if err != nil {
@@ -203,8 +215,16 @@ func LoginByUserID(userid string, ip string) (*LoginResponse, error) {
 		return nil, err
 	}
 
-	// Get MFA enabled status from user data
+	// Count User Teams
-	mfaEnabled := toBool(user["mfa_enabled"])
+	numTeams, err := countUserTeams(ctx, userid)
+	if err != nil {
+		return nil, err
+	}
+
+	status := LoginStatusSuccess
+	if numTeams > 0 {
+		status = LoginStatusTeamSelection
+	}
 
 	return &LoginResponse{
 		UserID:                userid,
@@ -217,6 +237,7 @@ func LoginByUserID(userid string, ip string) (*LoginResponse, error) {
 		TokenType:             "Bearer",
 		MFAEnabled:            mfaEnabled,
 		Scope:                 strings.Join(scopes, " "),
+		Status:                status,
 	}, nil
 }
 
@@ -228,12 +249,26 @@ func generateSessionID() string {
 // SendLoginCookies sends all necessary cookies for a successful login
 // This includes access token, refresh token, and session ID cookies with appropriate security settings
 func SendLoginCookies(c *gin.Context, loginResponse *LoginResponse, sessionID string) {
-	// Format tokens with Bearer prefix
+
+	// Send session ID cookie
+	expires := time.Now().Add(time.Duration(yaoClientConfig.ExpiresIn) * time.Second)
+	options := response.NewSecureCookieOptions().
+		WithExpires(expires).
+		WithSameSite("Strict")
+	response.SendSecureCookieWithOptions(c, "session_id", sessionID, options)
+
+	// MFA Temporary Access Token
+	if loginResponse.Status == LoginStatusMFA {
+		mfaToken := fmt.Sprintf("Bearer %s", loginResponse.MFAToken)
+		response.SendAccessTokenCookieWithExpiry(c, mfaToken, expires)
+		return
+	}
+
+	// Normal Access Token
 	accessToken := fmt.Sprintf("%s %s", loginResponse.TokenType, loginResponse.AccessToken)
 	refreshToken := fmt.Sprintf("%s %s", loginResponse.TokenType, loginResponse.RefreshToken)
 
 	// Calculate expiration times
-	expires := time.Now().Add(time.Duration(loginResponse.ExpiresIn) * time.Second)
 	refreshExpires := time.Now().Add(time.Duration(loginResponse.RefreshTokenExpiresIn) * time.Second)
 
 	// Send access token cookie
@@ -241,11 +276,4 @@ func SendLoginCookies(c *gin.Context, loginResponse *LoginResponse, sessionID st
 
 	// Send refresh token cookie
 	response.SendRefreshTokenCookieWithExpiry(c, refreshToken, refreshExpires)
-
-	// Send session ID cookie with the same expiration as access token
-	// Using HTTP-only flag for security
-	options := response.NewSecureCookieOptions().
-		WithExpires(expires).
-		WithSameSite("Strict")
-	response.SendSecureCookieWithOptions(c, "session_id", sessionID, options)
 }

openapi/user/oauth.go

@@ -178,14 +178,6 @@ func authback(c *gin.Context) {
 	// LoginThirdParty(providerID, userInfo)
 	loginResponse, err := LoginThirdParty(providerID, userInfo, userIPAddress(c))
 	if err != nil {
-
-		// Redirect to MFA required page
-		if err == response.ErrMFARequired {
-			response.RespondWithError(c, response.StatusUnauthorized, response.ErrMFARequired)
-			return
-		}
-
-		// Other errors
 		errorResp := &response.ErrorResponse{
 			Code:             response.ErrInvalidRequest.Code,
 			ErrorDescription: "Failed to login: " + err.Error(),
@@ -197,22 +189,18 @@ func authback(c *gin.Context) {
 	// Send all login cookies (access token, refresh token, and session ID)
 	SendLoginCookies(c, loginResponse, sid)
 
-	// Get Teams
+	// MFA Response
-	numTeams, err := countUserTeams(c.Request.Context(), loginResponse.UserID)
+	if loginResponse.Status == LoginStatusMFA {
-	if err != nil {
+		response.RespondWithSuccess(c, response.StatusOK, LoginSuccessResponse{
-		errorResp := &response.ErrorResponse{
+			SessionID:         sid,
-			Code:             response.ErrInvalidRequest.Code,
+			MFAEnabled:        loginResponse.MFAEnabled,
-			ErrorDescription: "Failed to count teams: " + err.Error(),
+			Status:            loginResponse.Status,
-		}
+			MFAToken:          loginResponse.MFAToken,
-		response.RespondWithError(c, response.StatusInternalServerError, errorResp)
+			MFATokenExpiresIn: loginResponse.MFATokenExpiresIn,
+		})
 		return
 	}
 
-	status := LoginStatusSuccess
-	if numTeams > 0 {
-		status = LoginStatusTeamSelection
-	}
-
 	// Send IDToken to the client
 	response.RespondWithSuccess(c, response.StatusOK, LoginSuccessResponse{
 		SessionID:             sid,
@@ -222,7 +210,7 @@ func authback(c *gin.Context) {
 		ExpiresIn:             loginResponse.ExpiresIn,
 		RefreshTokenExpiresIn: loginResponse.RefreshTokenExpiresIn,
 		MFAEnabled:            loginResponse.MFAEnabled,
-		Status:                status,
+		Status:                loginResponse.Status,
 	})
 }
 

openapi/user/types.go

@@ -2,16 +2,23 @@ package user
 
 import (
 	oauthtypes "github.com/yaoapp/yao/openapi/oauth/types"
-	"github.com/yaoapp/yao/openapi/response"
 )
 
+// LoginStatus represents the login status
+type LoginStatus string
+
 const (
 	// LoginStatusSuccess is the success status
-	LoginStatusSuccess = "ok"
+	LoginStatusSuccess LoginStatus = "ok"
 	// LoginStatusMFA is the MFA status
-	LoginStatusMFA = "mfa_required"
+	LoginStatusMFA LoginStatus = "mfa_required"
 	// LoginStatusTeamSelection is the team selection status
-	LoginStatusTeamSelection = "team_selection_required"
+	LoginStatusTeamSelection LoginStatus = "team_selection_required"
+)
+
+const (
+	// ScopeMFAVerification is the MFA verification scope for temporary access token
+	ScopeMFAVerification = "mfa_verification"
 )
 
 // Config represents the signin page configuration
@@ -181,21 +188,25 @@ type LoginResponse struct {
 	ExpiresIn             int         `json:"expires_in,omitempty"`
 	RefreshTokenExpiresIn int         `json:"refresh_token_expires_in,omitempty"`
 	TokenType             string      `json:"token_type,omitempty"`
+	MFAToken              string      `json:"mfa_token,omitempty"`            // MFA token verification code
+	MFATokenExpiresIn     int         `json:"mfa_token_expires_in,omitempty"` // MFA token verification code expires in
 	MFAEnabled            bool        `json:"mfa_enabled,omitempty"`
 	Scope                 string      `json:"scope,omitempty"`
+	Status                LoginStatus `json:"status,omitempty"`
 }
 
 // LoginSuccessResponse represents the response for login success
 type LoginSuccessResponse struct {
 	IDToken               string      `json:"id_token,omitempty"`
 	AccessToken           string      `json:"access_token,omitempty"`
+	MFAToken              string      `json:"mfa_token,omitempty"`            // MFA token verification code
+	MFATokenExpiresIn     int         `json:"mfa_token_expires_in,omitempty"` // MFA token verification code expires in
 	SessionID             string      `json:"session_id,omitempty"`
 	RefreshToken          string      `json:"refresh_token,omitempty"`
 	ExpiresIn             int         `json:"expires_in,omitempty"`
 	MFAEnabled            bool        `json:"mfa_enabled"`
 	RefreshTokenExpiresIn int         `json:"refresh_token_expires_in,omitempty"`
-	Status                string                  `json:"status,omitempty"`
+	Status                LoginStatus `json:"status,omitempty"`
-	Error                 *response.ErrorResponse `json:"error,omitempty"`
 }
 
 // Built-in preset mapping types