From 4a944603772dd62f018e9902deccc7919fb47cd0 Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 10 Oct 2025 11:14:52 +0800 Subject: [PATCH 1/2] Add MFA required error handling and update login response structure --- openapi/response/response.go | 1 + openapi/user/login.go | 12 ++++++++++++ openapi/user/oauth.go | 8 ++++++++ openapi/user/types.go | 2 ++ openapi/user/user.go | 16 ---------------- 5 files changed, 23 insertions(+), 16 deletions(-) diff --git a/openapi/response/response.go b/openapi/response/response.go index 0c25c84b..3c8643c0 100644 --- a/openapi/response/response.go +++ b/openapi/response/response.go @@ -112,6 +112,7 @@ 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."} // Configuration and service errors ErrInvalidConfiguration = types.ErrInvalidConfiguration diff --git a/openapi/user/login.go b/openapi/user/login.go index 69f3d4c5..29ac5d46 100644 --- a/openapi/user/login.go +++ b/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, diff --git a/openapi/user/oauth.go b/openapi/user/oauth.go index 47475847..1568a9c6 100644 --- a/openapi/user/oauth.go +++ b/openapi/user/oauth.go @@ -176,6 +176,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(), diff --git a/openapi/user/types.go b/openapi/user/types.go index bdc57abb..f9890952 100644 --- a/openapi/user/types.go +++ b/openapi/user/types.go @@ -163,6 +163,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"` diff --git a/openapi/user/user.go b/openapi/user/user.go index 2342bfc8..4440f3ac 100644 --- a/openapi/user/user.go +++ b/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!"}) } From 5e5b633fba5c60da24ca2fceb89a2588f130d880 Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 10 Oct 2025 11:57:47 +0800 Subject: [PATCH 2/2] Add team retrieval and counting methods for user membership - Introduced GetTeamsByMember and CountTeamsByMember methods in the DefaultUser struct to retrieve teams associated with a specific member and count the total number of teams, respectively. - Updated the UserProvider interface to include these new methods. - Enhanced the login response structure to include a status indicating whether team selection is required based on the user's team membership. - Added error handling for team retrieval in the authentication process. --- openapi/oauth/providers/user/team.go | 74 ++++++++++++++++++++++++++++ openapi/oauth/types/interfaces.go | 2 + openapi/response/response.go | 1 + openapi/user/oauth.go | 37 ++++++++++++++ openapi/user/types.go | 26 +++++++--- 5 files changed, 133 insertions(+), 7 deletions(-) diff --git a/openapi/oauth/providers/user/team.go b/openapi/oauth/providers/user/team.go index 84cc06de..f607a627 100644 --- a/openapi/oauth/providers/user/team.go +++ b/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{ diff --git a/openapi/oauth/types/interfaces.go b/openapi/oauth/types/interfaces.go index 7f85deec..adc2ba8d 100644 --- a/openapi/oauth/types/interfaces.go +++ b/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 diff --git a/openapi/response/response.go b/openapi/response/response.go index 3c8643c0..23e97c0c 100644 --- a/openapi/response/response.go +++ b/openapi/response/response.go @@ -113,6 +113,7 @@ var ( 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 diff --git a/openapi/user/oauth.go b/openapi/user/oauth.go index 1568a9c6..81b85c89 100644 --- a/openapi/user/oauth.go +++ b/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" @@ -195,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, @@ -204,6 +222,7 @@ func authback(c *gin.Context) { ExpiresIn: loginResponse.ExpiresIn, RefreshTokenExpiresIn: loginResponse.RefreshTokenExpiresIn, MFAEnabled: loginResponse.MFAEnabled, + Status: status, }) } @@ -453,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) diff --git a/openapi/user/types.go b/openapi/user/types.go index f9890952..8258cbb7 100644 --- a/openapi/user/types.go +++ b/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 @@ -177,13 +187,15 @@ type LoginResponse struct { // LoginSuccessResponse represents the response for login success type LoginSuccessResponse struct { - IDToken string `json:"id_token,omitempty"` - AccessToken string `json:"access_token,omitempty"` - 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"` + IDToken string `json:"id_token,omitempty"` + AccessToken string `json:"access_token,omitempty"` + 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"` + Error *response.ErrorResponse `json:"error,omitempty"` } // Built-in preset mapping types