Inherit sso roles (#15903)

* Use Roles from Provider

* Add Cache For Roles From Users

* Fix Failing Test

* Revert secrets

Update Impl to use loggedInUser and CreateUser for Role assignment and resync

* Revert Expiry check

conf/openmetadata.yaml

@@ -161,6 +161,7 @@ authorizerConfiguration:
   principalDomain: ${AUTHORIZER_PRINCIPAL_DOMAIN:-"openmetadata.org"}
   enforcePrincipalDomain: ${AUTHORIZER_ENFORCE_PRINCIPAL_DOMAIN:-false}
   enableSecureSocketConnection : ${AUTHORIZER_ENABLE_SECURE_SOCKET:-false}
+  useRolesFromProvider: ${AUTHORIZER_USE_ROLES_FROM_PROVIDER:-false}
 
 authenticationConfiguration:
   clientType: ${AUTHENTICATION_CLIENT_TYPE:-public}

openmetadata-service/src/main/java/org/openmetadata/service/Entity.java

@@ -217,6 +217,7 @@ public final class Entity {
   //
   // Reserved names in OpenMetadata
   //
+  public static final String ADMIN_ROLE = "Admin";
   public static final String ADMIN_USER_NAME = "admin";
   public static final String ORGANIZATION_NAME = "Organization";
   public static final String ORGANIZATION_POLICY_NAME = "OrganizationPolicy";

openmetadata-service/src/main/java/org/openmetadata/service/OpenMetadataApplication.java

@@ -276,8 +276,8 @@ public class OpenMetadataApplication extends Application<OpenMetadataApplication
                   "oauth_login",
                   new AuthLoginServlet(
                       oidcClient,
-                      config.getAuthenticationConfiguration().getOidcConfiguration().getServerUrl(),
+                      config.getAuthenticationConfiguration(),
-                      config.getAuthenticationConfiguration().getJwtPrincipalClaims()));
+                      config.getAuthorizerConfiguration()));
       authLogin.addMapping("/api/v1/auth/login");
       ServletRegistration.Dynamic authCallback =
           environment
@@ -286,8 +286,8 @@ public class OpenMetadataApplication extends Application<OpenMetadataApplication
                   "auth_callback",
                   new AuthCallbackServlet(
                       oidcClient,
-                      config.getAuthenticationConfiguration().getOidcConfiguration().getServerUrl(),
+                      config.getAuthenticationConfiguration(),
-                      config.getAuthenticationConfiguration().getJwtPrincipalClaims()));
+                      config.getAuthorizerConfiguration()));
       authCallback.addMapping("/callback");
 
       ServletRegistration.Dynamic authLogout =
@@ -364,7 +364,11 @@ public class OpenMetadataApplication extends Application<OpenMetadataApplication
           environment.servlets().addServlet("saml_login", new SamlLoginServlet());
       samlRedirectServlet.addMapping("/api/v1/saml/login");
       ServletRegistration.Dynamic samlReceiverServlet =
-          environment.servlets().addServlet("saml_acs", new SamlAssertionConsumerServlet());
+          environment
+              .servlets()
+              .addServlet(
+                  "saml_acs",
+                  new SamlAssertionConsumerServlet(catalogConfig.getAuthorizerConfiguration()));
       samlReceiverServlet.addMapping("/api/v1/saml/acs");
       ServletRegistration.Dynamic samlMetadataServlet =
           environment.servlets().addServlet("saml_metadata", new SamlMetadataServlet());

openmetadata-service/src/main/java/org/openmetadata/service/resources/teams/UserResource.java

@@ -24,6 +24,10 @@ import static org.openmetadata.schema.entity.teams.AuthenticationMechanism.AuthT
 import static org.openmetadata.service.exception.CatalogExceptionMessage.EMAIL_SENDING_ISSUE;
 import static org.openmetadata.service.jdbi3.UserRepository.AUTH_MECHANISM_FIELD;
 import static org.openmetadata.service.security.jwt.JWTTokenGenerator.getExpiryDate;
+import static org.openmetadata.service.util.UserUtil.getRoleListFromUser;
+import static org.openmetadata.service.util.UserUtil.getRolesFromAuthorizationToken;
+import static org.openmetadata.service.util.UserUtil.reSyncUserRolesFromToken;
+import static org.openmetadata.service.util.UserUtil.validateAndGetRolesRef;
 
 import at.favre.lib.crypto.bcrypt.BCrypt;
 import freemarker.template.TemplateException;
@@ -64,6 +68,7 @@ import javax.ws.rs.Path;
 import javax.ws.rs.PathParam;
 import javax.ws.rs.Produces;
 import javax.ws.rs.QueryParam;
+import javax.ws.rs.container.ContainerRequestContext;
 import javax.ws.rs.core.Context;
 import javax.ws.rs.core.MediaType;
 import javax.ws.rs.core.Response;
@@ -77,6 +82,7 @@ import org.openmetadata.schema.EntityInterface;
 import org.openmetadata.schema.TokenInterface;
 import org.openmetadata.schema.api.data.RestoreEntity;
 import org.openmetadata.schema.api.security.AuthenticationConfiguration;
+import org.openmetadata.schema.api.security.AuthorizerConfiguration;
 import org.openmetadata.schema.api.teams.CreateUser;
 import org.openmetadata.schema.auth.BasicAuthMechanism;
 import org.openmetadata.schema.auth.ChangePasswordRequest;
@@ -160,6 +166,7 @@ public class UserResource extends EntityResource<User, UserRepository> {
   private final TokenRepository tokenRepository;
   private boolean isEmailServiceEnabled;
   private AuthenticationConfiguration authenticationConfiguration;
+  private AuthorizerConfiguration authorizerConfiguration;
   private final AuthenticatorHandler authHandler;
   static final String FIELDS = "profile,roles,teams,follows,owns,domain,personas,defaultPersona";
 
@@ -194,6 +201,7 @@ public class UserResource extends EntityResource<User, UserRepository> {
   public void initialize(OpenMetadataApplicationConfig config) throws IOException {
     super.initialize(config);
     this.authenticationConfiguration = config.getAuthenticationConfiguration();
+    this.authorizerConfiguration = config.getAuthorizerConfiguration();
     SmtpSettings smtpSettings = config.getSmtpSettings();
     this.isEmailServiceEnabled = smtpSettings != null && smtpSettings.getEnableSmtpServer();
     this.repository.initializeUsers(config);
@@ -416,6 +424,7 @@ public class UserResource extends EntityResource<User, UserRepository> {
   public User getCurrentLoggedInUser(
       @Context UriInfo uriInfo,
       @Context SecurityContext securityContext,
+      @Context ContainerRequestContext containerRequestContext,
       @Parameter(
               description = "Fields requested in the returned resource",
               schema = @Schema(type = "string", example = FIELDS))
@@ -424,6 +433,13 @@ public class UserResource extends EntityResource<User, UserRepository> {
     Fields fields = getFields(fieldsParam);
     String currentUserName = securityContext.getUserPrincipal().getName();
     User user = repository.getByName(uriInfo, currentUserName, fields);
+
+    // Sync the Roles from token to User
+    if (Boolean.TRUE.equals(authorizerConfiguration.getUseRolesFromProvider())
+        && Boolean.FALSE.equals(user.getIsBot() != null && user.getIsBot())) {
+      reSyncUserRolesFromToken(
+          uriInfo, user, getRolesFromAuthorizationToken(containerRequestContext));
+    }
     return addHref(uriInfo, user);
   }
 
@@ -529,6 +545,7 @@ public class UserResource extends EntityResource<User, UserRepository> {
   public Response createUser(
       @Context UriInfo uriInfo,
       @Context SecurityContext securityContext,
+      @Context ContainerRequestContext containerRequestContext,
       @Valid CreateUser create) {
     User user = getUser(securityContext.getUserPrincipal().getName(), create);
     if (Boolean.TRUE.equals(create.getIsAdmin())) {
@@ -558,7 +575,16 @@ public class UserResource extends EntityResource<User, UserRepository> {
       }
       // else the user will get a mail if configured smtp
     }
+
+    // Add the roles on user creation
+    if (Boolean.TRUE.equals(authorizerConfiguration.getUseRolesFromProvider())
+        && Boolean.FALSE.equals(user.getIsBot() != null && user.getIsBot())) {
+      user.setRoles(
+          validateAndGetRolesRef(getRolesFromAuthorizationToken(containerRequestContext)));
+    }
+
     // TODO do we need to authenticate user is creating himself?
+
     addHref(uriInfo, repository.create(uriInfo, user));
     if (isBasicAuth() && isEmailServiceEnabled) {
       try {
@@ -1254,13 +1280,16 @@ public class UserResource extends EntityResource<User, UserRepository> {
       @Valid CreatePersonalToken tokenRequest) {
     String userName = securityContext.getUserPrincipal().getName();
     User user =
-        repository.getByName(null, userName, getFields("email,isBot"), Include.NON_DELETED, false);
+        repository.getByName(
+            null, userName, getFields("roles,email,isBot"), Include.NON_DELETED, false);
     if (Boolean.FALSE.equals(user.getIsBot())) {
       // Create Personal Access Token
       JWTAuthMechanism authMechanism =
           JWTTokenGenerator.getInstance()
               .getJwtAuthMechanism(
                   userName,
+                  getRoleListFromUser(user),
+                  user.getIsAdmin(),
                   user.getEmail(),
                   false,
                   ServiceTokenType.PERSONAL_ACCESS,

openmetadata-service/src/main/java/org/openmetadata/service/security/AuthCallbackServlet.java

@@ -45,6 +45,8 @@ import javax.servlet.http.HttpServlet;
 import javax.servlet.http.HttpServletRequest;
 import javax.servlet.http.HttpServletResponse;
 import lombok.extern.slf4j.Slf4j;
+import org.openmetadata.schema.api.security.AuthenticationConfiguration;
+import org.openmetadata.schema.api.security.AuthorizerConfiguration;
 import org.pac4j.core.exception.TechnicalException;
 import org.pac4j.core.util.CommonHelper;
 import org.pac4j.oidc.client.OidcClient;
@@ -57,13 +59,19 @@ public class AuthCallbackServlet extends HttpServlet {
   private final ClientAuthentication clientAuthentication;
   private final List<String> claimsOrder;
   private final String serverUrl;
+  private final String principalDomain;
 
-  public AuthCallbackServlet(OidcClient oidcClient, String serverUrl, List<String> claimsOrder) {
+  public AuthCallbackServlet(
-    CommonHelper.assertNotBlank("ServerUrl", serverUrl);
+      OidcClient oidcClient,
+      AuthenticationConfiguration authenticationConfiguration,
+      AuthorizerConfiguration authorizerConfiguration) {
+    CommonHelper.assertNotBlank(
+        "ServerUrl", authenticationConfiguration.getOidcConfiguration().getServerUrl());
     this.client = oidcClient;
-    this.claimsOrder = claimsOrder;
+    this.claimsOrder = authenticationConfiguration.getJwtPrincipalClaims();
-    this.serverUrl = serverUrl;
+    this.serverUrl = authenticationConfiguration.getOidcConfiguration().getServerUrl();
     this.clientAuthentication = getClientAuthentication(client.getConfiguration());
+    this.principalDomain = authorizerConfiguration.getPrincipalDomain();
   }
 
   @Override
@@ -110,7 +118,7 @@ public class AuthCallbackServlet extends HttpServlet {
       req.getSession().setAttribute(OIDC_CREDENTIAL_PROFILE, credentials);
 
       // Redirect
-      sendRedirectWithToken(resp, credentials, serverUrl, claimsOrder);
+      sendRedirectWithToken(resp, credentials, serverUrl, claimsOrder, principalDomain);
     } catch (Exception e) {
       getErrorMessage(resp, e);
     }

openmetadata-service/src/main/java/org/openmetadata/service/security/AuthLoginServlet.java

@@ -22,6 +22,8 @@ import javax.servlet.http.HttpServlet;
 import javax.servlet.http.HttpServletRequest;
 import javax.servlet.http.HttpServletResponse;
 import lombok.extern.slf4j.Slf4j;
+import org.openmetadata.schema.api.security.AuthenticationConfiguration;
+import org.openmetadata.schema.api.security.AuthorizerConfiguration;
 import org.pac4j.core.exception.TechnicalException;
 import org.pac4j.core.util.CommonHelper;
 import org.pac4j.oidc.client.GoogleOidcClient;
@@ -36,11 +38,16 @@ public class AuthLoginServlet extends HttpServlet {
   private final OidcClient client;
   private final List<String> claimsOrder;
   private final String serverUrl;
+  private final String principalDomain;
 
-  public AuthLoginServlet(OidcClient oidcClient, String serverUrl, List<String> claimsOrder) {
+  public AuthLoginServlet(
+      OidcClient oidcClient,
+      AuthenticationConfiguration authenticationConfiguration,
+      AuthorizerConfiguration authorizerConfiguration) {
     this.client = oidcClient;
-    this.serverUrl = serverUrl;
+    this.serverUrl = authenticationConfiguration.getOidcConfiguration().getServerUrl();
-    this.claimsOrder = claimsOrder;
+    this.claimsOrder = authenticationConfiguration.getJwtPrincipalClaims();
+    this.principalDomain = authorizerConfiguration.getPrincipalDomain();
   }
 
   @Override
@@ -50,7 +57,7 @@ public class AuthLoginServlet extends HttpServlet {
       Optional<OidcCredentials> credentials = getUserCredentialsFromSession(req, client);
       if (credentials.isPresent()) {
         LOG.debug("Auth Tokens Located from Session: {} ", req.getSession().getId());
-        sendRedirectWithToken(resp, credentials.get(), serverUrl, claimsOrder);
+        sendRedirectWithToken(resp, credentials.get(), serverUrl, claimsOrder, principalDomain);
       } else {
         LOG.debug("Performing Auth Code Flow to Idp: {} ", req.getSession().getId());
         Map<String, String> params = buildParams();
@@ -105,8 +112,13 @@ public class AuthLoginServlet extends HttpServlet {
     }
 
     CodeChallengeMethod pkceMethod = client.getConfiguration().findPkceMethod();
+
+    // Use Default PKCE method if not disabled
+    if (pkceMethod == null && !client.getConfiguration().isDisablePkce()) {
+      pkceMethod = CodeChallengeMethod.S256;
+    }
     if (pkceMethod != null) {
-      CodeVerifier verfifier = new CodeVerifier(CommonHelper.randomString(10));
+      CodeVerifier verfifier = new CodeVerifier(CommonHelper.randomString(43));
       request.getSession().setAttribute(client.getCodeVerifierSessionAttributeName(), verfifier);
       params.put(
           OidcConfiguration.CODE_CHALLENGE,

openmetadata-service/src/main/java/org/openmetadata/service/security/CatalogOpenIdAuthorizationRequestFilter.java

@@ -15,6 +15,7 @@ package org.openmetadata.service.security;
 
 import static org.openmetadata.common.utils.CommonUtil.nullOrEmpty;
 
+import java.util.HashSet;
 import javax.annotation.Priority;
 import javax.ws.rs.container.ContainerRequestContext;
 import javax.ws.rs.container.ContainerRequestFilter;
@@ -50,7 +51,8 @@ public class CatalogOpenIdAuthorizationRequestFilter implements ContainerRequest
     CatalogPrincipal catalogPrincipal = new CatalogPrincipal(principal);
     String scheme = containerRequestContext.getUriInfo().getRequestUri().getScheme();
     CatalogSecurityContext catalogSecurityContext =
-        new CatalogSecurityContext(catalogPrincipal, scheme, CatalogSecurityContext.OPENID_AUTH);
+        new CatalogSecurityContext(
+            catalogPrincipal, scheme, CatalogSecurityContext.OPENID_AUTH, new HashSet<>());
     LOG.debug("SecurityContext {}", catalogSecurityContext);
     containerRequestContext.setSecurityContext(catalogSecurityContext);
   }

openmetadata-service/src/main/java/org/openmetadata/service/security/JwtFilter.java

@@ -14,6 +14,7 @@
 package org.openmetadata.service.security;
 
 import static org.openmetadata.common.utils.CommonUtil.nullOrEmpty;
+import static org.openmetadata.service.security.jwt.JWTTokenGenerator.ROLES_CLAIM;
 import static org.openmetadata.service.security.jwt.JWTTokenGenerator.TOKEN_TYPE;
 
 import com.auth0.jwk.Jwk;
@@ -59,6 +60,7 @@ public class JwtFilter implements ContainerRequestFilter {
   private String principalDomain;
   private boolean enforcePrincipalDomain;
   private AuthProvider providerType;
+  private boolean useRolesFromProvider = false;
 
   private static final List<String> DEFAULT_PUBLIC_KEY_URLS =
       Arrays.asList(
@@ -104,6 +106,7 @@ public class JwtFilter implements ContainerRequestFilter {
     this.jwkProvider = new MultiUrlJwkProvider(publicKeyUrlsBuilder.build());
     this.principalDomain = authorizerConfiguration.getPrincipalDomain();
     this.enforcePrincipalDomain = authorizerConfiguration.getEnforcePrincipalDomain();
+    this.useRolesFromProvider = authorizerConfiguration.getUseRolesFromProvider();
   }
 
   @VisibleForTesting
@@ -144,8 +147,19 @@ public class JwtFilter implements ContainerRequestFilter {
 
     String userName = validateAndReturnUsername(claims);
 
+    Set<String> userRoles = new HashSet<>();
+    boolean isBot =
+        claims.containsKey(BOT_CLAIM) && Boolean.TRUE.equals(claims.get(BOT_CLAIM).asBoolean());
+    // Re-sync user roles from token
+    if (useRolesFromProvider && !isBot && claims.containsKey(ROLES_CLAIM)) {
+      List<String> roles = claims.get(ROLES_CLAIM).asList(String.class);
+      if (!nullOrEmpty(roles)) {
+        userRoles = new HashSet<>(claims.get(ROLES_CLAIM).asList(String.class));
+      }
+    }
+
     // validate bot token
-    if (claims.containsKey(BOT_CLAIM) && Boolean.TRUE.equals(claims.get(BOT_CLAIM).asBoolean())) {
+    if (isBot) {
       validateBotToken(tokenFromHeader, userName);
     }
 
@@ -159,7 +173,8 @@ public class JwtFilter implements ContainerRequestFilter {
     CatalogPrincipal catalogPrincipal = new CatalogPrincipal(userName);
     String scheme = requestContext.getUriInfo().getRequestUri().getScheme();
     CatalogSecurityContext catalogSecurityContext =
-        new CatalogSecurityContext(catalogPrincipal, scheme, SecurityContext.DIGEST_AUTH);
+        new CatalogSecurityContext(
+            catalogPrincipal, scheme, SecurityContext.DIGEST_AUTH, userRoles);
     LOG.debug("SecurityContext {}", catalogSecurityContext);
     requestContext.setSecurityContext(catalogSecurityContext);
   }

openmetadata-service/src/main/java/org/openmetadata/service/security/NoopFilter.java

@@ -13,6 +13,7 @@
 
 package org.openmetadata.service.security;
 
+import java.util.HashSet;
 import javax.ws.rs.container.ContainerRequestContext;
 import javax.ws.rs.container.ContainerRequestFilter;
 import javax.ws.rs.core.Context;
@@ -35,7 +36,8 @@ public class NoopFilter implements ContainerRequestFilter {
     CatalogPrincipal catalogPrincipal = new CatalogPrincipal("anonymous");
     String scheme = containerRequestContext.getUriInfo().getRequestUri().getScheme();
     CatalogSecurityContext catalogSecurityContext =
-        new CatalogSecurityContext(catalogPrincipal, scheme, SecurityContext.BASIC_AUTH);
+        new CatalogSecurityContext(
+            catalogPrincipal, scheme, SecurityContext.BASIC_AUTH, new HashSet<>());
     LOG.debug("SecurityContext {}", catalogSecurityContext);
     containerRequestContext.setSecurityContext(catalogSecurityContext);
   }

openmetadata-service/src/main/java/org/openmetadata/service/security/SecurityUtil.java

@@ -327,7 +327,8 @@ public final class SecurityUtil {
       HttpServletResponse response,
       OidcCredentials credentials,
       String serverUrl,
-      List<String> claimsOrder)
+      List<String> claimsOrder,
+      String defaultDomain)
       throws ParseException, IOException {
     JWT jwt = credentials.getIdToken();
     Map<String, Object> claims = new TreeMap<>(String.CASE_INSENSITIVE_ORDER);
@@ -344,13 +345,13 @@ public final class SecurityUtil {
                         "Invalid JWT token, none of the following claims are present "
                             + claimsOrder));
 
-    String email = (String) jwt.getJWTClaimsSet().getClaim("email");
     String userName;
     if (preferredJwtClaim.contains("@")) {
       userName = preferredJwtClaim.split("@")[0];
     } else {
       userName = preferredJwtClaim;
     }
+    String email = String.format("%s@%s", userName, defaultDomain);
 
     String url =
         String.format(

openmetadata-service/src/main/java/org/openmetadata/service/security/auth/AuthenticatorHandler.java

@@ -1,6 +1,7 @@
 package org.openmetadata.service.security.auth;
 
 import static org.openmetadata.service.exception.CatalogExceptionMessage.NOT_IMPLEMENTED_METHOD;
+import static org.openmetadata.service.util.UserUtil.getRoleListFromUser;
 
 import freemarker.template.TemplateException;
 import java.io.IOException;
@@ -104,6 +105,8 @@ public interface AuthenticatorHandler {
         JWTTokenGenerator.getInstance()
             .generateJWTToken(
                 storedUser.getName(),
+                getRoleListFromUser(storedUser),
+                storedUser.getIsAdmin(),
                 storedUser.getEmail(),
                 expireInSeconds,
                 false,

openmetadata-service/src/main/java/org/openmetadata/service/security/auth/BasicAuthenticator.java

@@ -37,6 +37,7 @@ import static org.openmetadata.service.exception.CatalogExceptionMessage.TOKEN_E
 import static org.openmetadata.service.exception.CatalogExceptionMessage.TOKEN_EXPIRY_ERROR;
 import static org.openmetadata.service.resources.teams.UserResource.USER_PROTECTED_FIELDS;
 import static org.openmetadata.service.util.EmailUtil.getSmtpSettings;
+import static org.openmetadata.service.util.UserUtil.getRoleListFromUser;
 
 import at.favre.lib.crypto.bcrypt.BCrypt;
 import freemarker.template.TemplateException;
@@ -387,6 +388,8 @@ public class BasicAuthenticator implements AuthenticatorHandler {
         JWTTokenGenerator.getInstance()
             .generateJWTToken(
                 storedUser.getName(),
+                getRoleListFromUser(storedUser),
+                storedUser.getIsAdmin(),
                 storedUser.getEmail(),
                 loginConfiguration.getJwtTokenExpiryTime(),
                 false,
@@ -527,13 +530,15 @@ public class BasicAuthenticator implements AuthenticatorHandler {
             userRepository.getByEmail(
                 null,
                 userName,
-                new EntityUtil.Fields(Set.of(USER_PROTECTED_FIELDS), USER_PROTECTED_FIELDS));
+                new EntityUtil.Fields(
+                    Set.of(USER_PROTECTED_FIELDS, "roles"), "authenticationMechanism,roles"));
       } else {
         storedUser =
             userRepository.getByName(
                 null,
                 userName,
-                new EntityUtil.Fields(Set.of(USER_PROTECTED_FIELDS), USER_PROTECTED_FIELDS));
+                new EntityUtil.Fields(
+                    Set.of(USER_PROTECTED_FIELDS, "roles"), "authenticationMechanism,roles"));
       }
 
       if (storedUser != null && Boolean.TRUE.equals(storedUser.getIsBot())) {

openmetadata-service/src/main/java/org/openmetadata/service/security/auth/CatalogSecurityContext.java

@@ -13,14 +13,19 @@
 
 package org.openmetadata.service.security.auth;
 
+import static org.openmetadata.common.utils.CommonUtil.nullOrEmpty;
+
 import java.security.Principal;
+import java.util.HashSet;
+import java.util.Set;
 import javax.ws.rs.core.SecurityContext;
 import lombok.extern.slf4j.Slf4j;
 
 /** Holds authenticated principal and security context which is passed to the JAX-RS request methods */
 @Slf4j
 public record CatalogSecurityContext(
-    Principal principal, String scheme, String authenticationScheme) implements SecurityContext {
+    Principal principal, String scheme, String authenticationScheme, Set<String> userRoles)
+    implements SecurityContext {
   public static final String OPENID_AUTH = "openid";
 
   @Override
@@ -28,6 +33,13 @@ public record CatalogSecurityContext(
     return principal;
   }
 
+  public Set<String> getUserRoles() {
+    if (nullOrEmpty(userRoles)) {
+      return new HashSet<>();
+    }
+    return userRoles;
+  }
+
   @Override
   public boolean isUserInRole(String role) {
     LOG.debug("isUserInRole user: {}, role: {}", principal, role);

openmetadata-service/src/main/java/org/openmetadata/service/security/auth/CatalogSecurityContextRequestFilter.java

@@ -1,63 +0,0 @@
-/*
- *  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.auth;
-
-import java.security.Principal;
-import javax.servlet.http.HttpServletRequest;
-import javax.ws.rs.container.ContainerRequestContext;
-import javax.ws.rs.container.ContainerRequestFilter;
-import javax.ws.rs.core.Context;
-import javax.ws.rs.core.SecurityContext;
-import javax.ws.rs.ext.Provider;
-import lombok.extern.slf4j.Slf4j;
-import org.openmetadata.schema.api.security.AuthenticationConfiguration;
-import org.openmetadata.service.security.AuthenticationException;
-
-@Slf4j
-@Provider
-public class CatalogSecurityContextRequestFilter implements ContainerRequestFilter {
-
-  @Context private HttpServletRequest httpRequest;
-
-  @SuppressWarnings("unused")
-  private CatalogSecurityContextRequestFilter() {}
-
-  public CatalogSecurityContextRequestFilter(
-      AuthenticationConfiguration authenticationConfiguration) {
-    /* used for testing */
-  }
-
-  @Override
-  public void filter(ContainerRequestContext requestContext) {
-    Principal principal = httpRequest.getUserPrincipal();
-    String scheme = requestContext.getUriInfo().getRequestUri().getScheme();
-
-    LOG.debug(
-        "Method: {}, AuthType: {}, RemoteUser: {}, UserPrincipal: {}, Scheme: {}",
-        httpRequest.getMethod(),
-        httpRequest.getAuthType(),
-        httpRequest.getRemoteUser(),
-        principal,
-        scheme);
-
-    if (principal == null) {
-      throw new AuthenticationException("Not authorized. Principal is not available");
-    }
-
-    SecurityContext securityContext =
-        new CatalogSecurityContext(principal, scheme, httpRequest.getAuthType());
-    LOG.debug("SecurityContext {}", securityContext);
-    requestContext.setSecurityContext(securityContext);
-  }
-}

openmetadata-service/src/main/java/org/openmetadata/service/security/jwt/JWTTokenGenerator.java

@@ -13,6 +13,10 @@
 
 package org.openmetadata.service.security.jwt;
 
+import static org.openmetadata.common.utils.CommonUtil.nullOrEmpty;
+import static org.openmetadata.service.Entity.ADMIN_ROLE;
+import static org.openmetadata.service.util.UserUtil.getRoleListFromUser;
+
 import com.auth0.jwt.JWT;
 import com.auth0.jwt.algorithms.Algorithm;
 import com.auth0.jwt.exceptions.JWTCreationException;
@@ -30,6 +34,7 @@ import java.time.ZoneId;
 import java.util.Base64;
 import java.util.Date;
 import java.util.List;
+import java.util.Set;
 import lombok.Getter;
 import lombok.extern.slf4j.Slf4j;
 import org.openmetadata.schema.api.security.jwt.JWTTokenConfiguration;
@@ -41,6 +46,7 @@ import org.openmetadata.service.security.AuthenticationException;
 
 @Slf4j
 public class JWTTokenGenerator {
+  public static final String ROLES_CLAIM = "roles";
   public static final String SUBJECT_CLAIM = "sub";
   private static final String EMAIL_CLAIM = "email";
   private static final String IS_BOT_CLAIM = "isBot";
@@ -86,27 +92,54 @@ public class JWTTokenGenerator {
 
   public JWTAuthMechanism generateJWTToken(User user, JWTTokenExpiry expiry) {
     return getJwtAuthMechanism(
-        user.getName(), user.getEmail(), true, ServiceTokenType.BOT, getExpiryDate(expiry), expiry);
+        user.getName(),
+        getRoleListFromUser(user),
+        user.getIsAdmin(),
+        user.getEmail(),
+        true,
+        ServiceTokenType.BOT,
+        getExpiryDate(expiry),
+        expiry);
   }
 
   public JWTAuthMechanism generateJWTToken(
       String userName,
+      Set<String> roles,
+      boolean isAdmin,
       String email,
       long expiryInSeconds,
       boolean isBot,
       ServiceTokenType tokenType) {
     return getJwtAuthMechanism(
-        userName, email, isBot, tokenType, getCustomExpiryDate(expiryInSeconds), null);
+        userName,
+        roles,
+        isAdmin,
+        email,
+        isBot,
+        tokenType,
+        getCustomExpiryDate(expiryInSeconds),
+        null);
   }
 
   public JWTAuthMechanism getJwtAuthMechanism(
       String userName,
+      Set<String> roles,
+      boolean isAdmin,
       String email,
       boolean isBot,
       ServiceTokenType tokenType,
       Date expires,
       JWTTokenExpiry expiry) {
     try {
+      // Handle the Admin Role Here Since there is no Admin Role as such , just a isAdmin flag in
+      // User Schema
+      if (isAdmin) {
+        if (nullOrEmpty(roles)) {
+          roles = Set.of(ADMIN_ROLE);
+        } else {
+          roles.add(ADMIN_ROLE);
+        }
+      }
       JWTAuthMechanism jwtAuthMechanism = new JWTAuthMechanism().withJWTTokenExpiry(expiry);
       Algorithm algorithm = Algorithm.RSA256(null, privateKey);
       String token =
@@ -114,6 +147,7 @@ public class JWTTokenGenerator {
               .withIssuer(issuer)
               .withKeyId(kid)
               .withClaim(SUBJECT_CLAIM, userName)
+              .withClaim(ROLES_CLAIM, roles.stream().toList())
               .withClaim(EMAIL_CLAIM, email)
               .withClaim(IS_BOT_CLAIM, isBot)
               .withClaim(TOKEN_TYPE, tokenType.value())

openmetadata-service/src/main/java/org/openmetadata/service/security/saml/SamlAssertionConsumerServlet.java

@@ -13,15 +13,23 @@
 
 package org.openmetadata.service.security.saml;
 
+import static org.openmetadata.service.util.UserUtil.getRoleListFromUser;
+
 import com.onelogin.saml2.Auth;
+import java.util.HashSet;
 import java.util.List;
+import java.util.Set;
 import javax.servlet.annotation.WebServlet;
 import javax.servlet.http.HttpServlet;
 import javax.servlet.http.HttpServletRequest;
 import javax.servlet.http.HttpServletResponse;
 import lombok.extern.slf4j.Slf4j;
+import org.openmetadata.schema.api.security.AuthorizerConfiguration;
 import org.openmetadata.schema.auth.JWTAuthMechanism;
 import org.openmetadata.schema.auth.ServiceTokenType;
+import org.openmetadata.schema.entity.teams.User;
+import org.openmetadata.schema.type.Include;
+import org.openmetadata.service.Entity;
 import org.openmetadata.service.security.jwt.JWTTokenGenerator;
 
 /**
@@ -32,6 +40,12 @@ import org.openmetadata.service.security.jwt.JWTTokenGenerator;
 @WebServlet("/api/v1/saml/acs")
 @Slf4j
 public class SamlAssertionConsumerServlet extends HttpServlet {
+  private Set<String> admins;
+
+  public SamlAssertionConsumerServlet(AuthorizerConfiguration configuration) {
+    admins = configuration.getAdminPrincipals();
+  }
+
   @Override
   protected void doPost(HttpServletRequest req, HttpServletResponse resp) {
     try {
@@ -68,14 +82,32 @@ public class SamlAssertionConsumerServlet extends HttpServlet {
         email = String.format("%s@%s", username, SamlSettingsHolder.getInstance().getDomain());
       }
 
-      JWTAuthMechanism jwtAuthMechanism =
+      JWTAuthMechanism jwtAuthMechanism;
+      try {
+        User user = Entity.getEntityByName(Entity.USER, username, "id,roles", Include.NON_DELETED);
+        jwtAuthMechanism =
             JWTTokenGenerator.getInstance()
                 .generateJWTToken(
                     username,
+                    getRoleListFromUser(user),
+                    user.getIsAdmin(),
                     email,
                     SamlSettingsHolder.getInstance().getTokenValidity(),
                     false,
                     ServiceTokenType.OM_USER);
+      } catch (Exception e) {
+        LOG.error("[SAML ACS] User not found: " + username);
+        jwtAuthMechanism =
+            JWTTokenGenerator.getInstance()
+                .generateJWTToken(
+                    username,
+                    new HashSet<>(),
+                    admins.contains(username),
+                    email,
+                    SamlSettingsHolder.getInstance().getTokenValidity(),
+                    false,
+                    ServiceTokenType.OM_USER);
+      }
 
       String url =
           SamlSettingsHolder.getInstance().getRelayState()

openmetadata-service/src/main/java/org/openmetadata/service/util/UserUtil.java

@@ -14,22 +14,31 @@
 package org.openmetadata.service.util;
 
 import static org.openmetadata.common.utils.CommonUtil.listOf;
+import static org.openmetadata.common.utils.CommonUtil.listOrEmpty;
 import static org.openmetadata.common.utils.CommonUtil.nullOrEmpty;
 import static org.openmetadata.schema.entity.teams.AuthenticationMechanism.AuthType.JWT;
 import static org.openmetadata.schema.type.Include.NON_DELETED;
+import static org.openmetadata.service.Entity.ADMIN_ROLE;
 import static org.openmetadata.service.Entity.ADMIN_USER_NAME;
 
 import at.favre.lib.crypto.bcrypt.BCrypt;
+import java.util.ArrayList;
+import java.util.Collections;
 import java.util.HashSet;
 import java.util.List;
 import java.util.Objects;
 import java.util.Set;
 import java.util.UUID;
+import java.util.stream.Collectors;
+import javax.json.JsonPatch;
+import javax.ws.rs.container.ContainerRequestContext;
+import javax.ws.rs.core.UriInfo;
 import lombok.extern.slf4j.Slf4j;
 import org.openmetadata.schema.auth.BasicAuthMechanism;
 import org.openmetadata.schema.auth.JWTAuthMechanism;
 import org.openmetadata.schema.auth.JWTTokenExpiry;
 import org.openmetadata.schema.entity.teams.AuthenticationMechanism;
+import org.openmetadata.schema.entity.teams.Role;
 import org.openmetadata.schema.entity.teams.User;
 import org.openmetadata.schema.security.client.OpenMetadataJWTClientConfig;
 import org.openmetadata.schema.services.connections.metadata.AuthProvider;
@@ -40,6 +49,7 @@ import org.openmetadata.service.exception.EntityNotFoundException;
 import org.openmetadata.service.jdbi3.EntityRepository;
 import org.openmetadata.service.jdbi3.UserRepository;
 import org.openmetadata.service.resources.teams.RoleResource;
+import org.openmetadata.service.security.auth.CatalogSecurityContext;
 import org.openmetadata.service.security.jwt.JWTTokenGenerator;
 import org.openmetadata.service.util.EntityUtil.Fields;
 import org.openmetadata.service.util.RestUtil.PutResponse;
@@ -242,4 +252,98 @@ public final class UserUtil {
     }
     return userOrBot;
   }
+
+  public static Set<String> getRoleListFromUser(User user) {
+    if (nullOrEmpty(user.getRoles())) {
+      return new HashSet<>();
+    }
+    return listOrEmpty(user.getRoles()).stream()
+        .map(EntityReference::getName)
+        .collect(Collectors.toSet());
+  }
+
+  public static List<EntityReference> validateAndGetRolesRef(Set<String> rolesList) {
+    if (nullOrEmpty(rolesList)) {
+      return Collections.emptyList();
+    }
+    List<EntityReference> references = new ArrayList<>();
+
+    // Fetch the roles from the database
+    for (String role : rolesList) {
+      // Admin role is not present in the roles table, it is just a flag in the user table
+      if (!role.equals(ADMIN_ROLE)) {
+        try {
+          Role fetchedRole = Entity.getEntityByName(Entity.ROLE, role, "id", NON_DELETED, true);
+          references.add(fetchedRole.getEntityReference());
+        } catch (EntityNotFoundException ex) {
+          LOG.error("[ReSyncRoles] Role not found: {}", role, ex);
+        }
+      }
+    }
+    return references;
+  }
+
+  public static Set<String> getRolesFromAuthorizationToken(
+      ContainerRequestContext containerRequestContext) {
+    CatalogSecurityContext catalogSecurityContext =
+        (CatalogSecurityContext) containerRequestContext.getSecurityContext();
+    return catalogSecurityContext.getUserRoles();
+  }
+
+  public static boolean isRolesSyncNeeded(Set<String> fromToken, Set<String> fromDB) {
+    // Check if there are roles in the token that are not present in the DB
+    for (String role : fromToken) {
+      if (!fromDB.contains(role)) {
+        return true;
+      }
+    }
+
+    // Check if there are roles in the DB that are not present in the token
+    for (String role : fromDB) {
+      if (!fromToken.contains(role)) {
+        return true;
+      }
+    }
+
+    return false;
+  }
+
+  public static boolean reSyncUserRolesFromToken(
+      UriInfo uriInfo, User user, Set<String> rolesFromToken) {
+    boolean syncUser = false;
+
+    User updatedUser = JsonUtils.deepCopy(user, User.class);
+    // Check if Admin User
+    if (rolesFromToken.contains(ADMIN_ROLE)) {
+      if (Boolean.FALSE.equals(user.getIsAdmin())) {
+        syncUser = true;
+        updatedUser.setIsAdmin(true);
+      }
+
+      // Remove the Admin Role from the list
+      rolesFromToken.remove(ADMIN_ROLE);
+    }
+
+    Set<String> rolesFromUser = getRoleListFromUser(user);
+
+    // Check if roles are different
+    if (!nullOrEmpty(rolesFromToken) && isRolesSyncNeeded(rolesFromToken, rolesFromUser)) {
+      syncUser = true;
+      List<EntityReference> rolesReferenceFromToken = validateAndGetRolesRef(rolesFromToken);
+      updatedUser.setRoles(rolesReferenceFromToken);
+    }
+
+    if (syncUser) {
+      LOG.info("Syncing User Roles for User: {}", user.getName());
+      JsonPatch patch = JsonUtils.getJsonPatch(user, updatedUser);
+
+      UserRepository userRepository = (UserRepository) Entity.getEntityRepository(Entity.USER);
+      userRepository.patch(uriInfo, user.getId(), user.getName(), patch);
+
+      // Set the updated roles to the original user
+      user.setRoles(updatedUser.getRoles());
+    }
+
+    return syncUser;
+  }
 }

openmetadata-spec/src/main/resources/json/schema/configuration/authorizerConfiguration.json

@@ -58,6 +58,11 @@
     "enableSecureSocketConnection": {
       "description": "Enable Secure Socket Connection.",
       "type": "boolean"
+    },
+    "useRolesFromProvider": {
+      "description": "Use Roles from Provider",
+      "type": "boolean",
+      "default": false
     }
   },
   "required": ["className", "containerRequestFilter", "adminPrincipals", "principalDomain", "enforcePrincipalDomain", "enableSecureSocketConnection"],