From 85c914ef4750585671ca7ec4796bbd7fb4ce2bd5 Mon Sep 17 00:00:00 2001
From: Janak Amarasena <janakamarasena@gmail.com>
Date: Tue, 17 Oct 2023 17:20:34 +0530
Subject: [PATCH] Add check to determine API based auth flow

---
 .../handler/step/impl/DefaultStepHandler.java | 37 +++++++++++++------
 .../service/AuthServiceRequestWrapper.java    |  1 +
 .../framework/util/FrameworkConstants.java    |  1 +
 3 files changed, 27 insertions(+), 12 deletions(-)

diff --git a/components/authentication-framework/org.wso2.carbon.identity.application.authentication.framework/src/main/java/org/wso2/carbon/identity/application/authentication/framework/handler/step/impl/DefaultStepHandler.java b/components/authentication-framework/org.wso2.carbon.identity.application.authentication.framework/src/main/java/org/wso2/carbon/identity/application/authentication/framework/handler/step/impl/DefaultStepHandler.java
index 3521dc3e39fd..ffed4bd651de 100644
--- a/components/authentication-framework/org.wso2.carbon.identity.application.authentication.framework/src/main/java/org/wso2/carbon/identity/application/authentication/framework/handler/step/impl/DefaultStepHandler.java
+++ b/components/authentication-framework/org.wso2.carbon.identity.application.authentication.framework/src/main/java/org/wso2/carbon/identity/application/authentication/framework/handler/step/impl/DefaultStepHandler.java
@@ -725,18 +725,8 @@ protected void doAuthentication(HttpServletRequest request, HttpServletResponse
             request.setAttribute(FrameworkConstants.RequestParams.FLOW_STATUS, status);
             /* If this is an authentication initiation and the authenticator supports API based authentication
              we need to send the auth initiation data in order to support performing API based authentication.*/
-            if (status == AuthenticatorFlowStatus.INCOMPLETE && authenticator.isAPIBasedAuthenticationSupported()) {
-                authenticator.getAuthInitiationData(context).ifPresent(authInitiationData -> {
-                    List<AuthenticatorData> authInitiationDataList =
-                            (List<AuthenticatorData>) request
-                                    .getAttribute(AuthServiceConstants.AUTH_SERVICE_AUTH_INITIATION_DATA);
-                    if (authInitiationDataList == null) {
-                        authInitiationDataList = new ArrayList<>();
-                        request.setAttribute(AuthServiceConstants.AUTH_SERVICE_AUTH_INITIATION_DATA,
-                                authInitiationDataList);
-                    }
-                    authInitiationDataList.add(authInitiationData);
-                });
+            if (status == AuthenticatorFlowStatus.INCOMPLETE) {
+                handleAPIBasedAuthenticationData(request, authenticator, context);
             }
 
             if (LOG.isDebugEnabled()) {
@@ -1449,4 +1439,27 @@ private String resolveUserResidentOrganization(AuthenticatedUser authenticatedUs
         }
         throw new FrameworkException("User resident organization could not found");
     }
+
+    private void handleAPIBasedAuthenticationData(HttpServletRequest request, ApplicationAuthenticator authenticator,
+                                                  AuthenticationContext context) {
+
+        if (isAPIBasedAuthenticationFlow(request) && authenticator.isAPIBasedAuthenticationSupported()) {
+            authenticator.getAuthInitiationData(context).ifPresent(authInitiationData -> {
+                List<AuthenticatorData> authInitiationDataList =
+                        (List<AuthenticatorData>) request
+                                .getAttribute(AuthServiceConstants.AUTH_SERVICE_AUTH_INITIATION_DATA);
+                if (authInitiationDataList == null) {
+                    authInitiationDataList = new ArrayList<>();
+                    request.setAttribute(AuthServiceConstants.AUTH_SERVICE_AUTH_INITIATION_DATA,
+                            authInitiationDataList);
+                }
+                authInitiationDataList.add(authInitiationData);
+            });
+        }
+    }
+
+    private boolean isAPIBasedAuthenticationFlow(HttpServletRequest request) {
+
+        return Boolean.TRUE.equals(request.getAttribute(FrameworkConstants.IS_API_BASED_AUTH_FLOW));
+    }
 }
diff --git a/components/authentication-framework/org.wso2.carbon.identity.application.authentication.framework/src/main/java/org/wso2/carbon/identity/application/authentication/framework/model/auth/service/AuthServiceRequestWrapper.java b/components/authentication-framework/org.wso2.carbon.identity.application.authentication.framework/src/main/java/org/wso2/carbon/identity/application/authentication/framework/model/auth/service/AuthServiceRequestWrapper.java
index 995c2c3866f1..fbb3acd82882 100644
--- a/components/authentication-framework/org.wso2.carbon.identity.application.authentication.framework/src/main/java/org/wso2/carbon/identity/application/authentication/framework/model/auth/service/AuthServiceRequestWrapper.java
+++ b/components/authentication-framework/org.wso2.carbon.identity.application.authentication.framework/src/main/java/org/wso2/carbon/identity/application/authentication/framework/model/auth/service/AuthServiceRequestWrapper.java
@@ -47,6 +47,7 @@ public AuthServiceRequestWrapper(HttpServletRequest request, Map<String, String[
         this.parameters = parameters;
         setSessionDataKey(parameters);
         skipNonceCookieValidation();
+        this.setAttribute(FrameworkConstants.IS_API_BASED_AUTH_FLOW, true);
     }
 
     @Override
diff --git a/components/authentication-framework/org.wso2.carbon.identity.application.authentication.framework/src/main/java/org/wso2/carbon/identity/application/authentication/framework/util/FrameworkConstants.java b/components/authentication-framework/org.wso2.carbon.identity.application.authentication.framework/src/main/java/org/wso2/carbon/identity/application/authentication/framework/util/FrameworkConstants.java
index 67d860677ca3..46e7f7aac12f 100644
--- a/components/authentication-framework/org.wso2.carbon.identity.application.authentication.framework/src/main/java/org/wso2/carbon/identity/application/authentication/framework/util/FrameworkConstants.java
+++ b/components/authentication-framework/org.wso2.carbon.identity.application.authentication.framework/src/main/java/org/wso2/carbon/identity/application/authentication/framework/util/FrameworkConstants.java
@@ -78,6 +78,7 @@ public abstract class FrameworkConstants {
     public static final String SKIP_NONCE_COOKIE_VALIDATION = "SkipNonceCookieValidation";
     public static final String IS_MULTI_OPS_RESPONSE = "isMultiOptionsResponse";
     public static final String IS_AUTH_FLOW_CONCLUDED = "isAuthFlowConcluded";
+    public static final String IS_API_BASED_AUTH_FLOW = "isAPIBasedAuthFlow";
 
     // This is to support sign-up form to be displayed in the provisioning flow, as when trying to displaying the
     // sign-up form, we validate whether self-sign up is enabled.