Avoid rehandling api based auth response

Author: janakamarasena

components/org.wso2.carbon.identity.oauth.endpoint/src/main/java/org/wso2/carbon/identity/oauth/endpoint/authz/OAuth2AuthzEndpoint.java

@@ -4341,14 +4341,19 @@ private Response handleApiBasedAuthenticationResponse(OAuthMessage oAuthMessage,
                 ObjectMapper objectMapper = new ObjectMapper();
                 objectMapper.setSerializationInclusion(JsonInclude.Include.NON_EMPTY);
                 String jsonString = objectMapper.writeValueAsString(authResponse);
+                oAuthMessage.getRequest().setAttribute(IS_API_BASED_AUTH_HANDLED, true);
                 return Response.ok().entity(jsonString).build();
 
             } else {
-                String location = oauthResponse.getMetadata().get("Location").get(0).toString();
-                if (StringUtils.isNotBlank(location)) {
-                    Map<String, String> queryParams = getQueryParamsFromUrl(location);
-                    String jsonPayload = new Gson().toJson(queryParams);
-                    return Response.status(HttpServletResponse.SC_OK).entity(jsonPayload).build();
+                List<Object> locationHeader = oauthResponse.getMetadata().get("Location");
+                if (CollectionUtils.isNotEmpty(locationHeader)) {
+                    String location = locationHeader.get(0).toString();
+                    if (StringUtils.isNotBlank(location)) {
+                        Map<String, String> queryParams = getQueryParamsFromUrl(location);
+                        String jsonPayload = new Gson().toJson(queryParams);
+                        oAuthMessage.getRequest().setAttribute(IS_API_BASED_AUTH_HANDLED, true);
+                        return Response.status(HttpServletResponse.SC_OK).entity(jsonPayload).build();
+                    }
                 }
             }
         } catch (AuthServiceException | JsonProcessingException | UnsupportedEncodingException | URISyntaxException e) {
@@ -4357,6 +4362,7 @@ private Response handleApiBasedAuthenticationResponse(OAuthMessage oAuthMessage,
             params.put(OAuthConstants.OAUTH_ERROR, OAuth2ErrorCodes.SERVER_ERROR);
             params.put(OAuthConstants.OAUTH_ERROR_DESCRIPTION, "Server error occurred while performing authorization.");
             String jsonString = new Gson().toJson(params);
+            oAuthMessage.getRequest().setAttribute(IS_API_BASED_AUTH_HANDLED, true);
             return Response.status(HttpServletResponse.SC_INTERNAL_SERVER_ERROR).entity(jsonString).build();
         }
 

components/org.wso2.carbon.identity.oauth.endpoint/src/test/java/org/wso2/carbon/identity/oauth/endpoint/authz/OAuth2AuthzEndpointTest.java

@@ -1856,7 +1856,7 @@ public Object answer(InvocationOnMock invocation) {
 
         Method sendRequestToFramework =
                 authzEndpointObject.getClass().getDeclaredMethod("handleAuthFlowThroughFramework",
-                        OAuthMessage.class, String.class);
+                        OAuthMessage.class, String.class, String.class);
         sendRequestToFramework.setAccessible(true);
 
         when(oAuthMessage.getRequest()).thenReturn(httpServletRequest);
@@ -1880,7 +1880,7 @@ public Object answer(InvocationOnMock invocation) {
 
         Method sendRequestToFramework2 =
                 authzEndpointObject.getClass().getDeclaredMethod("handleAuthFlowThroughFramework",
-                        OAuthMessage.class, String.class);
+                        OAuthMessage.class, String.class, String.class);
         sendRequestToFramework2.setAccessible(true);
         try {
             response = (Response) sendRequestToFramework.invoke(authzEndpointObject, oAuthMessage, "type");