From fd7dc661684ede3b039294ddca980d480a774101 Mon Sep 17 00:00:00 2001
From: Hasanthi Dissanayake <hasanthi.disanayakka@yahoo.com>
Date: Tue, 14 Nov 2023 21:16:11 +0530
Subject: [PATCH] Handle errors in the totp flow

---
 .../authenticator/totp/TOTPAuthenticator.java | 34 +++++++++++++++++++
 pom.xml                                       |  2 +-
 2 files changed, 35 insertions(+), 1 deletion(-)

diff --git a/component/authenticator/src/main/java/org/wso2/carbon/identity/application/authenticator/totp/TOTPAuthenticator.java b/component/authenticator/src/main/java/org/wso2/carbon/identity/application/authenticator/totp/TOTPAuthenticator.java
index a8a8e0d5b..eb24f1a0d 100644
--- a/component/authenticator/src/main/java/org/wso2/carbon/identity/application/authenticator/totp/TOTPAuthenticator.java
+++ b/component/authenticator/src/main/java/org/wso2/carbon/identity/application/authenticator/totp/TOTPAuthenticator.java
@@ -37,6 +37,7 @@
 import org.wso2.carbon.identity.application.authentication.framework.exception.UserIdNotFoundException;
 import org.wso2.carbon.identity.application.authentication.framework.model.AuthenticatedUser;
 import org.wso2.carbon.identity.application.authentication.framework.model.AuthenticatorData;
+import org.wso2.carbon.identity.application.authentication.framework.model.AuthenticatorMessage;
 import org.wso2.carbon.identity.application.authentication.framework.model.AuthenticatorParamMetadata;
 import org.wso2.carbon.identity.application.authentication.framework.util.FrameworkConstants;
 import org.wso2.carbon.identity.application.authentication.framework.util.FrameworkUtils;
@@ -98,6 +99,7 @@ public class TOTPAuthenticator extends AbstractApplicationAuthenticator
     private static final long serialVersionUID = 2009231028659744926L;
     private static final Log log = LogFactory.getLog(TOTPAuthenticator.class);
     private static final String IS_API_BASED = "IS_API_BASED";
+    private static final String AUTHENTICATOR_MESSAGE = "authenticatorMessage";
 
     /**
      * Check whether token or action are in request.
@@ -304,6 +306,14 @@ protected void initiateAuthenticationRequest(HttpServletRequest request,
                                     String.valueOf(Math.round((double) timeToUnlock / 1000 / 60)));
                         }
                         errorParam = buildErrorParamString(paramMap);
+                        Map<String, String> messageContext = getMessageContext("lockedReason",
+                                String.valueOf(reason));
+                        AuthenticatorMessage authenticatorMessage =
+                                new AuthenticatorMessage(FrameworkConstants.AuthenticatorMessageType.ERROR,
+                                        UserCoreConstants.ErrorCode.USER_IS_LOCKED,
+                                        "The authenticated user account is locked.",
+                                        messageContext);
+                        setAuthenticatorMessage(authenticatorMessage, context);
                     }
                 }
             }
@@ -451,6 +461,13 @@ private long getUnlockTimeInMilliSeconds(AuthenticatedUser authenticatedUser) th
         return Long.parseLong(claimValues.get(TOTPAuthenticatorConstants.ACCOUNT_UNLOCK_TIME_CLAIM));
     }
 
+    private static Map<String, String> getMessageContext(String key, String value) {
+
+        Map <String,String> messageContext = new HashMap<>();
+        messageContext.put(key, value);
+        return messageContext;
+    }
+
     private String buildTOTPLoginPageURL(AuthenticationContext context, String username, String retryParam,
                                          String errorParam, String multiOptionURI)
             throws AuthenticationFailedException, URISyntaxException, URLBuilderException {
@@ -671,11 +688,22 @@ private void validateAccountLockStatusForLocalUser(AuthenticationContext context
             }
             IdentityErrorMsgContext customErrorMessageContext = new IdentityErrorMsgContext(
                     UserCoreConstants.ErrorCode.USER_IS_LOCKED + ":" + accountLockedReason);
+            AuthenticatorMessage authenticatorMessage =
+                    new AuthenticatorMessage(FrameworkConstants.AuthenticatorMessageType.ERROR,
+                            UserCoreConstants.ErrorCode.USER_IS_LOCKED,
+                            "The authenticated user account is locked.",
+                            null);
+            setAuthenticatorMessage(authenticatorMessage, context);
             IdentityUtil.setIdentityErrorMsg(customErrorMessageContext);
             throw new AuthenticationFailedException(errorMessage);
         }
     }
 
+    private static void setAuthenticatorMessage(AuthenticatorMessage errorMessage, AuthenticationContext context) {
+
+        context.setProperty(AUTHENTICATOR_MESSAGE, errorMessage);
+    }
+
     /**
      * Check whether status of retrying authentication.
      *
@@ -1018,6 +1046,12 @@ private void handleTotpVerificationFail(AuthenticationContext context) throws Au
             setUserClaimValues(authenticatedUser, updatedClaims);
             String errorMessage = String.format("User account: %s is locked.", (LoggerUtils.isLogMaskingEnable ?
                     LoggerUtils.getMaskedContent(authenticatedUser.getUserName()) : authenticatedUser.getUserName()));
+            AuthenticatorMessage authenticatorMessage =
+                    new AuthenticatorMessage(FrameworkConstants.AuthenticatorMessageType.ERROR,
+                            UserCoreConstants.ErrorCode.USER_IS_LOCKED,
+                            errorMessage,
+                            null);
+            setAuthenticatorMessage(authenticatorMessage, context);
             IdentityErrorMsgContext customErrorMessageContext = new IdentityErrorMsgContext(
                     UserCoreConstants.ErrorCode.USER_IS_LOCKED + ":" +
                             TOTPAuthenticatorConstants.MAX_TOTP_ATTEMPTS_EXCEEDED);
diff --git a/pom.xml b/pom.xml
index 01ec96207..e9349143b 100644
--- a/pom.xml
+++ b/pom.xml
@@ -158,7 +158,7 @@
         <carbon.commons.version>4.8.7</carbon.commons.version>
         <carbon.commons.imp.pkg.version>[4.4.0, 5.0.0)</carbon.commons.imp.pkg.version>
         <!--Carbon identity version-->
-        <carbon.identity.framework.version>5.25.491</carbon.identity.framework.version>
+        <carbon.identity.framework.version>5.25.503</carbon.identity.framework.version>
 
         <org.wso2.carbon.identity.organization.management.core.version>1.0.0
         </org.wso2.carbon.identity.organization.management.core.version>