Capping heap size used by StatefulIndexPrivileges

Signed-off-by: Nils Bandener <[email protected]>

build.gradle

@@ -613,11 +613,8 @@ dependencies {
     implementation 'com.rfksystems:blake2b:2.0.0'
     implementation 'com.password4j:password4j:1.8.2'
 
-    // Action privileges
-    implementation 'com.selectivem.collections:checklists:1.2.1'
-    implementation 'com.selectivem.collections:compact-subsets:1.2.1'
-    implementation 'com.selectivem.collections:interfaces:1.2.1'
-    runtimeOnly 'com.selectivem.collections:backing-collections:1.2.1'
+    // Action privileges: check tables and compact collections
+    implementation 'com.selectivem.collections:special-collections-complete:1.3.0'
 
     //JWT
     implementation "io.jsonwebtoken:jjwt-api:${jjwt_version}"

src/main/java/org/opensearch/security/OpenSearchSecurityPlugin.java

@@ -165,6 +165,7 @@
 import org.opensearch.security.http.NonSslHttpServerTransport;
 import org.opensearch.security.http.XFFResolver;
 import org.opensearch.security.identity.SecurityTokenManager;
+import org.opensearch.security.privileges.ActionPrivileges;
 import org.opensearch.security.privileges.PrivilegesEvaluator;
 import org.opensearch.security.privileges.PrivilegesInterceptor;
 import org.opensearch.security.privileges.RestLayerPrivilegesEvaluator;
@@ -2009,6 +2010,10 @@ public List<Setting<?>> getSettings() {
                     Property.Filtered
                 )
             );
+
+            // Privileges evaluation
+            settings.add(ActionPrivileges.PRECOMPUTED_PRIVILEGES_MAX_HEAP_SIZE);
+            settings.add(ActionPrivileges.PRECOMPUTED_PRIVILEGES_INCLUDE_INDICES);
         }
 
         return settings;

src/main/java/org/opensearch/security/privileges/ActionPrivileges.java

@@ -20,6 +20,7 @@
 import java.util.function.Supplier;
 import java.util.stream.Collectors;
 
+import com.google.common.base.Strings;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
@@ -31,13 +32,18 @@
 import org.opensearch.cluster.metadata.IndexAbstraction;
 import org.opensearch.cluster.metadata.IndexMetadata;
 import org.opensearch.cluster.metadata.IndexNameExpressionResolver;
+import org.opensearch.common.settings.Setting;
+import org.opensearch.common.settings.Settings;
+import org.opensearch.core.common.unit.ByteSizeUnit;
+import org.opensearch.core.common.unit.ByteSizeValue;
 import org.opensearch.security.resolver.IndexResolverReplacer;
 import org.opensearch.security.securityconf.FlattenedActionGroups;
 import org.opensearch.security.securityconf.impl.SecurityDynamicConfiguration;
 import org.opensearch.security.securityconf.impl.v7.RoleV7;
 import org.opensearch.security.support.WildcardMatcher;
 
 import com.selectivem.collections.CheckTable;
+import com.selectivem.collections.CompactMapGroupBuilder;
 import com.selectivem.collections.DeduplicatingCompactSubSetBuilder;
 import com.selectivem.collections.ImmutableCompactSubSet;
 
@@ -59,6 +65,38 @@
  * then removed.
  */
 public class ActionPrivileges {
+
+    /**
+     * This setting controls the allowed heap size of the precomputed index privileges (in the inner class StatefulIndexPrivileges).
+     * If the size of the indices exceed the amount of bytes configured here, it will be truncated. Privileges evaluation will
+     * continue to work correctly, but it will be slower.
+     * <p>
+     * This settings defaults to 10 MB. This is a generous limit. Experiments have shown that an example setup with
+     * 10,000 indices and 1,000 roles requires about 1 MB of heap. 100,000 indices and 100 roles require about 9 MB of heap.
+     * (Of course, these numbers can vary widely based on the actual role configuration).
+     * <p>
+     * The setting plugins.security.privileges_evaluation.precomputed_privileges.include_indices can be used to control
+     * for which indices the precomputed privileges shall be created. This allows to lower the heap utilization.
+     */
+    public static Setting<ByteSizeValue> PRECOMPUTED_PRIVILEGES_MAX_HEAP_SIZE = Setting.memorySizeSetting(
+        "plugins.security.privileges_evaluation.precomputed_privileges.max_heap_size",
+        new ByteSizeValue(10, ByteSizeUnit.MB),
+        Setting.Property.NodeScope
+    );
+
+    /**
+     * Determines the indices which shall be included in the precomputed index privileges. Included indices get
+     * the fasted privilege evaluation.
+     * <p>
+     * You can use patterns like "index_*".
+     * <p>
+     * Defaults to all indices.
+     */
+    public static Setting<String> PRECOMPUTED_PRIVILEGES_INCLUDE_INDICES = Setting.simpleString(
+        "plugins.security.privileges_evaluation.precomputed_privileges.include_indices",
+        Setting.Property.NodeScope
+    );
+
     private static final Logger log = LogManager.getLogger(ActionPrivileges.class);
 
     private final ClusterPrivileges cluster;
@@ -68,6 +106,8 @@ public class ActionPrivileges {
     private final ImmutableSet<String> wellKnownClusterActions;
     private final ImmutableSet<String> wellKnownIndexActions;
     private final Supplier<Map<String, IndexAbstraction>> indexMetadataSupplier;
+    private final ByteSizeValue statefulIndexMaxHeapSize;
+    private final WildcardMatcher statefulIndexIncludeIndices;
 
     private volatile StatefulIndexPrivileges statefulIndex;
 
@@ -78,6 +118,7 @@ public ActionPrivileges(
         SecurityDynamicConfiguration<?> rolesUnsafe,
         FlattenedActionGroups actionGroups,
         Supplier<Map<String, IndexAbstraction>> indexMetadataSupplier,
+        Settings settings,
         ImmutableSet<String> wellKnownClusterActions,
         ImmutableSet<String> wellKnownIndexActions,
         ImmutableSet<String> explicitlyRequiredIndexActions
@@ -91,17 +132,24 @@ public ActionPrivileges(
         this.wellKnownClusterActions = wellKnownClusterActions;
         this.wellKnownIndexActions = wellKnownIndexActions;
         this.indexMetadataSupplier = indexMetadataSupplier;
+        this.statefulIndexMaxHeapSize = PRECOMPUTED_PRIVILEGES_MAX_HEAP_SIZE.get(settings);
+        String statefulIndexIncludeIndices = PRECOMPUTED_PRIVILEGES_INCLUDE_INDICES.get(settings);
+        this.statefulIndexIncludeIndices = Strings.isNullOrEmpty(statefulIndexIncludeIndices)
+            ? null
+            : WildcardMatcher.from(statefulIndexIncludeIndices);
     }
 
     public ActionPrivileges(
         SecurityDynamicConfiguration<?> roles,
         FlattenedActionGroups actionGroups,
-        Supplier<Map<String, IndexAbstraction>> indexMetadataSupplier
+        Supplier<Map<String, IndexAbstraction>> indexMetadataSupplier,
+        Settings settings
     ) {
         this(
             roles,
             actionGroups,
             indexMetadataSupplier,
+            settings,
             WellKnownActions.CLUSTER_ACTIONS,
             WellKnownActions.INDEX_ACTIONS,
             WellKnownActions.EXPLICITLY_REQUIRED_INDEX_ACTIONS
@@ -175,10 +223,20 @@ public PrivilegesEvaluatorResponse hasExplicitIndexPrivilege(
     public void updateStatefulIndexPrivileges(Map<String, IndexAbstraction> indices) {
         StatefulIndexPrivileges statefulIndex = this.statefulIndex;
 
-        indices = StatefulIndexPrivileges.relevantOnly(indices);
+        indices = StatefulIndexPrivileges.relevantOnly(indices, statefulIndexIncludeIndices);
 
         if (statefulIndex == null || !statefulIndex.indices.equals(indices)) {
-            this.statefulIndex = new StatefulIndexPrivileges(roles, actionGroups, wellKnownIndexActions, indices);
+            this.statefulIndex = new StatefulIndexPrivileges(roles, actionGroups, wellKnownIndexActions, indices, statefulIndexMaxHeapSize);
+        }
+    }
+
+    int getEstimatedStatefulIndexByteSize() {
+        StatefulIndexPrivileges statefulIndex = this.statefulIndex;
+
+        if (statefulIndex != null) {
+            return statefulIndex.estimatedByteSize;
+        } else {
+            return 0;
         }
     }
 
@@ -238,7 +296,6 @@ static class ClusterPrivileges {
                 roles.getCEntries().keySet()
             );
             Map<String, DeduplicatingCompactSubSetBuilder.SubSetBuilder<String>> actionToRoles = new HashMap<>();
-            Map<String, Set<String>> explicitPermissionToRoles = new HashMap<>();
             ImmutableSet.Builder<String> rolesWithWildcardPermissions = ImmutableSet.builder();
             ImmutableMap.Builder<String, WildcardMatcher> rolesToActionMatcher = ImmutableMap.builder();
 
@@ -787,6 +844,8 @@ static class StatefulIndexPrivileges {
          */
         private final ImmutableSet<String> wellKnownIndexActions;
 
+        private final int estimatedByteSize;
+
         /**
          * Creates pre-computed index privileges based on the given parameters.
          * <p>
@@ -798,16 +857,20 @@ static class StatefulIndexPrivileges {
             SecurityDynamicConfiguration<RoleV7> roles,
             FlattenedActionGroups actionGroups,
             ImmutableSet<String> wellKnownIndexActions,
-            Map<String, IndexAbstraction> indices
+            Map<String, IndexAbstraction> indices,
+            ByteSizeValue statefulIndexMaxHeapSize
         ) {
+            Map<
+                String,
+                CompactMapGroupBuilder.MapBuilder<String, DeduplicatingCompactSubSetBuilder.SubSetBuilder<String>>> actionToIndexToRoles =
+                    new HashMap<>();
+            CompactMapGroupBuilder<String, DeduplicatingCompactSubSetBuilder.SubSetBuilder<String>> indexMapBuilder =
+                new CompactMapGroupBuilder<>(indices.keySet());
             DeduplicatingCompactSubSetBuilder<String> roleSetBuilder = new DeduplicatingCompactSubSetBuilder<>(
                 roles.getCEntries().keySet()
             );
 
-            Map<String, Map<String, DeduplicatingCompactSubSetBuilder.SubSetBuilder<String>>> actionToIndexToRoles = new HashMap<>();
-            int leafs = 0; // This counts the number of leafs in the actionToIndexToRoles data structure. Useful for estimating the size.
-
-            for (Map.Entry<String, RoleV7> entry : roles.getCEntries().entrySet()) {
+            top: for (Map.Entry<String, RoleV7> entry : roles.getCEntries().entrySet()) {
                 try {
                     String roleName = entry.getKey();
                     RoleV7 role = entry.getValue();
@@ -839,23 +902,34 @@ static class StatefulIndexPrivileges {
                                 Map.Entry::getKey
                             )) {
                                 for (String action : matchedActions) {
-                                    Map<String, DeduplicatingCompactSubSetBuilder.SubSetBuilder<String>> indexToRoles = actionToIndexToRoles
-                                        .computeIfAbsent(action, k -> new HashMap<>());
+                                    CompactMapGroupBuilder.MapBuilder<
+                                        String,
+                                        DeduplicatingCompactSubSetBuilder.SubSetBuilder<String>> indexToRoles = actionToIndexToRoles
+                                            .computeIfAbsent(
+                                                action,
+                                                k -> indexMapBuilder.createMapBuilder((k2) -> roleSetBuilder.createSubSetBuilder())
+                                            );
 
-                                    indexToRoles.computeIfAbsent(indicesEntry.getKey(), k -> roleSetBuilder.createSubSetBuilder())
-                                        .add(roleName);
-                                    leafs++;
+                                    indexToRoles.get(indicesEntry.getKey()).add(roleName);
 
                                     if (indicesEntry.getValue() instanceof IndexAbstraction.Alias) {
                                         // For aliases we additionally add the sub-indices to the privilege map
                                         for (IndexMetadata subIndex : indicesEntry.getValue().getIndices()) {
-                                            indexToRoles.computeIfAbsent(
-                                                subIndex.getIndex().getName(),
-                                                k -> roleSetBuilder.createSubSetBuilder()
-                                            ).add(roleName);
-                                            leafs++;
+                                            indexToRoles.get(subIndex.getIndex().getName()).add(roleName);
                                         }
                                     }
+
+                                    if (roleSetBuilder.getEstimatedByteSize() + indexMapBuilder
+                                        .getEstimatedByteSize() > statefulIndexMaxHeapSize.getBytes()) {
+                                        log.info(
+                                            "Size of precomputed index privileges exceeds configured limit ({}). Using capped data structure."
+                                                + "This might lead to slightly lower performance during privilege evaluation. Consider raising {} or limiting the performance critical indices using {}.",
+                                            statefulIndexMaxHeapSize,
+                                            PRECOMPUTED_PRIVILEGES_MAX_HEAP_SIZE.getKey(),
+                                            PRECOMPUTED_PRIVILEGES_INCLUDE_INDICES.getKey()
+                                        );
+                                        break top;
+                                    }
                                 }
                             }
                         }
@@ -867,22 +941,15 @@ static class StatefulIndexPrivileges {
 
             DeduplicatingCompactSubSetBuilder.Completed<String> completedRoleSetBuilder = roleSetBuilder.build();
 
-            log.debug("StatefulIndexPermissions data structure contains {} leafs", leafs);
+            this.estimatedByteSize = roleSetBuilder.getEstimatedByteSize() + indexMapBuilder.getEstimatedByteSize();
+            log.debug("Estimated size of StatefulIndexPermissions data structure: ", this.estimatedByteSize);
 
             this.actionToIndexToRoles = actionToIndexToRoles.entrySet()
                 .stream()
                 .collect(
                     ImmutableMap.toImmutableMap(
                         entry -> entry.getKey(),
-                        entry -> entry.getValue()
-                            .entrySet()
-                            .stream()
-                            .collect(
-                                ImmutableMap.toImmutableMap(
-                                    entry2 -> entry2.getKey(),
-                                    entry2 -> entry2.getValue().build(completedRoleSetBuilder)
-                                )
-                            )
+                        entry -> entry.getValue().build(subSetBuilder -> subSetBuilder.build(completedRoleSetBuilder))
                     )
                 );
 
@@ -966,10 +1033,15 @@ private String backingIndexToDataStream(String index, Map<String, IndexAbstracti
          * Filters the given index abstraction map to only contain entries that are relevant the for stateful class.
          * This removes data stream backing indices and closes indices.
          */
-        static Map<String, IndexAbstraction> relevantOnly(Map<String, IndexAbstraction> indices) {
+        static Map<String, IndexAbstraction> relevantOnly(Map<String, IndexAbstraction> indices, WildcardMatcher includeIndices) {
             boolean doFilter = false;
 
             for (IndexAbstraction indexAbstraction : indices.values()) {
+                if (includeIndices != null && !includeIndices.test(indexAbstraction.getName())) {
+                    doFilter = true;
+                    break;
+                }
+
                 if (indexAbstraction instanceof IndexAbstraction.Index) {
                     if (indexAbstraction.getParentDataStream() != null
                         || indexAbstraction.getWriteIndex().getState() == IndexMetadata.State.CLOSE) {
@@ -986,6 +1058,10 @@ static Map<String, IndexAbstraction> relevantOnly(Map<String, IndexAbstraction>
             ImmutableMap.Builder<String, IndexAbstraction> builder = ImmutableMap.builder();
 
             for (IndexAbstraction indexAbstraction : indices.values()) {
+                if (includeIndices != null && !includeIndices.test(indexAbstraction.getName())) {
+                    continue;
+                }
+
                 if (indexAbstraction instanceof IndexAbstraction.Index) {
                     if (indexAbstraction.getParentDataStream() == null
                         && indexAbstraction.getWriteIndex().getState() != IndexMetadata.State.CLOSE) {

src/main/java/org/opensearch/security/privileges/PrivilegesEvaluator.java

@@ -149,6 +149,7 @@ public class PrivilegesEvaluator {
     private final PitPrivilegesEvaluator pitPrivilegesEvaluator;
     private DynamicConfigModel dcm;
     private final NamedXContentRegistry namedXContentRegistry;
+    private final Settings settings;
     private volatile ActionPrivileges actionPrivileges;
 
     public PrivilegesEvaluator(
@@ -172,6 +173,7 @@ public PrivilegesEvaluator(
         this.threadContext = threadContext;
         this.privilegesInterceptor = privilegesInterceptor;
         this.clusterStateSupplier = clusterStateSupplier;
+        this.settings = settings;
 
         this.checkSnapshotRestoreWritePrivileges = settings.getAsBoolean(
             ConfigConstants.SECURITY_CHECK_SNAPSHOT_RESTORE_WRITE_PRIVILEGES,
@@ -235,7 +237,8 @@ void updateConfiguration(
             ActionPrivileges actionPrivileges = new ActionPrivileges(
                 DynamicConfigFactory.addStatics(rolesConfiguration.clone()),
                 flattenedActionGroups,
-                () -> clusterStateSupplier.get().metadata().getIndicesLookup()
+                () -> clusterStateSupplier.get().metadata().getIndicesLookup(),
+                settings
             );
             actionPrivileges.updateStatefulIndexPrivileges(clusterStateSupplier.get().metadata().getIndicesLookup());
             this.actionPrivileges = actionPrivileges;

src/main/java/org/opensearch/security/privileges/TermsAggregationEvaluator.java

@@ -54,7 +54,8 @@ public class TermsAggregationEvaluator {
         MultiGetAction.NAME,
         GetAction.NAME,
         SearchAction.NAME,
-        FieldCapabilitiesAction.NAME );
+        FieldCapabilitiesAction.NAME
+    );
 
     private static final QueryBuilder NONE_QUERY = new MatchNoneQueryBuilder();