Addressing comments

ubuntu · Aug 30, 2023 · c7f4073 · c7f4073
1 parent eb58b32
commit c7f4073
Showing 1 changed file with 52 additions and 43 deletions.
diff --git a/internal/brokers/examplebroker/broker.go b/internal/brokers/examplebroker/broker.go
@@ -21,15 +21,22 @@ import (
 
 const maxAttempts int = 5
 
+type passwdReset int
+
+const (
+	noReset passwdReset = iota
+	canReset
+	mustReset
+)
+
 type sessionInfo struct {
-	username       string
-	selectedMode   string
-	lang           string
-	allModes       map[string]map[string]string
-	currentAttempt int
+	username        string
+	selectedMode    string
+	lang            string
+	allModes        map[string]map[string]string
+	attemptsPerMode map[string]int
 
-	needsPasswdReset bool
-	canPasswdReset   bool
+	pwdChange passwdReset
 
 	neededMfa      int
 	currentMfaStep int
@@ -143,16 +150,20 @@ func New(name string) (b *Broker, fullName, brandIcon string) {
 func (b *Broker) NewSession(ctx context.Context, username, lang string) (sessionID, encryptionKey string, err error) {
 	sessionID = uuid.New().String()
 	info := sessionInfo{
-		username: username,
-		lang:     lang,
+		username:        username,
+		lang:            lang,
+		pwdChange:       noReset,
+		attemptsPerMode: make(map[string]int),
 	}
 	switch username {
 	case "user-mfa":
-		info.neededMfa = 2
+		info.neededMfa = 3
 	case "user-needs-reset":
-		info.needsPasswdReset = true
+		info.neededMfa = 1
+		info.pwdChange = mustReset
 	case "user-can-reset":
-		info.canPasswdReset = true
+		info.neededMfa = 1
+		info.pwdChange = canReset
 	}
 
 	b.currentSessionsMu.Lock()
@@ -176,7 +187,7 @@ func (b *Broker) GetAuthenticationModes(ctx context.Context, sessionID string, s
 		allModes = 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 {
+	if sessionInfo.currentMfaStep > 0 && sessionInfo.pwdChange != noReset {
 		allModes = getPasswdResetModes(sessionInfo, sessionInfo.allModes)
 	}
 
@@ -337,7 +348,7 @@ func getSupportedModes(sessionInfo sessionInfo, supportedUILayouts []map[string]
 				break
 			}
 			allModes["mandatoryreset"] = map[string]string{
-				"selection_label": "Password reset",
+				"selection_label": "Password reset (3 days until mandatory)",
 				"ui": mapToJSON(map[string]string{
 					"type":  "newpassword",
 					"label": "Enter your new password",
@@ -365,26 +376,29 @@ func getSupportedModes(sessionInfo sessionInfo, supportedUILayouts []map[string]
 func 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
-	}
-	if authMode, exists := supportedModes["totp_with_button"]; exists && info.selectedMode != "totp_with_button" {
-		mfaModes["totp_with_button"] = authMode
-	}
-	if authMode, exists := supportedModes["fidodevice1"]; exists && info.selectedMode != "fidodevice1" {
-		mfaModes["fidodevice1"] = authMode
+	supportedMfaModes := []string{"phoneack1", "totp_with_button", "fidodevice1"}
+	for _, mode := range supportedMfaModes {
+		if _, exists := supportedModes[mode]; exists && info.selectedMode != mode {
+			mfaModes[mode] = supportedModes[mode]
+		}
 	}
 	return mfaModes
 }
 
 func 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
+
+	var mode string
+	switch info.pwdChange {
+	case canReset:
+		mode = "optionalreset"
+	case mustReset:
+		mode = "mandatoryreset"
 	}
-	if authMode, exists := supportedModes["optionalreset"]; info.canPasswdReset && exists {
-		passwdResetModes["optionalreset"] = authMode
+	if authMode, exists := supportedModes[mode]; exists {
+		passwdResetModes[mode] = authMode
 	}
+
 	return passwdResetModes
 }
 
@@ -467,18 +481,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++
+	if access == responses.AuthAllowed {
 		switch sessionInfo.username {
 		case "user-needs-reset":
-			sessionInfo.needsPasswdReset = true
+			fallthrough
 		case "user-can-reset":
-			sessionInfo.canPasswdReset = true
+			fallthrough
+		case "user-mfa":
+			if sessionInfo.currentMfaStep < sessionInfo.neededMfa {
+				sessionInfo.currentMfaStep++
+				access = responses.AuthNext
+			}
 		}
-	}
-	if access == responses.AuthRetry {
-		sessionInfo.currentAttempt++
-		if sessionInfo.currentAttempt >= maxAttempts {
+	} else if access == responses.AuthRetry {
+		sessionInfo.attemptsPerMode[sessionInfo.selectedMode]++
+		if sessionInfo.attemptsPerMode[sessionInfo.selectedMode] >= maxAttempts {
 			access = responses.AuthDenied
 		}
 	}
@@ -501,14 +518,6 @@ 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.AuthRetry, "", nil
 		}
@@ -521,7 +530,7 @@ func (b *Broker) handleIsAuthorized(ctx context.Context, sessionInfo sessionInfo
 	case "totp_with_button", "totp":
 		wantedCode := sessionInfo.allModes[sessionInfo.selectedMode]["wantedCode"]
 		if authData["challenge"] != wantedCode {
-			return responses.AuthDenied, "", nil
+			return responses.AuthRetry, "", nil
 		}
 
 	case "phoneack1":