Support intersecting multi-sets of queries with DocumentPermissions (#91151)

Since #81403, the Role class has been able to support multi-levels of
limiting (intersections). However, it was an oversight that the
underlying DocumentPermissions and FieldPermissions still do not support
it. They are still hardcoded to support up to 2 levels of intersection.
This PR now updates DocumentPermissions so it can support multi-level of
intersections. The similar change for FieldPermissions will be done in a
separate PR.

Author: ywangd

x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/security/authz/permission/DocumentPermissions.java

@@ -14,6 +14,7 @@
 import org.elasticsearch.common.bytes.BytesReference;
 import org.elasticsearch.common.io.stream.StreamOutput;
 import org.elasticsearch.common.lucene.search.Queries;
+import org.elasticsearch.core.Nullable;
 import org.elasticsearch.index.mapper.NestedLookup;
 import org.elasticsearch.index.query.QueryBuilder;
 import org.elasticsearch.index.query.QueryRewriteContext;
@@ -36,6 +37,7 @@
 import java.util.SortedSet;
 import java.util.TreeSet;
 import java.util.function.Function;
+import java.util.stream.Stream;
 
 import static org.apache.lucene.search.BooleanClause.Occur.FILTER;
 import static org.apache.lucene.search.BooleanClause.Occur.SHOULD;
@@ -46,58 +48,56 @@
  * queries are used as an additional filter.
  */
 public final class DocumentPermissions implements CacheKey {
-    // SortedSet because orders are important when they get serialised for request cache key
-    private final SortedSet<BytesReference> queries;
-    private final SortedSet<BytesReference> limitedByQueries;
-    private List<String> evaluatedQueries;
-    private List<String> evaluatedLimitedByQueries;
 
-    private static DocumentPermissions ALLOW_ALL = new DocumentPermissions();
+    @Nullable
+    private final List<Set<BytesReference>> listOfQueries;
+    @Nullable
+    private List<List<String>> listOfEvaluatedQueries;
 
-    DocumentPermissions() {
-        this.queries = null;
-        this.limitedByQueries = null;
+    private static final DocumentPermissions ALLOW_ALL = new DocumentPermissions();
+
+    private DocumentPermissions() {
+        this.listOfQueries = null;
     }
 
-    DocumentPermissions(Set<BytesReference> queries) {
-        this(queries, null);
+    private DocumentPermissions(Set<BytesReference> queries) {
+        assert queries != null && false == queries.isEmpty() : "null or empty queries not permitted";
+        this.listOfQueries = List.of(new TreeSet<>(queries));
     }
 
-    DocumentPermissions(Set<BytesReference> queries, Set<BytesReference> scopedByQueries) {
-        if (queries == null && scopedByQueries == null) {
-            throw new IllegalArgumentException("one of the queries or scoped queries must be provided");
-        }
-        this.queries = (queries != null) ? new TreeSet<>(queries) : null;
-        this.limitedByQueries = (scopedByQueries != null) ? new TreeSet<>(scopedByQueries) : null;
+    private DocumentPermissions(List<Set<BytesReference>> listOfQueries) {
+        assert listOfQueries != null && false == listOfQueries.isEmpty() : "null or empty list of queries not permitted";
+        assert listOfQueries.stream().allMatch(queries -> queries != null && false == queries.isEmpty())
+            : "null or empty queries not permitted";
+        // SortedSet because orders are important when they get serialised for request cache key
+        this.listOfQueries = listOfQueries.stream()
+            .map(queries -> queries instanceof SortedSet<BytesReference> ? queries : new TreeSet<>(queries))
+            .toList();
     }
 
-    public Set<BytesReference> getQueries() {
-        return queries == null ? null : Set.copyOf(queries);
+    public List<Set<BytesReference>> getListOfQueries() {
+        return listOfQueries;
     }
 
-    public Set<BytesReference> getLimitedByQueries() {
-        return limitedByQueries == null ? null : Set.copyOf(limitedByQueries);
+    public Set<BytesReference> getSingleSetOfQueries() {
+        assert listOfQueries != null && listOfQueries.size() == 1 : "the list of queries does not have a single member";
+        return listOfQueries.get(0);
     }
 
     /**
      * @return {@code true} if either queries or scoped queries are present for document level security else returns {@code false}
      */
     public boolean hasDocumentLevelPermissions() {
-        return queries != null || limitedByQueries != null;
+        return listOfQueries != null;
     }
 
     public boolean hasStoredScript() throws IOException {
-        if (queries != null) {
-            for (BytesReference q : queries) {
-                if (DLSRoleQueryValidator.hasStoredScript(q, NamedXContentRegistry.EMPTY)) {
-                    return true;
-                }
-            }
-        }
-        if (limitedByQueries != null) {
-            for (BytesReference q : limitedByQueries) {
-                if (DLSRoleQueryValidator.hasStoredScript(q, NamedXContentRegistry.EMPTY)) {
-                    return true;
+        if (listOfQueries != null) {
+            for (Set<BytesReference> queries : listOfQueries) {
+                for (BytesReference q : queries) {
+                    if (DLSRoleQueryValidator.hasStoredScript(q, NamedXContentRegistry.EMPTY)) {
+                        return true;
+                    }
                 }
             }
         }
@@ -125,35 +125,25 @@ public BooleanQuery filter(
     ) throws IOException {
         if (hasDocumentLevelPermissions()) {
             evaluateQueries(SecurityQueryTemplateEvaluator.wrap(user, scriptService));
-            BooleanQuery.Builder filter;
-            if (evaluatedQueries != null && evaluatedLimitedByQueries != null) {
-                filter = new BooleanQuery.Builder();
-                BooleanQuery.Builder scopedFilter = new BooleanQuery.Builder();
-                buildRoleQuery(shardId, searchExecutionContextProvider, evaluatedLimitedByQueries, scopedFilter);
-                filter.add(scopedFilter.build(), FILTER);
+            assert listOfEvaluatedQueries != null : "evaluated queries must not be null";
+            assert false == listOfEvaluatedQueries.isEmpty() : "evaluated queries must not be empty";
 
-                buildRoleQuery(shardId, searchExecutionContextProvider, evaluatedQueries, filter);
-            } else if (evaluatedQueries != null) {
-                filter = new BooleanQuery.Builder();
-                buildRoleQuery(shardId, searchExecutionContextProvider, evaluatedQueries, filter);
-            } else if (evaluatedLimitedByQueries != null) {
-                filter = new BooleanQuery.Builder();
-                buildRoleQuery(shardId, searchExecutionContextProvider, evaluatedLimitedByQueries, filter);
-            } else {
-                assert false : "one of queries and limited-by queries must be non-null";
-                return null;
+            BooleanQuery.Builder filter = new BooleanQuery.Builder();
+            for (int i = listOfEvaluatedQueries.size() - 1; i > 0; i--) {
+                final BooleanQuery.Builder scopedFilter = new BooleanQuery.Builder();
+                buildRoleQuery(shardId, searchExecutionContextProvider, listOfEvaluatedQueries.get(i), scopedFilter);
+                filter.add(scopedFilter.build(), FILTER);
             }
+            // TODO: All role queries can be filters
+            buildRoleQuery(shardId, searchExecutionContextProvider, listOfEvaluatedQueries.get(0), filter);
             return filter.build();
         }
         return null;
     }
 
     private void evaluateQueries(DlsQueryEvaluationContext context) {
-        if (queries != null && evaluatedQueries == null) {
-            evaluatedQueries = queries.stream().map(context::evaluate).toList();
-        }
-        if (limitedByQueries != null && evaluatedLimitedByQueries == null) {
-            evaluatedLimitedByQueries = limitedByQueries.stream().map(context::evaluate).toList();
+        if (listOfQueries != null && listOfEvaluatedQueries == null) {
+            listOfEvaluatedQueries = listOfQueries.stream().map(queries -> queries.stream().map(context::evaluate).toList()).toList();
         }
     }
 
@@ -216,18 +206,9 @@ static void failIfQueryUsesClient(QueryBuilder queryBuilder, QueryRewriteContext
      * @return {@link DocumentPermissions}
      */
     public static DocumentPermissions filteredBy(Set<BytesReference> queries) {
-        if (queries == null || queries.isEmpty()) {
-            throw new IllegalArgumentException("null or empty queries not permitted");
-        }
         return new DocumentPermissions(queries);
     }
 
-    /**
-     * Create {@link DocumentPermissions} with no restriction. The {@link #getQueries()}
-     * will return {@code null} in this case and {@link #hasDocumentLevelPermissions()}
-     * will be {@code false}
-     * @return {@link DocumentPermissions}
-     */
     public static DocumentPermissions allowAll() {
         return ALLOW_ALL;
     }
@@ -240,49 +221,41 @@ public static DocumentPermissions allowAll() {
      * @return instance of {@link DocumentPermissions}
      */
     public DocumentPermissions limitDocumentPermissions(DocumentPermissions limitedByDocumentPermissions) {
-        assert limitedByQueries == null && limitedByDocumentPermissions.limitedByQueries == null
-            : "nested scoping for document permissions is not permitted";
-        if (queries == null && limitedByDocumentPermissions.queries == null) {
+        if (hasDocumentLevelPermissions() && limitedByDocumentPermissions.hasDocumentLevelPermissions()) {
+            return new DocumentPermissions(
+                Stream.concat(getListOfQueries().stream(), limitedByDocumentPermissions.getListOfQueries().stream()).toList()
+            );
+        } else if (hasDocumentLevelPermissions()) {
+            return new DocumentPermissions(getListOfQueries());
+        } else if (limitedByDocumentPermissions.hasDocumentLevelPermissions()) {
+            return new DocumentPermissions(limitedByDocumentPermissions.getListOfQueries());
+        } else {
             return DocumentPermissions.allowAll();
         }
-        // TODO: should we apply the same logic here as FieldPermissions#limitFieldPermissions,
-        // i.e. treat limited-by as queries if original queries is null?
-        return new DocumentPermissions(queries, limitedByDocumentPermissions.queries);
     }
 
     @Override
     public String toString() {
-        return "DocumentPermissions [queries=" + queries + ", scopedByQueries=" + limitedByQueries + "]";
+        return "DocumentPermissions [listOfQueries=" + listOfQueries + "]";
     }
 
     @Override
     public void buildCacheKey(StreamOutput out, DlsQueryEvaluationContext context) throws IOException {
-        assert false == (queries == null && limitedByQueries == null) : "one of queries and limited-by queries must be non-null";
+        assert hasDocumentLevelPermissions() : "document permissions should not contribute to cache key when there is no DLS query";
         evaluateQueries(context);
-        if (evaluatedQueries != null) {
-            out.writeBoolean(true);
-            out.writeCollection(evaluatedQueries, StreamOutput::writeString);
-        } else {
-            out.writeBoolean(false);
-        }
-        if (evaluatedLimitedByQueries != null) {
-            out.writeBoolean(true);
-            out.writeCollection(evaluatedLimitedByQueries, StreamOutput::writeString);
-        } else {
-            out.writeBoolean(false);
-        }
+        out.writeCollection(listOfEvaluatedQueries, StreamOutput::writeStringCollection);
     }
 
     @Override
     public boolean equals(Object o) {
         if (this == o) return true;
         if (o == null || getClass() != o.getClass()) return false;
         DocumentPermissions that = (DocumentPermissions) o;
-        return Objects.equals(queries, that.queries) && Objects.equals(limitedByQueries, that.limitedByQueries);
+        return Objects.equals(listOfQueries, that.listOfQueries);
     }
 
     @Override
     public int hashCode() {
-        return Objects.hash(queries, limitedByQueries);
+        return Objects.hash(listOfQueries);
     }
 }

x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/termsenum/action/TransportTermsEnumAction.java

@@ -69,6 +69,7 @@
 
 import java.io.Closeable;
 import java.io.IOException;
+import java.io.UncheckedIOException;
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.HashMap;
@@ -429,13 +430,17 @@ private boolean canAccess(
                     Collections.emptyMap()
                 );
 
-                // Unfiltered access is allowed when both base role and limiting role allow it.
-                return hasMatchAllEquivalent(indexAccessControl.getDocumentPermissions().getQueries(), securityContext, queryShardContext)
-                    && hasMatchAllEquivalent(
-                        indexAccessControl.getDocumentPermissions().getLimitedByQueries(),
-                        securityContext,
-                        queryShardContext
-                    );
+                // Current user has potentially many roles and therefore potentially many queries
+                // defining sets of docs accessible
+                final List<Set<BytesReference>> listOfQueries = indexAccessControl.getDocumentPermissions().getListOfQueries();
+
+                // When the user is an API Key, its role is a limitedRole and its effective document permissions
+                // are intersections of the two sets of queries, one belongs to the API key itself and the other belongs
+                // to the owner user. To allow unfiltered access to termsDict, both sets of the queries must have
+                // the "all" permission, i.e. the query can be rewritten into a MatchAll query.
+                // The following code loop through both sets queries and returns true only when both of them
+                // have the "all" permission.
+                return listOfQueries.stream().allMatch(queries -> hasMatchAllEquivalent(queries, securityContext, queryShardContext));
             }
         }
         return true;
@@ -445,7 +450,7 @@ private boolean hasMatchAllEquivalent(
         Set<BytesReference> queries,
         SecurityContext securityContext,
         SearchExecutionContext queryShardContext
-    ) throws IOException {
+    ) {
         if (queries == null) {
             return true;
         }
@@ -458,7 +463,12 @@ private boolean hasMatchAllEquivalent(
                 queryShardContext.getParserConfig().registry(),
                 securityContext.getUser()
             );
-            QueryBuilder rewrittenQueryBuilder = Rewriteable.rewrite(queryBuilder, queryShardContext);
+            QueryBuilder rewrittenQueryBuilder;
+            try {
+                rewrittenQueryBuilder = Rewriteable.rewrite(queryBuilder, queryShardContext);
+            } catch (IOException e) {
+                throw new UncheckedIOException(e);
+            }
             if (rewrittenQueryBuilder instanceof MatchAllQueryBuilder) {
                 // One of the roles assigned has "all" permissions - allow unfettered access to termsDict
                 return true;