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

Merge pull request #1184 from trheyi/main

Browse source 
Implement MFA token generation and update login response handling
This commit is contained in:
Max 2025-10-10 14:25:50 +08:00 • committed by GitHub
commit 8fb0399641
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 93 additions and 66 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

70
openapi/user/login.go
View file

@ -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
mfaEnabled := toBool(user["mfa_enabled"])
// Count User Teams
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)
}

32
openapi/user/oauth.go
View file

@ -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
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)
// MFA Response
if loginResponse.Status == LoginStatusMFA {
response.RespondWithSuccess(c, response.StatusOK, LoginSuccessResponse{
SessionID: sid,
MFAEnabled: loginResponse.MFAEnabled,
Status: loginResponse.Status,
MFAToken: loginResponse.MFAToken,
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,
})
}

57
openapi/user/types.go
View file

@ -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
@ -173,29 +180,33 @@ 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"`
ExpiresIn int `json:"expires_in,omitempty"`
RefreshTokenExpiresIn int `json:"refresh_token_expires_in,omitempty"`
TokenType string `json:"token_type,omitempty"`
MFAEnabled bool `json:"mfa_enabled,omitempty"`
Scope string `json:"scope,omitempty"`
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"`
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"`
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"`
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 LoginStatus `json:"status,omitempty"`
}
// Built-in preset mapping types