opencode/yao
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Merge pull request #1183 from trheyi/main

Browse source 
Add MFA required error handling and update login response structure
This commit is contained in:
Max 2025-10-10 11:58:44 +08:00 • committed by GitHub
commit 8d146bb355
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
7 changed files with 156 additions and 23 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

74
openapi/oauth/providers/user/team.go
View file

@ -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{

2
openapi/oauth/types/interfaces.go
View file

@ -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

2
openapi/response/response.go
View file

@ -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

12
openapi/user/login.go
View file

@ -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,

45
openapi/user/oauth.go
View file

@ -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)

28
openapi/user/types.go
View file

@ -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"`
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

16
openapi/user/user.go
View file

@ -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!"})
}