Adds Saml Redirect Uri (#17936)

* Add Saml Redirect Uri In Session * Refactor * add redirect url in saml * add Saml logout servlet * fix for message * added logout method for saml --------- Co-authored-by: karanh37 <karanh37@gmail.com>
2025-07-24 09:50:01 +00:00 · 2024-09-23 22:20:48 +05:30 · 2024-09-23 22:20:48 +05:30 · d77cf36b38
commit d77cf36b38
parent 97b38f21ef
7 changed files with 132 additions and 6 deletions
--- a/openmetadata-service/src/main/java/org/openmetadata/service/OpenMetadataApplication.java
+++ b/openmetadata-service/src/main/java/org/openmetadata/service/OpenMetadataApplication.java
@ -116,6 +116,7 @@ import org.openmetadata.service.security.jwt.JWTTokenGenerator;
 import org.openmetadata.service.security.saml.OMMicrometerHttpFilter;
 import org.openmetadata.service.security.saml.SamlAssertionConsumerServlet;
 import org.openmetadata.service.security.saml.SamlLoginServlet;
+import org.openmetadata.service.security.saml.SamlLogoutServlet;
 import org.openmetadata.service.security.saml.SamlMetadataServlet;
 import org.openmetadata.service.security.saml.SamlSettingsHolder;
 import org.openmetadata.service.security.saml.SamlTokenRefreshServlet;
@ -350,6 +351,13 @@ public class OpenMetadataApplication extends Application<OpenMetadataApplication
      throws IOException, CertificateException, KeyStoreException, NoSuchAlgorithmException {
    if (catalogConfig.getAuthenticationConfiguration() != null
        && catalogConfig.getAuthenticationConfiguration().getProvider().equals(AuthProvider.SAML)) {
+
+      // Set up a Session Manager
+      MutableServletContextHandler contextHandler = environment.getApplicationContext();
+      if (contextHandler.getSessionHandler() == null) {
+        contextHandler.setSessionHandler(new SessionHandler());
+      }
+
      SamlSettingsHolder.getInstance().initDefaultSettings(catalogConfig);
      ServletRegistration.Dynamic samlRedirectServlet =
          environment.servlets().addServlet("saml_login", new SamlLoginServlet());
@ -368,6 +376,16 @@ public class OpenMetadataApplication extends Application<OpenMetadataApplication
      ServletRegistration.Dynamic samlRefreshServlet =
          environment.servlets().addServlet("saml_refresh_token", new SamlTokenRefreshServlet());
      samlRefreshServlet.addMapping("/api/v1/saml/refresh");
+
+      ServletRegistration.Dynamic samlLogoutServlet =
+          environment
+              .servlets()
+              .addServlet(
+                  "saml_logout_token",
+                  new SamlLogoutServlet(
+                      catalogConfig.getAuthenticationConfiguration(),
+                      catalogConfig.getAuthorizerConfiguration()));
+      samlLogoutServlet.addMapping("/api/v1/saml/logout");
    }
  }

--- a/openmetadata-service/src/main/java/org/openmetadata/service/security/AuthenticationCodeFlowHandler.java
+++ b/openmetadata-service/src/main/java/org/openmetadata/service/security/AuthenticationCodeFlowHandler.java
@ -281,7 +281,7 @@ public class AuthenticationCodeFlowHandler {
    }
  }

-  private void checkAndStoreRedirectUriInSession(HttpServletRequest request) {
+  public static void checkAndStoreRedirectUriInSession(HttpServletRequest request) {
    String redirectUri = request.getParameter(REDIRECT_URI_KEY);
    if (nullOrEmpty(redirectUri)) {
      throw new TechnicalException("Redirect URI is required");
--- a/openmetadata-service/src/main/java/org/openmetadata/service/security/saml/SamlAssertionConsumerServlet.java
+++ b/openmetadata-service/src/main/java/org/openmetadata/service/security/saml/SamlAssertionConsumerServlet.java
@ -14,6 +14,7 @@
 package org.openmetadata.service.security.saml;

 import static org.openmetadata.common.utils.CommonUtil.nullOrEmpty;
+import static org.openmetadata.service.security.AuthenticationCodeFlowHandler.SESSION_REDIRECT_URI;
 import static org.openmetadata.service.util.UserUtil.getRoleListFromUser;

 import com.onelogin.saml2.Auth;
@ -127,8 +128,9 @@ public class SamlAssertionConsumerServlet extends HttpServlet {
      resp.addCookie(refreshTokenCookie);

      // Redirect with JWT Token
+      String redirectUri = (String) req.getSession().getAttribute(SESSION_REDIRECT_URI);
      String url =
-          SamlSettingsHolder.getInstance().getRelayState()
+          redirectUri
              + "?id_token="
              + jwtAuthMechanism.getJWTToken()
              + "&email="
--- a/openmetadata-service/src/main/java/org/openmetadata/service/security/saml/SamlLoginServlet.java
+++ b/openmetadata-service/src/main/java/org/openmetadata/service/security/saml/SamlLoginServlet.java
@ -13,6 +13,8 @@

 package org.openmetadata.service.security.saml;

+import static org.openmetadata.service.security.AuthenticationCodeFlowHandler.checkAndStoreRedirectUriInSession;
+
 import com.onelogin.saml2.Auth;
 import java.io.IOException;
 import javax.servlet.annotation.WebServlet;
@ -33,6 +35,7 @@ public class SamlLoginServlet extends HttpServlet {
      throws IOException {
    Auth auth;
    try {
+      checkAndStoreRedirectUriInSession(req);
      auth = new Auth(SamlSettingsHolder.getInstance().getSaml2Settings(), req, resp);
      auth.login(SamlSettingsHolder.getInstance().getRelayState());
    } catch (Exception e) {
--- a/openmetadata-service/src/main/java/org/openmetadata/service/security/saml/SamlLogoutServlet.java
+++ b/openmetadata-service/src/main/java/org/openmetadata/service/security/saml/SamlLogoutServlet.java
@ -0,0 +1,101 @@
+/*
+ *  Copyright 2021 Collate
+ *  Licensed under the Apache License, Version 2.0 (the "License");
+ *  you may not use this file except in compliance with the License.
+ *  You may obtain a copy of the License at
+ *  http://www.apache.org/licenses/LICENSE-2.0
+ *  Unless required by applicable law or agreed to in writing, software
+ *  distributed under the License is distributed on an "AS IS" BASIS,
+ *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ *  See the License for the specific language governing permissions and
+ *  limitations under the License.
+ */
+
+package org.openmetadata.service.security.saml;
+
+import static org.openmetadata.common.utils.CommonUtil.listOrEmpty;
+import static org.openmetadata.service.security.AuthenticationCodeFlowHandler.getErrorMessage;
+import static org.openmetadata.service.security.SecurityUtil.findUserNameFromClaims;
+import static org.openmetadata.service.security.SecurityUtil.writeJsonResponse;
+
+import com.auth0.jwt.interfaces.Claim;
+import java.io.IOException;
+import java.time.LocalDateTime;
+import java.time.ZoneId;
+import java.util.Date;
+import java.util.List;
+import java.util.Map;
+import java.util.stream.Collectors;
+import javax.servlet.annotation.WebServlet;
+import javax.servlet.http.HttpServlet;
+import javax.servlet.http.HttpServletRequest;
+import javax.servlet.http.HttpServletResponse;
+import javax.servlet.http.HttpSession;
+import lombok.extern.slf4j.Slf4j;
+import org.openmetadata.schema.api.security.AuthenticationConfiguration;
+import org.openmetadata.schema.api.security.AuthorizerConfiguration;
+import org.openmetadata.schema.auth.LogoutRequest;
+import org.openmetadata.service.security.JwtFilter;
+
+/**
+ * This Servlet initiates a login and sends a login request to the IDP. After a successful processing it redirects user
+ * to the relayState which is the callback setup in the config.
+ */
+@WebServlet("/api/v1/saml/logout")
+@Slf4j
+public class SamlLogoutServlet extends HttpServlet {
+  private final JwtFilter jwtFilter;
+  private final List<String> jwtPrincipalClaims;
+  private final Map<String, String> jwtPrincipalClaimsMapping;
+
+  public SamlLogoutServlet(
+      AuthenticationConfiguration authenticationConfiguration,
+      AuthorizerConfiguration authorizerConf) {
+    jwtFilter = new JwtFilter(authenticationConfiguration, authorizerConf);
+    this.jwtPrincipalClaims = authenticationConfiguration.getJwtPrincipalClaims();
+    this.jwtPrincipalClaimsMapping =
+        listOrEmpty(authenticationConfiguration.getJwtPrincipalClaimsMapping()).stream()
+            .map(s -> s.split(":"))
+            .collect(Collectors.toMap(s -> s[0], s -> s[1]));
+  }
+
+  @Override
+  protected void doGet(
+      final HttpServletRequest httpServletRequest, final HttpServletResponse httpServletResponse)
+      throws IOException {
+    try {
+      LOG.debug("Performing application logout");
+      HttpSession session = httpServletRequest.getSession(false);
+      String token = JwtFilter.extractToken(httpServletRequest.getHeader("Authorization"));
+      if (session != null) {
+        LOG.debug("Invalidating the session for logout");
+        Map<String, Claim> claims = jwtFilter.validateJwtAndGetClaims(token);
+        String userName =
+            findUserNameFromClaims(jwtPrincipalClaimsMapping, jwtPrincipalClaims, claims);
+        Date logoutTime = Date.from(LocalDateTime.now().atZone(ZoneId.systemDefault()).toInstant());
+        // Mark the token invalid
+        JwtTokenCacheManager.getInstance()
+            .markLogoutEventForToken(
+                new LogoutRequest()
+                    .withUsername(userName)
+                    .withToken(token)
+                    .withLogoutTime(logoutTime));
+        // Invalidate the session
+        session.invalidate();
+
+        // Redirect to server
+        writeJsonResponse(httpServletResponse, "Logout successful");
+      } else {
+        LOG.error("No session store available for this web context");
+      }
+    } catch (Exception e) {
+      getErrorMessage(httpServletResponse, e);
+    }
+  }
+
+  public String getBaseUrl(HttpServletRequest request) {
+    String scheme = request.getScheme();
+    String serverName = request.getServerName();
+    return String.format("%s://%s", scheme, serverName);
+  }
+}
--- a/openmetadata-ui/src/main/resources/ui/src/components/Auth/AppAuthenticators/SamlAuthenticator.tsx
+++ b/openmetadata-ui/src/main/resources/ui/src/components/Auth/AppAuthenticators/SamlAuthenticator.tsx
@ -33,6 +33,7 @@ import { SamlSSOClientConfig } from '../../../generated/configuration/authentica
 import { postSamlLogout } from '../../../rest/miscAPI';
 import { showErrorToast } from '../../../utils/ToastUtils';

+import { ROUTES } from '../../../constants/constants';
 import { useApplicationStore } from '../../../hooks/useApplicationStore';
 import { AccessTokenResponse, refreshSAMLToken } from '../../../rest/auth-API';
 import { AuthenticatorRef } from '../AuthProviders/AuthProvider.interface';
@ -69,7 +70,8 @@ const SamlAuthenticator = forwardRef<AuthenticatorRef, Props>(

    const login = async () => {
      if (config.idp.authorityUrl) {
-        window.location.href = config.idp.authorityUrl;
+        const redirectUri = `${window.location.origin}${ROUTES.SAML_CALLBACK}`;
+        window.location.href = `${config.idp.authorityUrl}?redirectUri=${redirectUri}`;
      } else {
        showErrorToast('SAML IDP Authority URL is not configured.');
      }
@ -78,7 +80,7 @@ const SamlAuthenticator = forwardRef<AuthenticatorRef, Props>(
    const logout = () => {
      const token = getOidcToken();
      if (token) {
-        postSamlLogout({ token })
+        postSamlLogout()
          .then(() => {
            setIsAuthenticated(false);
            try {
--- a/openmetadata-ui/src/main/resources/ui/src/rest/miscAPI.ts
+++ b/openmetadata-ui/src/main/resources/ui/src/rest/miscAPI.ts
@ -94,8 +94,8 @@ export const getVersion = async () => {
  return response.data;
 };

-export const postSamlLogout = async (data: { token: string }) => {
-  const response = await APIClient.post(`/users/logout`, { ...data });
+export const postSamlLogout = async () => {
+  const response = await APIClient.get(`/saml/logout`);

  return response.data;
 };