Merge pull request #1183 from trheyi/main

Add MFA required error handling and update login response structure

openapi/oauth/providers/user/team.go

@@ -251,6 +251,80 @@ func (u *DefaultUser) GetTeamsByOwner(ctx context.Context, ownerID string) ([]ma
 	return u.GetTeams(ctx, param)
 }
 
+// GetTeamsByMember retrieves teams by member_id
+func (u *DefaultUser) GetTeamsByMember(ctx context.Context, memberID string) ([]maps.MapStr, error) {
+
+	// Set default select fields if not provided
+	param := model.QueryParam{
+		Select: []interface{}{"team_id", "user_id", "member_type"},
+		Wheres: []model.QueryWhere{
+			{Column: "user_id", Value: memberID},
+			{Column: "member_type", Value: "user"},
+			{Column: "status", Value: "active"},
+		},
+	}
+
+	if param.Select == nil {
+		param.Select = u.memberFields
+	}
+
+	m := model.Select(u.memberModel)
+	members, err := m.Get(param)
+	if err != nil {
+		return nil, fmt.Errorf(ErrFailedToGetTeam, err)
+	}
+
+	if len(members) == 0 {
+		return []maps.MapStr{}, nil
+	}
+
+	// Get team ids
+	teamIDs := []string{}
+	for _, member := range members {
+		teamIDs = append(teamIDs, member["team_id"].(string))
+	}
+
+	// Get teams
+	teams, err := u.GetTeams(ctx, model.QueryParam{
+		Select: u.teamFields,
+		Wheres: []model.QueryWhere{
+			{Column: "team_id", Value: teamIDs, Method: "in"},
+		},
+	})
+	if err != nil {
+		return nil, fmt.Errorf(ErrFailedToGetTeam, err)
+	}
+
+	return teams, nil
+}
+
+// CountTeamsByMember returns total count of teams by member_id
+func (u *DefaultUser) CountTeamsByMember(ctx context.Context, memberID string) (int64, error) {
+
+	param := model.QueryParam{
+		Select: []interface{}{"team_id", "user_id", "member_type"},
+		Wheres: []model.QueryWhere{
+			{Column: "user_id", Value: memberID},
+			{Column: "member_type", Value: "user"},
+			{Column: "status", Value: "active"},
+		},
+	}
+	// Use Paginate with a small page size to get the total count
+	// This is more reliable than manual COUNT(*) queries
+	m := model.Select(u.memberModel)
+	result, err := m.Paginate(param, 1, 1) // Get first page with 1 item to get total
+	if err != nil {
+		return 0, fmt.Errorf(ErrFailedToGetTeam, err)
+	}
+
+	// Extract total from pagination result using utility function
+	if totalInterface, ok := result["total"]; ok {
+		return parseIntFromDB(totalInterface)
+	}
+
+	return 0, fmt.Errorf("total not found in pagination result")
+}
+
 // GetTeamsByStatus retrieves teams by status
 func (u *DefaultUser) GetTeamsByStatus(ctx context.Context, status string) ([]maps.MapStr, error) {
 	param := model.QueryParam{

openapi/oauth/types/interfaces.go

@@ -266,7 +266,9 @@ type UserProvider interface {
 
 	// Team Query Methods
 	GetTeamsByOwner(ctx context.Context, ownerID string) ([]maps.MapStr, error)
+	GetTeamsByMember(ctx context.Context, memberID string) ([]maps.MapStr, error)
 	GetTeamsByStatus(ctx context.Context, status string) ([]maps.MapStr, error)
+	CountTeamsByMember(ctx context.Context, memberID string) (int64, error)
 
 	// Team Management
 	UpdateTeamStatus(ctx context.Context, teamID string, status string) error

openapi/response/response.go

@@ -112,6 +112,8 @@ var (
 	ErrInvalidClientMetadata    = &ErrorResponse{Code: "invalid_client_metadata", ErrorDescription: "The client metadata is invalid or contains unsupported values."}
 	ErrInvalidSoftwareStatement = &ErrorResponse{Code: "invalid_software_statement", ErrorDescription: "The software statement is invalid or cannot be verified."}
 	ErrUnapprovedSoftware       = &ErrorResponse{Code: "unapproved_software", ErrorDescription: "The software statement represents software that has been replaced or is otherwise invalid."}
+	ErrMFARequired              = &ErrorResponse{Code: "mfa_required", ErrorDescription: "Multi-factor authentication is required to access this resource."}
+	ErrTeamSelectionRequired    = &ErrorResponse{Code: "team_selection_required", ErrorDescription: "Team selection is required to access this resource."}
 
 	// Configuration and service errors
 	ErrInvalidConfiguration     = types.ErrInvalidConfiguration

openapi/user/login.go

@@ -135,6 +135,16 @@ 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)
 }
 
@@ -197,6 +207,8 @@ func LoginByUserID(userid string, ip string) (*LoginResponse, error) {
 	mfaEnabled := toBool(user["mfa_enabled"])
 
 	return &LoginResponse{
+		UserID:                userid,
+		Subject:               subject,
 		AccessToken:           accessToken,
 		IDToken:               oidcToken,
 		RefreshToken:          refreshToken,

openapi/user/oauth.go

@@ -1,6 +1,7 @@
 package user
 
 import (
+	"context"
 	"fmt"
 	"net"
 	"net/http"
@@ -13,6 +14,7 @@ import (
 	"github.com/google/uuid"
 	"github.com/yaoapp/gou/session"
 	"github.com/yaoapp/kun/log"
+	"github.com/yaoapp/kun/maps"
 	"github.com/yaoapp/yao/openapi/oauth"
 	"github.com/yaoapp/yao/openapi/response"
 	"github.com/yaoapp/yao/openapi/utils"
@@ -176,6 +178,14 @@ 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(),
@@ -187,6 +197,22 @@ func authback(c *gin.Context) {
 	// Send all login cookies (access token, refresh token, and session ID)
 	SendLoginCookies(c, loginResponse, sid)
 
+	// Get Teams
+	numTeams, err := countUserTeams(c.Request.Context(), loginResponse.UserID)
+	if err != nil {
+		errorResp := &response.ErrorResponse{
+			Code:             response.ErrInvalidRequest.Code,
+			ErrorDescription: "Failed to count teams: " + err.Error(),
+		}
+		response.RespondWithError(c, response.StatusInternalServerError, errorResp)
+		return
+	}
+
+	status := LoginStatusSuccess
+	if numTeams > 0 {
+		status = LoginStatusTeamSelection
+	}
+
 	// Send IDToken to the client
 	response.RespondWithSuccess(c, response.StatusOK, LoginSuccessResponse{
 		SessionID:             sid,
@@ -196,6 +222,7 @@ func authback(c *gin.Context) {
 		ExpiresIn:             loginResponse.ExpiresIn,
 		RefreshTokenExpiresIn: loginResponse.RefreshTokenExpiresIn,
 		MFAEnabled:            loginResponse.MFAEnabled,
+		Status:                status,
 	})
 }
 
@@ -445,6 +472,24 @@ func getUserInfo(providerID, state string) (string, error) {
 	return value.(string), nil
 }
 
+// getUserTeams gets the user teams
+func getUserTeams(ctx context.Context, userID string) ([]maps.MapStr, error) {
+	userProvider, err := oauth.OAuth.GetUserProvider()
+	if err != nil {
+		return nil, err
+	}
+	return userProvider.GetTeamsByMember(ctx, userID)
+}
+
+// countUserTeams counts the number of teams a user is a member of
+func countUserTeams(ctx context.Context, userID string) (int64, error) {
+	userProvider, err := oauth.OAuth.GetUserProvider()
+	if err != nil {
+		return 0, err
+	}
+	return userProvider.CountTeamsByMember(ctx, userID)
+}
+
 // removeUserInfo removes the user info from cache
 func removeUserInfo(providerID, state string) error {
 	key := userInfoKey(providerID, state)

openapi/user/types.go

@@ -2,6 +2,16 @@ package user
 
 import (
 	oauthtypes "github.com/yaoapp/yao/openapi/oauth/types"
+	"github.com/yaoapp/yao/openapi/response"
+)
+
+const (
+	// LoginStatusSuccess is the success status
+	LoginStatusSuccess = "ok"
+	// LoginStatusMFA is the MFA status
+	LoginStatusMFA = "mfa_required"
+	// LoginStatusTeamSelection is the team selection status
+	LoginStatusTeamSelection = "team_selection_required"
 )
 
 // Config represents the signin page configuration
@@ -163,6 +173,8 @@ type OIDCAddress = oauthtypes.OIDCAddress
 
 // LoginResponse represents the response for login
 type LoginResponse struct {
+	UserID                string `json:"user_id,omitempty"`
+	Subject               string `json:"subject,omitempty"`
 	AccessToken           string `json:"access_token"`
 	IDToken               string `json:"id_token,omitempty"`
 	RefreshToken          string `json:"refresh_token,omitempty"`
@@ -175,13 +187,15 @@ type LoginResponse struct {
 
 // LoginSuccessResponse represents the response for login success
 type LoginSuccessResponse struct {
-	IDToken               string `json:"id_token,omitempty"`
+	IDToken               string                  `json:"id_token,omitempty"`
-	AccessToken           string `json:"access_token,omitempty"`
+	AccessToken           string                  `json:"access_token,omitempty"`
-	SessionID             string `json:"session_id,omitempty"`
+	SessionID             string                  `json:"session_id,omitempty"`
-	RefreshToken          string `json:"refresh_token,omitempty"`
+	RefreshToken          string                  `json:"refresh_token,omitempty"`
-	ExpiresIn             int    `json:"expires_in,omitempty"`
+	ExpiresIn             int                     `json:"expires_in,omitempty"`
-	MFAEnabled            bool   `json:"mfa_enabled"`
+	MFAEnabled            bool                    `json:"mfa_enabled"`
-	RefreshTokenExpiresIn int    `json:"refresh_token_expires_in,omitempty"`
+	RefreshTokenExpiresIn int                     `json:"refresh_token_expires_in,omitempty"`
+	Status                string                  `json:"status,omitempty"`
+	Error                 *response.ErrorResponse `json:"error,omitempty"`
 }
 
 // Built-in preset mapping types

openapi/user/user.go

@@ -264,22 +264,6 @@ func attachThirdParty(group *gin.RouterGroup, oauth types.OAuth) {
 
 }
 
-// getTeamConfig returns the team configuration
-func getTeamConfig(c *gin.Context) {
-	locale := c.Query("locale")
-	if locale == "" {
-		locale = "en" // default locale
-	}
-
-	config := GetTeamConfig(locale)
-	if config == nil {
-		c.JSON(http.StatusNotFound, gin.H{"error": "Team configuration not found"})
-		return
-	}
-
-	c.JSON(http.StatusOK, config)
-}
-
 func placeholder(c *gin.Context) {
 	c.JSON(http.StatusOK, gin.H{"message": "Hello, World!"})
 }