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>

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");
     }
   }
 

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");

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="

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) {

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);
+  }
+}

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 {

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 }) => {
+export const postSamlLogout = async () => {
-  const response = await APIClient.post(`/users/logout`, { ...data });
+  const response = await APIClient.get(`/saml/logout`);
 
   return response.data;
 };