Update AuthenticationValve.java

Resolve SP_UUID from token and set to thread local.
wso2-extensions · Jun 14, 2024 · 5860f54 · 5860f54
1 parent a6e3f3a
commit 5860f54
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 41 deletions.
diff --git a/...ain/java/org/wso2/carbon/identity/auth/service/handler/impl/OAuth2AccessTokenHandler.java b/...ain/java/org/wso2/carbon/identity/auth/service/handler/impl/OAuth2AccessTokenHandler.java
@@ -173,13 +173,14 @@ protected AuthenticationResult doAuthenticate(MessageContext messageContext) {
                 ServiceProvider serviceProvider = null;
                 String serviceProviderName = null;
                 String serviceProviderUUID = null;
-
                 try {
                     serviceProvider = OAuth2Util.getServiceProvider(oAuth2IntrospectionResponseDTO.getClientId());
                     if (serviceProvider != null) {
                         serviceProviderName = serviceProvider.getApplicationName();
                         serviceProviderUUID = serviceProvider.getApplicationResourceId();
                     } else {
+                        log.debug("There is no associated Service provider for client Id "
+                                + oAuth2IntrospectionResponseDTO.getClientId());
                         throw new IdentityOAuth2Exception("There is no associated Service provider for client Id "
                                 + oAuth2IntrospectionResponseDTO.getClientId());
                     }
@@ -197,23 +198,22 @@ protected AuthenticationResult doAuthenticate(MessageContext messageContext) {
                             + oAuth2IntrospectionResponseDTO.getClientId(), e);
                 }
 
-                if (serviceProvider != null) {
-                    if (serviceProviderName != null){
-                        authenticationContext.addParameter(SERVICE_PROVIDER_NAME, serviceProviderName);
-                    }
-                    if (serviceProviderTenantDomain != null) {
-                        authenticationContext.addParameter(SERVICE_PROVIDER_TENANT_DOMAIN, serviceProviderTenantDomain);
-                    }
-                    if (serviceProviderUUID != null) {
-                        authenticationContext.addParameter(SERVICE_PROVIDER_UUID, serviceProviderUUID);
-                    }
-
-                    MDC.put(SERVICE_PROVIDER_NAME, serviceProviderName);
-                    MDC.put(SERVICE_PROVIDER_UUID, serviceProviderUUID);
-                    // Set OAuth service provider details to be consumed by the provisioning framework.
-                    setProvisioningServiceProviderThreadLocal(oAuth2IntrospectionResponseDTO.getClientId(),
-                            serviceProviderTenantDomain);
+                if (serviceProviderName != null){
+                    authenticationContext.addParameter(SERVICE_PROVIDER_NAME, serviceProviderName);
+                }
+                if (serviceProviderTenantDomain != null) {
+                    authenticationContext.addParameter(SERVICE_PROVIDER_TENANT_DOMAIN, serviceProviderTenantDomain);
+                }
+                if (serviceProviderUUID != null) {
+                    authenticationContext.addParameter(SERVICE_PROVIDER_UUID, serviceProviderUUID);
                 }
+
+                MDC.put(SERVICE_PROVIDER_NAME, serviceProviderName);
+                MDC.put(SERVICE_PROVIDER_UUID, serviceProviderUUID);
+                // Set OAuth service provider details to be consumed by the provisioning framework.
+                setProvisioningServiceProviderThreadLocal(oAuth2IntrospectionResponseDTO.getClientId(),
+                        serviceProviderTenantDomain);
+
             }
         }
         return authenticationResult;

diff --git a/...ity.auth.valve/src/main/java/org/wso2/carbon/identity/auth/valve/AuthenticationValve.java b/...ity.auth.valve/src/main/java/org/wso2/carbon/identity/auth/valve/AuthenticationValve.java
@@ -81,7 +81,7 @@ public class AuthenticationValve extends ValveBase {
     private static final String AUTH_CONTEXT = "auth-context";
     private static final String USER_AGENT = "User-Agent";
     private static final String REMOTE_ADDRESS = "remoteAddress";
-    private static final String SERVICE_PROVIDER = "serviceProvider";
+    private static final String SERVICE_PROVIDER_NAME = "serviceProvider";
     private static final String IMPERSONATOR = "impersonator";
     private static final String SERVICE_PROVIDER_UUID= "serviceProviderUUID";
     private final String CLIENT_COMPONENT = "clientComponent";
@@ -144,8 +144,6 @@ public void invoke(Request request, Response response) throws IOException, Servl
             if (authenticationStatus.equals(AuthenticationStatus.SUCCESS)) {
                 // Set service provider info used in authentication if any.
                 setThreadLocalServiceProvider(authenticationContext);
-                // Set service provider UUID info used in authentication if any.
-                setThreadLocalServiceProviderUUID(authenticationContext);
                 // Set authenticated user tenant domain.
                 setThreadLocalAuthUserTenantDomain(authenticationContext);
                 // Set client component in to MDC.
@@ -196,8 +194,6 @@ public void invoke(Request request, Response response) throws IOException, Servl
 
             // Clear thread local service provider info.
             unsetThreadLocalServiceProvider();
-            // Clear thread local service provider UUID info.
-            unsetThreadLocalServiceProviderUUID();
             // Clear thread local current session id.
             unsetCurrentSessionIdThreadLocal();
             // Clear thread local authenticated user tenant domain.
@@ -270,34 +266,21 @@ private boolean isUnauthorized(ResourceConfig securedResource) {
 
     private void setThreadLocalServiceProvider(AuthenticationContext authenticationContext) {
 
-        Object serviceProvider = authenticationContext.getParameter(SERVICE_PROVIDER);
-        Object serviceProviderTenantDomain = authenticationContext.getParameter(SERVICE_PROVIDER_TENANT_DOMAIN);
-        if (serviceProvider != null && serviceProviderTenantDomain != null) {
-            IdentityUtil.threadLocalProperties.get().put(SERVICE_PROVIDER, serviceProvider);
-            IdentityUtil.threadLocalProperties.get().put(SERVICE_PROVIDER_TENANT_DOMAIN, serviceProviderTenantDomain);
-        }
-    }
-
-    private void setThreadLocalServiceProviderUUID(AuthenticationContext authenticationContext) {
-
+        Object serviceProviderName = authenticationContext.getParameter(SERVICE_PROVIDER_NAME);
         Object serviceProviderUUID = authenticationContext.getParameter(SERVICE_PROVIDER_UUID);
         Object serviceProviderTenantDomain = authenticationContext.getParameter(SERVICE_PROVIDER_TENANT_DOMAIN);
-        if (serviceProviderUUID != null && serviceProviderTenantDomain != null) {
+        if(serviceProviderTenantDomain != null && serviceProviderName != null && serviceProviderUUID != null){
+            IdentityUtil.threadLocalProperties.get().put(SERVICE_PROVIDER_NAME, serviceProviderName);
             IdentityUtil.threadLocalProperties.get().put(SERVICE_PROVIDER_UUID, serviceProviderUUID);
             IdentityUtil.threadLocalProperties.get().put(SERVICE_PROVIDER_TENANT_DOMAIN, serviceProviderTenantDomain);
         }
     }
 
     private void unsetThreadLocalServiceProvider() {
 
-        IdentityUtil.threadLocalProperties.get().remove(SERVICE_PROVIDER);
+        IdentityUtil.threadLocalProperties.get().remove(SERVICE_PROVIDER_NAME);
         IdentityUtil.threadLocalProperties.get().remove(SERVICE_PROVIDER_TENANT_DOMAIN);
-    }
-
-    private void unsetThreadLocalServiceProviderUUID() {
-
         IdentityUtil.threadLocalProperties.get().remove(SERVICE_PROVIDER_UUID);
-        IdentityUtil.threadLocalProperties.get().remove(SERVICE_PROVIDER_TENANT_DOMAIN);
     }
 
     private void unsetThreadLocalAuthenticationType() {
@@ -323,7 +306,7 @@ private void unsetThreadLocalAuthUserTenantDomain() {
 
     private void setClientComponent() {
 
-        String serviceProvider = MDC.get(SERVICE_PROVIDER);
+        String serviceProvider = MDC.get(SERVICE_PROVIDER_NAME);
         if (serviceProvider != null) {
             MDC.put(CLIENT_COMPONENT, serviceProvider);
         } else {
@@ -336,7 +319,7 @@ private void unsetMDCThreadLocals() {
         MDC.remove(CLIENT_COMPONENT);
         MDC.remove(USER_AGENT);
         MDC.remove(REMOTE_ADDRESS);
-        MDC.remove(SERVICE_PROVIDER);
+        MDC.remove(SERVICE_PROVIDER_NAME);
         MDC.remove(IMPERSONATOR);
     }