Fixes #6171 Add SubjectContext for Policies (#6172)

catalog-rest-service/src/main/java/org/openmetadata/catalog/security/DefaultAuthorizer.java

@@ -15,10 +15,8 @@ package org.openmetadata.catalog.security;
 
 import static org.openmetadata.catalog.Entity.FIELD_OWNER;
 import static org.openmetadata.catalog.security.SecurityUtil.DEFAULT_PRINCIPAL_DOMAIN;
-import static org.openmetadata.common.utils.CommonUtil.listOrEmpty;
 
 import java.io.IOException;
-import java.util.ArrayList;
 import java.util.Collections;
 import java.util.HashSet;
 import java.util.List;
@@ -35,6 +33,7 @@ import org.openmetadata.catalog.entity.teams.User;
 import org.openmetadata.catalog.exception.EntityNotFoundException;
 import org.openmetadata.catalog.jdbi3.EntityRepository;
 import org.openmetadata.catalog.security.policyevaluator.RoleEvaluator;
+import org.openmetadata.catalog.security.policyevaluator.SubjectContext;
 import org.openmetadata.catalog.type.EntityReference;
 import org.openmetadata.catalog.type.Include;
 import org.openmetadata.catalog.type.MetadataOperation;
@@ -122,10 +121,8 @@ public class DefaultAuthorizer implements Authorizer {
       AuthenticationContext ctx, EntityReference entityReference, MetadataOperation operation) {
     validate(ctx);
     try {
-      EntityRepository<User> userRepository = Entity.getEntityRepository(Entity.USER);
+      SubjectContext subjectContext = SubjectContext.getSubjectContext(SecurityUtil.getUserName(ctx));
-      Fields fieldsRolesAndTeams = userRepository.getFields("roles, teams");
+      List<EntityReference> allRoles = subjectContext.getAllRoles();
-      User user = getUserFromAuthenticationContext(ctx, fieldsRolesAndTeams);
-      List<EntityReference> allRoles = getAllRoles(user);
       if (entityReference == null) {
         // In some cases there is no specific entity being acted upon. Eg: Lineage.
         return RoleEvaluator.getInstance().hasPermissions(allRoles, null, operation);
@@ -135,7 +132,7 @@ public class DefaultAuthorizer implements Authorizer {
           Entity.getEntity(entityReference, new Fields(List.of("tags", FIELD_OWNER)), Include.NON_DELETED);
       EntityReference owner = entity.getOwner();
 
-      if (Entity.shouldHaveOwner(entityReference.getType()) && owner != null && isOwnedByUser(user, owner)) {
+      if (Entity.shouldHaveOwner(entityReference.getType()) && owner != null && subjectContext.isOwner(owner)) {
         return true; // Entity is owned by the user.
       }
       return RoleEvaluator.getInstance().hasPermissions(allRoles, entity, operation);
@@ -154,17 +151,15 @@ public class DefaultAuthorizer implements Authorizer {
     }
 
     try {
-      EntityRepository<User> userRepository = Entity.getEntityRepository(Entity.USER);
+      SubjectContext subjectContext = SubjectContext.getSubjectContext(SecurityUtil.getUserName(ctx));
-      Fields fieldsRolesAndTeams = userRepository.getFields("roles, teams");
+      List<EntityReference> allRoles = subjectContext.getAllRoles();
-      User user = getUserFromAuthenticationContext(ctx, fieldsRolesAndTeams);
-      List<EntityReference> allRoles = getAllRoles(user);
       if (entityReference == null) {
         return RoleEvaluator.getInstance().getAllowedOperations(allRoles, null);
       }
       EntityInterface entity =
           Entity.getEntity(entityReference, new Fields(List.of("tags", FIELD_OWNER)), Include.NON_DELETED);
       EntityReference owner = entity.getOwner();
-      if (owner == null || isOwnedByUser(user, owner)) {
+      if (owner == null || subjectContext.isOwner(owner)) {
         // Entity does not have an owner or is owned by the user - allow all operations.
         return Stream.of(MetadataOperation.values()).collect(Collectors.toList());
       }
@@ -174,28 +169,13 @@ public class DefaultAuthorizer implements Authorizer {
     }
   }
 
-  /** Checks if the user is same as owner or part of the team that is the owner. */
-  private boolean isOwnedByUser(User user, EntityReference owner) {
-    if (owner.getType().equals(Entity.USER) && owner.getName().equals(user.getName())) {
-      return true; // Owner is same as user.
-    }
-    if (owner.getType().equals(Entity.TEAM)) {
-      for (EntityReference userTeam : user.getTeams()) {
-        if (userTeam.getName().equals(owner.getName())) {
-          return true; // Owner is a team, and the user is part of this team.
-        }
-      }
-    }
-    return false;
-  }
-
   @Override
   public boolean isAdmin(AuthenticationContext ctx) {
     validate(ctx);
     try {
-      User user = getUserFromAuthenticationContext(ctx, Fields.EMPTY_FIELDS);
+      SubjectContext subjectContext = SubjectContext.getSubjectContext(SecurityUtil.getUserName(ctx));
-      return Boolean.TRUE.equals(user.getIsAdmin());
+      return subjectContext.isAdmin();
-    } catch (IOException | EntityNotFoundException ex) {
+    } catch (EntityNotFoundException ex) {
       return false;
     }
   }
@@ -204,9 +184,9 @@ public class DefaultAuthorizer implements Authorizer {
   public boolean isBot(AuthenticationContext ctx) {
     validate(ctx);
     try {
-      User user = getUserFromAuthenticationContext(ctx, Fields.EMPTY_FIELDS);
+      SubjectContext subjectContext = SubjectContext.getSubjectContext(SecurityUtil.getUserName(ctx));
-      return Boolean.TRUE.equals(user.getIsBot());
+      return subjectContext.isBot();
-    } catch (IOException | EntityNotFoundException ex) {
+    } catch (EntityNotFoundException ex) {
       return false;
     }
   }
@@ -218,11 +198,9 @@ public class DefaultAuthorizer implements Authorizer {
     }
     validate(ctx);
     try {
-      EntityRepository<User> userRepository = Entity.getEntityRepository(Entity.USER);
+      SubjectContext subjectContext = SubjectContext.getSubjectContext(SecurityUtil.getUserName(ctx));
-      Fields fieldsTeams = userRepository.getFields("teams");
+      return subjectContext.isOwner(owner);
-      User user = getUserFromAuthenticationContext(ctx, fieldsTeams);
+    } catch (EntityNotFoundException ex) {
-      return isOwnedByUser(user, owner);
-    } catch (IOException | EntityNotFoundException ex) {
       return false;
     }
   }
@@ -233,25 +211,6 @@ public class DefaultAuthorizer implements Authorizer {
     }
   }
 
-  private User getUserFromAuthenticationContext(AuthenticationContext ctx, Fields fields) throws IOException {
-    EntityRepository<User> userRepository = Entity.getEntityRepository(Entity.USER);
-    if (ctx.getUser() != null) {
-      // If a requested field is not present in the user, then add it
-      for (String field : fields.getFieldList()) {
-        if (!ctx.getUserFields().contains(field)) {
-          userRepository.setFields(ctx.getUser(), userRepository.getFields(field));
-          ctx.getUserFields().add(fields);
-        }
-      }
-      return ctx.getUser();
-    }
-    String userName = SecurityUtil.getUserName(ctx);
-    User user = userRepository.getByName(null, userName, fields);
-    ctx.setUser(user);
-    ctx.setUserFields(fields);
-    return user;
-  }
-
   private void addOrUpdateUser(User user) {
     EntityRepository<User> userRepository = Entity.getEntityRepository(Entity.USER);
     try {
@@ -263,10 +222,4 @@ public class DefaultAuthorizer implements Authorizer {
       LOG.debug("User entry: {} already exists.", user);
     }
   }
-
-  private List<EntityReference> getAllRoles(User user) {
-    List<EntityReference> allRoles = new ArrayList<>(listOrEmpty(user.getRoles()));
-    allRoles.addAll(listOrEmpty(user.getInheritedRoles()));
-    return allRoles.stream().distinct().collect(Collectors.toList()); // Remove duplicates
-  }
 }

catalog-rest-service/src/main/java/org/openmetadata/catalog/security/policyevaluator/RoleEvaluator.java

@@ -59,7 +59,13 @@ public class RoleEvaluator {
     try {
       Fields roleFields = roleRepository.getFields("policies");
       ResultList<Role> roles = roleRepository.listAfter(null, roleFields, filter, Short.MAX_VALUE, null);
-      roles.getData().forEach(r -> roleToPolicies.put(r.getId(), r.getPolicies()));
+      roles
+          .getData()
+          .forEach(
+              r -> {
+                LOG.info("Adding to role {}:{} policies {}", r.getName(), r.getId(), r.getPolicies());
+                roleToPolicies.put(r.getId(), r.getPolicies());
+              });
     } catch (IOException e) {
       LOG.error("Failed to load roles", e);
     }
@@ -67,8 +73,8 @@ public class RoleEvaluator {
 
   public boolean hasPermissions(List<EntityReference> roles, EntityInterface entity, MetadataOperation operation) {
     // Role based permission
-    for (EntityReference roleRef : roles) {
+    for (EntityReference role : roles) {
-      List<EntityReference> policies = roleToPolicies.get(roleRef.getId());
+      List<EntityReference> policies = roleToPolicies.get(role.getId());
       for (EntityReference policy : policies) {
         if (PolicyEvaluator.getInstance().hasPermission(policy.getId(), entity, operation)) {
           return true;
@@ -80,8 +86,8 @@ public class RoleEvaluator {
 
   public List<MetadataOperation> getAllowedOperations(List<EntityReference> roles, EntityInterface entity) {
     List<MetadataOperation> list = new ArrayList<>();
-    for (EntityReference roleRef : roles) {
+    for (EntityReference role : roles) {
-      List<EntityReference> policies = roleToPolicies.get(roleRef.getId());
+      List<EntityReference> policies = roleToPolicies.get(role.getId());
       for (EntityReference policy : policies) {
         list.addAll(PolicyEvaluator.getInstance().getAllowedOperations(policy.getId(), entity));
       }
@@ -91,9 +97,11 @@ public class RoleEvaluator {
 
   public void update(Role role) {
     roleToPolicies.put(role.getId(), role.getPolicies());
+    LOG.info("Updating to role {}:{} policies {}", role.getName(), role.getId(), role.getPolicies());
   }
 
   public void delete(Role role) {
     roleToPolicies.remove(role.getId());
+    LOG.info("Removing to role {}:{}", role.getName(), role.getId());
   }
 }

catalog-rest-service/src/main/java/org/openmetadata/catalog/security/policyevaluator/SubjectContext.java

@@ -0,0 +1,97 @@
+/*
+ *  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.catalog.security.policyevaluator;
+
+import static org.openmetadata.common.utils.CommonUtil.listOrEmpty;
+
+import com.google.common.cache.CacheBuilder;
+import com.google.common.cache.CacheLoader;
+import com.google.common.cache.LoadingCache;
+import com.google.common.util.concurrent.UncheckedExecutionException;
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.TimeUnit;
+import java.util.stream.Collectors;
+import org.openmetadata.catalog.Entity;
+import org.openmetadata.catalog.entity.teams.User;
+import org.openmetadata.catalog.exception.EntityNotFoundException;
+import org.openmetadata.catalog.jdbi3.EntityRepository;
+import org.openmetadata.catalog.type.EntityReference;
+import org.openmetadata.catalog.util.EntityUtil.Fields;
+
+/** Subject context used for Access Control Policies */
+public class SubjectContext {
+  // Cache used for caching subject context
+  private static final LoadingCache<String, SubjectContext> CACHE =
+      CacheBuilder.newBuilder().maximumSize(1000).expireAfterAccess(1, TimeUnit.MINUTES).build(new SubjectLoader());
+  private static final EntityRepository<User> USER_REPOSITORY = Entity.getEntityRepository(Entity.USER);
+  private static final Fields SUBJECT_FIELDS = USER_REPOSITORY.getFields("roles, teams");
+
+  private final User user;
+
+  public SubjectContext(User user) {
+    this.user = user;
+  }
+
+  public static SubjectContext getSubjectContext(String userName) throws EntityNotFoundException {
+    try {
+      return CACHE.get(userName);
+    } catch (ExecutionException | UncheckedExecutionException ex) {
+      throw new EntityNotFoundException(ex.getMessage());
+    }
+  }
+
+  public static void cleanup() {
+    CACHE.invalidateAll();
+  }
+
+  public boolean isAdmin() {
+    return Boolean.TRUE.equals(user.getIsAdmin());
+  }
+
+  public boolean isBot() {
+    return Boolean.TRUE.equals(user.getIsBot());
+  }
+
+  public boolean isOwner(EntityReference entityOwner) {
+    if (entityOwner.getType().equals(Entity.USER) && entityOwner.getName().equals(user.getName())) {
+      return true; // Owner is same as user.
+    }
+    if (entityOwner.getType().equals(Entity.TEAM)) {
+      for (EntityReference userTeam : user.getTeams()) {
+        if (userTeam.getName().equals(entityOwner.getName())) {
+          return true; // Owner is a team, and the user is part of this team.
+        }
+      }
+    }
+    return false;
+  }
+
+  public List<EntityReference> getAllRoles() {
+    List<EntityReference> allRoles = new ArrayList<>(listOrEmpty(user.getRoles()));
+    allRoles.addAll(listOrEmpty(user.getInheritedRoles()));
+    return allRoles.stream().distinct().collect(Collectors.toList()); // Remove duplicates
+  }
+
+  static class SubjectLoader extends CacheLoader<String, SubjectContext> {
+    @Override
+    public SubjectContext load(String userName) throws IOException {
+      User user = USER_REPOSITORY.getByName(null, userName, SUBJECT_FIELDS);
+      user.getRoles().forEach(r -> System.out.println("User " + user.getRoles() + " role " + r));
+      return new SubjectContext(user);
+    }
+  }
+}

catalog-rest-service/src/test/java/org/openmetadata/catalog/CatalogApplicationTest.java

@@ -29,6 +29,7 @@ import org.junit.jupiter.api.BeforeAll;
 import org.openmetadata.catalog.fernet.Fernet;
 import org.openmetadata.catalog.resources.CollectionRegistry;
 import org.openmetadata.catalog.resources.events.WebhookCallbackResource;
+import org.openmetadata.catalog.security.policyevaluator.SubjectContext;
 import org.testcontainers.containers.JdbcDatabaseContainer;
 
 @Slf4j
@@ -92,6 +93,7 @@ public abstract class CatalogApplicationTest {
     if (APP != null) {
       APP.after();
     }
+    SubjectContext.cleanup();
   }
 
   public static Client getClient() {

catalog-rest-service/src/test/java/org/openmetadata/catalog/resources/EntityResourceTest.java

@@ -277,7 +277,6 @@ public abstract class EntityResourceTest<T extends EntityInterface, K extends Cr
     new StorageServiceResourceTest().setupStorageServices();
     new DashboardServiceResourceTest().setupDashboardServices(test);
     new MlModelServiceResourceTest().setupMlModelServices(test);
-
     new TableResourceTest().setupDatabaseSchemas(test);
   }