Add MFA functionality to the examplebroker

internal/brokers/examplebroker/broker.go

@@ -142,11 +142,21 @@ func New(name string) (b *Broker, fullName, brandIcon string) {
 // NewSession creates a new session for the specified user.
 func (b *Broker) NewSession(ctx context.Context, username, lang string) (sessionID, encryptionKey string, err error) {
 	sessionID = uuid.New().String()
-	b.currentSessionsMu.Lock()
-	b.currentSessions[sessionID] = sessionInfo{
+	info := sessionInfo{
 		username: username,
 		lang:     lang,
 	}
+	switch username {
+	case "user-mfa":
+		info.neededMfa = 2
+	case "user-needs-reset":
+		info.needsPasswdReset = true
+	case "user-can-reset":
+		info.canPasswdReset = true
+	}
+
+	b.currentSessionsMu.Lock()
+	b.currentSessions[sessionID] = info
 	b.currentSessionsMu.Unlock()
 	return sessionID, brokerEncryptionKey, nil
 }
@@ -159,6 +169,58 @@ func (b *Broker) GetAuthenticationModes(ctx context.Context, sessionID string, s
 	}
 
 	//var candidatesAuthenticationModes []map[string]string
+	allModes := b.getSupportedModes(sessionInfo, supportedUILayouts)
+
+	// If the user needs mfa, we remove the last used mode from the list of available modes.
+	if sessionInfo.currentMfaStep > 0 && sessionInfo.currentMfaStep < sessionInfo.neededMfa {
+		allModes = b.getMfaModes(sessionInfo, sessionInfo.allModes)
+	}
+	// If the user needs or can reset the password, we only show those authentication modes.
+	if sessionInfo.currentMfaStep > 0 && sessionInfo.needsPasswdReset || sessionInfo.canPasswdReset {
+		allModes = b.getPasswdResetModes(sessionInfo, sessionInfo.allModes)
+	}
+
+	b.userLastSelectedModeMu.Lock()
+	lastSelection := b.userLastSelectedMode[sessionInfo.username]
+	b.userLastSelectedModeMu.Unlock()
+	// Sort in preference order. We want by default password as first and potentially last selection too.
+	if _, exists := allModes[lastSelection]; !exists {
+		lastSelection = ""
+	}
+
+	var allModeIDs []string
+	for n := range allModes {
+		if n == "password" || n == lastSelection {
+			continue
+		}
+		allModeIDs = append(allModeIDs, n)
+	}
+	sort.Strings(allModeIDs)
+
+	if _, exists := allModes["password"]; exists {
+		allModeIDs = append([]string{"password"}, allModeIDs...)
+	}
+	if lastSelection != "" && lastSelection != "password" {
+		allModeIDs = append([]string{lastSelection}, allModeIDs...)
+	}
+
+	for _, id := range allModeIDs {
+		authMode := allModes[id]
+		authenticationModes = append(authenticationModes, map[string]string{
+			"id":    id,
+			"label": authMode["selection_label"],
+		})
+	}
+	sessionInfo.allModes = allModes
+
+	if err := b.updateSession(sessionID, sessionInfo); err != nil {
+		return nil, err
+	}
+
+	return authenticationModes, nil
+}
+
+func (b *Broker) getSupportedModes(sessionInfo sessionInfo, supportedUILayouts []map[string]string) map[string]map[string]string {
 	allModes := make(map[string]map[string]string)
 	for _, layout := range supportedUILayouts {
 		switch layout["type"] {
@@ -269,50 +331,60 @@ func (b *Broker) GetAuthenticationModes(ctx context.Context, sessionID string, s
 
 		case "webview":
 			// This broker does not support webview
+
+		case "newpassword":
+			if layout["entry"] != "" {
+				allModes["mandatoryreset"] = map[string]string{
+					"selection_label": "Password reset",
+					"ui": mapToJSON(map[string]string{
+						"type":  "newpassword",
+						"label": "Enter your new password",
+						"entry": "chars_password",
+					}),
+				}
+
+				if layout["skip-button"] != "" {
+					allModes["optionalreset"] = map[string]string{
+						"selection_label": "Password reset",
+						"ui": mapToJSON(map[string]string{
+							"type":        "newpassword",
+							"label":       "Enter your new password",
+							"entry":       "chars_password",
+							"skip-button": "Skip",
+						}),
+					}
+				}
+			}
 		}
 	}
 
-	// Sort in preference order. We want by default password as first and potentially last selection too.
-	b.userLastSelectedModeMu.Lock()
-	lastSelection := b.userLastSelectedMode[sessionInfo.username]
-	if _, exists := allModes[lastSelection]; !exists {
-		lastSelection = ""
-	}
-	b.userLastSelectedModeMu.Unlock()
+	return allModes
+}
 
-	var allModeIDs []string
-	for n := range allModes {
-		if n == "password" || n == lastSelection {
-			continue
-		}
-		allModeIDs = append(allModeIDs, n)
+func (b *Broker) getMfaModes(info sessionInfo, supportedModes map[string]map[string]string) map[string]map[string]string {
+	mfaModes := make(map[string]map[string]string)
+
+	if authMode, exists := supportedModes["phoneack1"]; exists && info.selectedMode != "phoneack1" {
+		mfaModes["phoneack1"] = authMode
 	}
-	sort.Strings(allModeIDs)
-	if lastSelection != "" && lastSelection != "password" {
-		allModeIDs = append([]string{lastSelection, "password"}, allModeIDs...)
-	} else {
-		allModeIDs = append([]string{"password"}, allModeIDs...)
+	if authMode, exists := supportedModes["totp_with_button"]; exists && info.selectedMode != "totp_with_button" {
+		mfaModes["totp_with_button"] = authMode
 	}
-
-	for _, id := range allModeIDs {
-		authMode := allModes[id]
-		authenticationModes = append(authenticationModes, map[string]string{
-			"id":    id,
-			"label": authMode["selection_label"],
-		})
+	if authMode, exists := supportedModes["fidodevice1"]; exists && info.selectedMode != "fidodevice1" {
+		mfaModes["fidodevice1"] = authMode
 	}
-	sessionInfo.allModes = allModes
+	return mfaModes
+}
 
-	// Checks if the session was ended in the meantime, otherwise we would just accidentally create a new one.
-	if _, err := b.sessionInfo(sessionID); err != nil {
-		return nil, err
+func (b *Broker) getPasswdResetModes(info sessionInfo, supportedModes map[string]map[string]string) map[string]map[string]string {
+	passwdResetModes := make(map[string]map[string]string)
+	if authMode, exists := supportedModes["mandatoryreset"]; info.needsPasswdReset && exists {
+		passwdResetModes["mandatoryreset"] = authMode
 	}
-
-	b.currentSessionsMu.Lock()
-	defer b.currentSessionsMu.Unlock()
-	b.currentSessions[sessionID] = sessionInfo
-
-	return authenticationModes, nil
+	if authMode, exists := supportedModes["optionalreset"]; info.canPasswdReset && exists {
+		passwdResetModes["optionalreset"] = authMode
+	}
+	return passwdResetModes
 }
 
 // SelectAuthenticationMode returns the UI layout information for the selected authentication mode.
@@ -395,6 +467,21 @@ func (b *Broker) IsAuthorized(ctx context.Context, sessionID, authenticationData
 	}()
 
 	access, data, err = b.handleIsAuthorized(b.isAuthorizedCalls[sessionID].ctx, sessionInfo, authData)
+	if access == responses.AuthNext {
+		sessionInfo.currentMfaStep++
+		switch sessionInfo.username {
+		case "user-needs-reset":
+			sessionInfo.needsPasswdReset = true
+		case "user-can-reset":
+			sessionInfo.canPasswdReset = true
+		}
+	}
+	if access == responses.AuthRetry {
+		sessionInfo.currentAttempt++
+		if sessionInfo.currentAttempt >= maxAttempts {
+			access = responses.AuthDenied
+		}
+	}
 
 	// Store last successful authentication mode for this user in the broker.
 	b.userLastSelectedModeMu.Lock()
@@ -414,13 +501,21 @@ func (b *Broker) handleIsAuthorized(ctx context.Context, sessionInfo sessionInfo
 	// Take into account the cancellation.
 	switch sessionInfo.selectedMode {
 	case "password":
+		switch sessionInfo.username {
+		case "user-mfa":
+			fallthrough
+		case "user-needs-reset":
+			fallthrough
+		case "user-can-reset":
+			return responses.AuthNext, "", nil
+		}
 		if authData["challenge"] != "goodpass" {
-			return responses.AuthDenied, "", nil
+			return responses.AuthRetry, "", nil
 		}
 
 	case "pincode":
 		if authData["challenge"] != "4242" {
-			return responses.AuthDenied, "", nil
+			return responses.AuthRetry, "", nil
 		}
 
 	case "totp_with_button", "totp":