ai-workspace-services/accounts
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Adjust login MFA handling for new users (#401)

Browse Source
This commit is contained in:
shenlan 2025-10-05 08:00:16 +08:00 committed by GitHub
parent 060c11107e
commit e3c3780496
2 changed files with 62 additions and 43 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

66
account/api/api.go
View File

@ -527,37 +527,39 @@ func (h *handler) login(c *gin.Context) {
return
}
if !user.MFAEnabled {
challengeToken, err := h.createMFAChallenge(user.ID)
if err != nil {
respondError(c, http.StatusInternalServerError, "mfa_challenge_failed", "failed to prepare mfa challenge")
if user.MFAEnabled {
if totpCode == "" {
respondError(c, http.StatusBadRequest, "mfa_code_required", "totp code is required")
return
}
c.JSON(http.StatusUnauthorized, gin.H{
"error": "mfa_setup_required",
"mfaToken": challengeToken,
"user": sanitizeUser(user),
valid, err := totp.ValidateCustom(totpCode, user.MFATOTPSecret, time.Now().UTC(), totp.ValidateOpts{
Period: 30,
Skew: 1,
Digits: otp.DigitsSix,
Algorithm: otp.AlgorithmSHA1,
})
return
}
if err != nil {
respondError(c, http.StatusInternalServerError, "invalid_mfa_code", "invalid totp code")
return
}
if !valid {
respondError(c, http.StatusUnauthorized, "invalid_mfa_code", "invalid totp code")
return
}
if totpCode == "" {
respondError(c, http.StatusBadRequest, "mfa_code_required", "totp code is required")
return
}
token, expiresAt, err := h.createSession(user.ID)
if err != nil {
respondError(c, http.StatusInternalServerError, "session_creation_failed", "failed to create session")
return
}
valid, err := totp.ValidateCustom(totpCode, user.MFATOTPSecret, time.Now().UTC(), totp.ValidateOpts{
Period: 30,
Skew: 1,
Digits: otp.DigitsSix,
Algorithm: otp.AlgorithmSHA1,
})
if err != nil {
respondError(c, http.StatusInternalServerError, "invalid_mfa_code", "invalid totp code")
return
}
if !valid {
respondError(c, http.StatusUnauthorized, "invalid_mfa_code", "invalid totp code")
c.JSON(http.StatusOK, gin.H{
"message": "login successful",
"token": token,
"expiresAt": expiresAt.UTC(),
"user": sanitizeUser(user),
})
return
}
@ -567,12 +569,20 @@ func (h *handler) login(c *gin.Context) {
return
}
c.JSON(http.StatusOK, gin.H{
response := gin.H{
"message": "login successful",
"token": token,
"expiresAt": expiresAt.UTC(),
"user": sanitizeUser(user),
})
}
if challengeToken, err := h.createMFAChallenge(user.ID); err != nil {
slog.Error("failed to create mfa challenge during login", "err", err, "userID", user.ID)
} else {
response["mfaToken"] = challengeToken
}
c.JSON(http.StatusOK, response)
}
func (h *handler) findUserByIdentifier(ctx context.Context, identifier string) (*store.User, error) {

39
account/api/api_test.go
View File

@ -295,15 +295,18 @@ func TestMFATOTPFlow(t *testing.T) {
rr = httptest.NewRecorder()
router.ServeHTTP(rr, req)
if rr.Code != http.StatusUnauthorized {
t.Fatalf("expected login to require mfa setup, got %d", rr.Code)
if rr.Code != http.StatusOK {
t.Fatalf("expected login success for new user, got %d: %s", rr.Code, rr.Body.String())
}
resp := decodeResponse(t, rr)
if resp.Error != "mfa_setup_required" {
t.Fatalf("expected mfa_setup_required error, got %q", resp.Error)
if resp.Token == "" {
t.Fatalf("expected session token in login response")
}
if resp.MFAToken == "" {
t.Fatalf("expected mfa token in response")
t.Fatalf("expected mfa token in login response")
}
if resp.User == nil {
t.Fatalf("expected user object in login response")
}
provisionPayload := map[string]string{
@ -525,11 +528,14 @@ func TestDisableMFA(t *testing.T) {
req.Header.Set("Content-Type", "application/json")
rr = httptest.NewRecorder()
router.ServeHTTP(rr, req)
if rr.Code != http.StatusUnauthorized {
t.Fatalf("expected login to require mfa setup, got %d: %s", rr.Code, rr.Body.String())
if rr.Code != http.StatusOK {
t.Fatalf("expected login success for new user, got %d: %s", rr.Code, rr.Body.String())
}
resp := decodeResponse(t, rr)
if resp.Token == "" {
t.Fatalf("expected session token in login response")
}
if resp.MFAToken == "" {
t.Fatalf("expected mfa token in login response")
}
@ -621,12 +627,15 @@ func TestDisableMFA(t *testing.T) {
req.Header.Set("Content-Type", "application/json")
rr = httptest.NewRecorder()
router.ServeHTTP(rr, req)
if rr.Code != http.StatusUnauthorized {
t.Fatalf("expected login to require setup again, got %d: %s", rr.Code, rr.Body.String())
if rr.Code != http.StatusOK {
t.Fatalf("expected login success after disable, got %d: %s", rr.Code, rr.Body.String())
}
resp = decodeResponse(t, rr)
if resp.Error != "mfa_setup_required" {
t.Fatalf("expected mfa_setup_required error after disable, got %q", resp.Error)
if resp.Token == "" {
t.Fatalf("expected session token after disable login")
}
if resp.MFAToken == "" {
t.Fatalf("expected mfa token after disable login")
}
}
@ -758,12 +767,12 @@ func TestPasswordResetFlow(t *testing.T) {
req.Header.Set("Content-Type", "application/json")
rr = httptest.NewRecorder()
router.ServeHTTP(rr, req)
if rr.Code != http.StatusUnauthorized {
t.Fatalf("expected login to prompt for mfa setup, got %d: %s", rr.Code, rr.Body.String())
if rr.Code != http.StatusOK {
t.Fatalf("expected login success after password reset, got %d: %s", rr.Code, rr.Body.String())
}
resp = decodeResponse(t, rr)
if resp.Error != "mfa_setup_required" {
t.Fatalf("expected mfa_setup_required after password reset, got %q", resp.Error)
if resp.Token == "" {
t.Fatalf("expected session token after password reset")
}
loginPayload["password"] = registerPayload["password"]