Skip 2fa recovery code step if disabled or exists in settings v2

IniZio · Oct 10, 2024 · e292f74 · e292f74
1 parent b58e4d1
commit e292f74
Show file tree

Hide file tree

Showing 6 changed files with 114 additions and 44 deletions.
diff --git a/pkg/auth/handler/webapp/authflowv2/settings_mfa_enter_oob_otp.go b/pkg/auth/handler/webapp/authflowv2/settings_mfa_enter_oob_otp.go
@@ -177,6 +177,8 @@ func (h *AuthflowV2SettingsMFAEnterOOBOTPHandler) ServeHTTP(w http.ResponseWrite
 			return err
 		}
 
+		s := session.GetSession(r.Context())
+
 		tokenString := r.Form.Get("q_token")
 		code := r.Form.Get("x_code")
 
@@ -187,13 +189,29 @@ func (h *AuthflowV2SettingsMFAEnterOOBOTPHandler) ServeHTTP(w http.ResponseWrite
 			return err
 		}
 
-		redirectURI, err := url.Parse(AuthflowV2RouteSettingsMFAViewRecoveryCode)
-		if err != nil {
-			return err
+		var redirectURI *url.URL
+		if output.RecoveryCodesCreated {
+			redirectURI, err = url.Parse(AuthflowV2RouteSettingsMFAViewRecoveryCode)
+			if err != nil {
+				return err
+			}
+			q := redirectURI.Query()
+			q.Set("q_token", output.Token)
+			redirectURI.RawQuery = q.Encode()
+		} else {
+			_, err = h.AccountManagement.FinishAddOOBOTPAuthenticator(s, output.Token, &accountmanagement.FinishAddOOBOTPAuthenticatorInput{})
+			if err != nil {
+				return err
+			}
+
+			redirectURI, err = url.Parse(AuthflowV2RouteSettingsMFA)
+			if err != nil {
+				return err
+			}
+			q := redirectURI.Query()
+			q.Set("q_token", output.Token)
+			redirectURI.RawQuery = q.Encode()
 		}
-		q := redirectURI.Query()
-		q.Set("q_token", output.Token)
-		redirectURI.RawQuery = q.Encode()
 
 		result := webapp.Result{RedirectURI: redirectURI.String()}
 		result.WriteResponse(w, r)

diff --git a/pkg/auth/handler/webapp/authflowv2/settings_mfa_enter_totp.go b/pkg/auth/handler/webapp/authflowv2/settings_mfa_enter_totp.go
@@ -114,16 +114,30 @@ func (h *AuthflowV2SettingsMFAEnterTOTPHandler) ServeHTTP(w http.ResponseWriter,
 			return err
 		}
 
-		redirectURI, err := url.Parse(AuthflowV2RouteSettingsMFAViewRecoveryCode)
-		if err != nil {
-			return err
+		var redirectURI *url.URL
+		if output.RecoveryCodesCreated {
+			redirectURI, err = url.Parse(AuthflowV2RouteSettingsMFAViewRecoveryCode)
+			if err != nil {
+				return err
+			}
+			q := redirectURI.Query()
+			q.Set("q_token", output.Token)
+			redirectURI.RawQuery = q.Encode()
+		} else {
+			_, err = h.AccountManagement.FinishAddTOTPAuthenticator(s, output.Token, &accountmanagement.FinishAddTOTPAuthenticatorInput{})
+			if err != nil {
+				return err
+			}
+
+			redirectURI, err = url.Parse(AuthflowV2RouteSettingsMFA)
+			if err != nil {
+				return err
+			}
+			q := redirectURI.Query()
+			q.Set("q_token", output.Token)
+			redirectURI.RawQuery = q.Encode()
 		}
 
-		q := redirectURI.Query()
-		q.Set("q_token", output.Token)
-
-		redirectURI.RawQuery = q.Encode()
-
 		result := webapp.Result{RedirectURI: redirectURI.String()}
 		result.WriteResponse(w, r)
 

diff --git a/pkg/lib/accountmanagement/redis_store.go b/pkg/lib/accountmanagement/redis_store.go
@@ -39,9 +39,10 @@ type GenerateTokenOptions struct {
 	IdentityID string
 
 	// AuthenticatorID for updating authenticator
-	AuthenticatorID            string
-	AuthenticatorRecoveryCodes []string
-	AuthenticatorType          model.AuthenticatorType
+	AuthenticatorID                   string
+	AuthenticatorRecoveryCodes        []string
+	AuthenticatorRecoveryCodesCreated bool
+	AuthenticatorType                 model.AuthenticatorType
 
 	// TOTP
 	AuthenticatorTOTPIssuer           string

diff --git a/pkg/lib/accountmanagement/service.go b/pkg/lib/accountmanagement/service.go
@@ -80,6 +80,7 @@ type AuthenticationInfoService interface {
 type MFAService interface {
 	GenerateRecoveryCodes() []string
 	ReplaceRecoveryCodes(userID string, codes []string) ([]*mfa.RecoveryCode, error)
+	ListRecoveryCodes(userID string) ([]*mfa.RecoveryCode, error)
 }
 
 type PasskeyService interface {

diff --git a/pkg/lib/accountmanagement/service_authenticator.go b/pkg/lib/accountmanagement/service_authenticator.go
@@ -210,7 +210,8 @@ type ResumeAddTOTPAuthenticatorInput struct {
 	Code        string
 }
 type ResumeAddTOTPAuthenticatorOutput struct {
-	Token string
+	Token                string
+	RecoveryCodesCreated bool
 }
 
 func (s *Service) ResumeAddTOTPAuthenticator(resolvedSession session.ResolvedSession, tokenString string, input *ResumeAddTOTPAuthenticatorInput) (output *ResumeAddTOTPAuthenticatorOutput, err error) {
@@ -265,22 +266,24 @@ func (s *Service) ResumeAddTOTPAuthenticator(resolvedSession session.ResolvedSes
 		return
 	}
 
-	recoveryCodes := s.MFA.GenerateRecoveryCodes()
+	recoveryCodes, recoveryCodesCreated, err := s.generateRecoveryCodes(userID)
 
 	newToken, err := s.Store.GenerateToken(GenerateTokenOptions{
-		UserID:                       userID,
-		AuthenticatorType:            model.AuthenticatorType(token.Authenticator.AuthenticatorType),
-		AuthenticatorTOTPDisplayName: input.DisplayName,
-		AuthenticatorTOTPSecret:      token.Authenticator.TOTPSecret,
-		AuthenticatorTOTPVerified:    true,
-		AuthenticatorRecoveryCodes:   recoveryCodes,
+		UserID:                            userID,
+		AuthenticatorRecoveryCodes:        recoveryCodes,
+		AuthenticatorRecoveryCodesCreated: recoveryCodesCreated,
+		AuthenticatorType:                 model.AuthenticatorType(token.Authenticator.AuthenticatorType),
+		AuthenticatorTOTPDisplayName:      input.DisplayName,
+		AuthenticatorTOTPSecret:           token.Authenticator.TOTPSecret,
+		AuthenticatorTOTPVerified:         true,
 	})
 	if err != nil {
 		return
 	}
 
 	output = &ResumeAddTOTPAuthenticatorOutput{
-		Token: newToken,
+		Token:                newToken,
+		RecoveryCodesCreated: recoveryCodesCreated,
 	}
 	return
 }
@@ -331,9 +334,11 @@ func (s *Service) FinishAddTOTPAuthenticator(resolvedSession session.ResolvedSes
 			return err
 		}
 
-		_, err = s.MFA.ReplaceRecoveryCodes(userID, token.Authenticator.RecoveryCodes)
-		if err != nil {
-			return err
+		if token.Authenticator.RecoveryCodesCreated {
+			_, err = s.MFA.ReplaceRecoveryCodes(userID, token.Authenticator.RecoveryCodes)
+			if err != nil {
+				return err
+			}
 		}
 
 		return nil
@@ -516,7 +521,8 @@ type ResumeAddOOBOTPAuthenticatorInput struct {
 	Code string
 }
 type ResumeAddOOBOTPAuthenticatorOutput struct {
-	Token string
+	Token                string
+	RecoveryCodesCreated bool
 }
 
 func (s *Service) ResumeAddOOBOTPAuthenticator(resolvedSession session.ResolvedSession, tokenString string, input *ResumeAddOOBOTPAuthenticatorInput) (output *ResumeAddOOBOTPAuthenticatorOutput, err error) {
@@ -548,22 +554,24 @@ func (s *Service) ResumeAddOOBOTPAuthenticator(resolvedSession session.ResolvedS
 		return
 	}
 
-	recoveryCodes := s.MFA.GenerateRecoveryCodes()
+	recoveryCodes, recoveryCodesCreated, err := s.generateRecoveryCodes(userID)
 
 	newToken, err := s.Store.GenerateToken(GenerateTokenOptions{
-		UserID:                      userID,
-		AuthenticatorRecoveryCodes:  recoveryCodes,
-		AuthenticatorType:           model.AuthenticatorType(token.Authenticator.AuthenticatorType),
-		AuthenticatorOOBOTPChannel:  token.Authenticator.OOBOTPChannel,
-		AuthenticatorOOBOTPTarget:   token.Authenticator.OOBOTPTarget,
-		AuthenticatorOOBOTPVerified: true,
+		UserID:                            userID,
+		AuthenticatorRecoveryCodes:        recoveryCodes,
+		AuthenticatorRecoveryCodesCreated: recoveryCodesCreated,
+		AuthenticatorType:                 model.AuthenticatorType(token.Authenticator.AuthenticatorType),
+		AuthenticatorOOBOTPChannel:        token.Authenticator.OOBOTPChannel,
+		AuthenticatorOOBOTPTarget:         token.Authenticator.OOBOTPTarget,
+		AuthenticatorOOBOTPVerified:       true,
 	})
 	if err != nil {
 		return
 	}
 
 	output = &ResumeAddOOBOTPAuthenticatorOutput{
-		Token: newToken,
+		Token:                newToken,
+		RecoveryCodesCreated: recoveryCodesCreated,
 	}
 	return
 }
@@ -624,9 +632,11 @@ func (s *Service) FinishAddOOBOTPAuthenticator(resolvedSession session.ResolvedS
 			return err
 		}
 
-		_, err = s.MFA.ReplaceRecoveryCodes(userID, token.Authenticator.RecoveryCodes)
-		if err != nil {
-			return err
+		if token.Authenticator.RecoveryCodesCreated {
+			_, err = s.MFA.ReplaceRecoveryCodes(userID, token.Authenticator.RecoveryCodes)
+			if err != nil {
+				return err
+			}
 		}
 
 		return nil
@@ -670,3 +680,26 @@ func (s *Service) DeleteOOBOTPAuthenticator(resolvedSession session.ResolvedSess
 	}
 	return
 }
+
+func (s *Service) generateRecoveryCodes(userID string) (recoveryCodes []string, isCreated bool, err error) {
+	if *s.Config.Authentication.RecoveryCode.Disabled {
+		return nil, false, nil
+	}
+
+	err = s.Database.WithTx(func() error {
+		existing, err := s.MFA.ListRecoveryCodes(userID)
+		if err != nil {
+			return err
+		}
+
+		if len(existing) == 0 {
+			isCreated = true
+			recoveryCodes = s.MFA.GenerateRecoveryCodes()
+			return nil
+		}
+
+		return nil
+	})
+
+	return recoveryCodes, isCreated, err
+}
diff --git a/pkg/lib/accountmanagement/token.go b/pkg/lib/accountmanagement/token.go
@@ -35,9 +35,12 @@ type TokenIdentity struct {
 }
 
 type TokenAuthenticator struct {
-	AuthenticatorID   string   `json:"authenticator_id,omitempty"`
-	AuthenticatorType string   `json:"authenticator_type,omitempty"`
-	RecoveryCodes     []string `json:"recovery_codes,omitempty"`
+	AuthenticatorID   string `json:"authenticator_id,omitempty"`
+	AuthenticatorType string `json:"authenticator_type,omitempty"`
+
+	// Recovery Codes
+	RecoveryCodes        []string `json:"recovery_codes,omitempty"`
+	RecoveryCodesCreated bool     `json:"recovery_codes_created,omitempty"`
 
 	// TOTP
 	TOTPIssuer           string `json:"totp_issuer,omitempty"`