From 9003970648d3461b0694c6a958ce37b74c984ded Mon Sep 17 00:00:00 2001
From: Wathsara Wishwantha Daluwatta
 <31571237+Wathsara@users.noreply.github.com>
Date: Wed, 26 Jul 2023 11:31:19 +0530
Subject: [PATCH]  Fix token revocation on session expiry for SAML SLO from
 Federated IDP

---
 .../TokenBindingExpiryEventHandler.java       | 34 ++++++++++++++++++-
 1 file changed, 33 insertions(+), 1 deletion(-)

diff --git a/components/org.wso2.carbon.identity.oauth/src/main/java/org/wso2/carbon/identity/oauth2/token/bindings/handlers/TokenBindingExpiryEventHandler.java b/components/org.wso2.carbon.identity.oauth/src/main/java/org/wso2/carbon/identity/oauth2/token/bindings/handlers/TokenBindingExpiryEventHandler.java
index a4b84cdecf..95aa0979a8 100644
--- a/components/org.wso2.carbon.identity.oauth/src/main/java/org/wso2/carbon/identity/oauth2/token/bindings/handlers/TokenBindingExpiryEventHandler.java
+++ b/components/org.wso2.carbon.identity.oauth/src/main/java/org/wso2/carbon/identity/oauth2/token/bindings/handlers/TokenBindingExpiryEventHandler.java
@@ -105,7 +105,7 @@ public void handleEvent(Event event) throws IdentityEventException {
                     revokeTokensForCommonAuthCookie(request, context.getLastAuthenticatedUser());
                 }
             } else {
-                revokeTokensForCommonAuthCookie(request, context.getLastAuthenticatedUser());
+                revokeTokensForCommonAuthCookie(request, getAuthenticatedUser(eventProperties, context));
             }
         } catch (IdentityOAuth2Exception | OAuthSystemException  e) {
             log.error("Error while revoking the tokens on session termination.", e);
@@ -418,4 +418,36 @@ private void revokeFederatedTokens(String consumerKey, AuthenticatedUser user, A
                 .revokeAccessTokens(new String[]{accessTokenDO.getAccessToken()}, OAuth2Util.isHashEnabled());
         OAuthUtil.invokePostRevocationBySystemListeners(accessTokenDO, Collections.emptyMap());
     }
+
+    /**
+     * Retrieve the authenticated user from the session context identifier in the event if it is not available in the
+     * authentication context.
+     *
+     * @param eventProperties Event properties.
+     * @return context Authentication context.
+     */
+    private AuthenticatedUser getAuthenticatedUser(Map<String, Object> eventProperties, AuthenticationContext context) {
+
+        AuthenticatedUser authenticatedUser = context.getLastAuthenticatedUser();
+        if (authenticatedUser != null) {
+            return authenticatedUser;
+        }
+        Map<String, Object> paramMap = (Map<String, Object>) eventProperties.get(IdentityEventConstants
+                .EventProperty.PARAMS);
+        String sessionContextIdentifier = getSessionIdentifier(paramMap);
+        if (StringUtils.isNotBlank(sessionContextIdentifier)) {
+            SessionContext sessionContext = (SessionContext) eventProperties.get(IdentityEventConstants
+                    .EventProperty.SESSION_CONTEXT);
+            if (sessionContext != null) {
+                authenticatedUser = (AuthenticatedUser) sessionContext
+                        .getProperty(FrameworkConstants.AUTHENTICATED_USER);
+            } else {
+                if (log.isDebugEnabled()) {
+                    log.debug("Session context for session context identifier: " + sessionContextIdentifier +
+                            " is not found.");
+                }
+            }
+        }
+        return authenticatedUser;
+    }
 }