Introduce index.query.max_nested_depth. fix opensearch-project#3268

Signed-off-by: Mikhail Khludnev <mkhl@apache.org>

server/src/main/java/org/opensearch/common/settings/IndexScopedSettings.java

@@ -150,6 +150,7 @@ public final class IndexScopedSettings extends AbstractScopedSettings {
                 IndexSettings.MAX_ADJACENCY_MATRIX_FILTERS_SETTING,
                 IndexSettings.MAX_ANALYZED_OFFSET_SETTING,
                 IndexSettings.MAX_TERMS_COUNT_SETTING,
+                IndexSettings.MAX_NESTED_QUERY_DEPTH_SETTING,
                 IndexSettings.INDEX_TRANSLOG_SYNC_INTERVAL_SETTING,
                 IndexSettings.DEFAULT_FIELD_SETTING,
                 IndexSettings.QUERY_STRING_LENIENT_SETTING,

server/src/main/java/org/opensearch/index/IndexSettings.java

@@ -269,6 +269,17 @@ public static IndexMergePolicy fromString(String text) {
         Property.IndexScope
     );
 
+    /**
+     * Index setting describing the maximum number of nested scopes in queries.
+     * The default maximum of 20. 0 means no nesting allowed.
+     */
+    public static final Setting<Integer> MAX_NESTED_QUERY_DEPTH_SETTING = Setting.intSetting(
+        "index.query.max_nested_depth",
+        20,
+        0,
+        Property.Dynamic,
+        Property.IndexScope
+    );
     /**
      * Index setting describing for NGramTokenizer and NGramTokenFilter
      * the maximum difference between
@@ -747,6 +758,8 @@ private void setRetentionLeaseMillis(final TimeValue retentionLease) {
     private volatile TimeValue searchIdleAfter;
     private volatile int maxAnalyzedOffset;
     private volatile int maxTermsCount;
+
+    private volatile int maxNestedQueryDepth;
     private volatile String defaultPipeline;
     private volatile String requiredPipeline;
     private volatile boolean searchThrottled;
@@ -902,6 +915,7 @@ public IndexSettings(final IndexMetadata indexMetadata, final Settings nodeSetti
         maxSlicesPerPit = scopedSettings.get(MAX_SLICES_PER_PIT);
         maxAnalyzedOffset = scopedSettings.get(MAX_ANALYZED_OFFSET_SETTING);
         maxTermsCount = scopedSettings.get(MAX_TERMS_COUNT_SETTING);
+        maxNestedQueryDepth = scopedSettings.get(MAX_NESTED_QUERY_DEPTH_SETTING);
         maxRegexLength = scopedSettings.get(MAX_REGEX_LENGTH_SETTING);
         this.tieredMergePolicyProvider = new TieredMergePolicyProvider(logger, this);
         this.logByteSizeMergePolicyProvider = new LogByteSizeMergePolicyProvider(logger, this);
@@ -1007,6 +1021,7 @@ public IndexSettings(final IndexMetadata indexMetadata, final Settings nodeSetti
         scopedSettings.addSettingsUpdateConsumer(MAX_REFRESH_LISTENERS_PER_SHARD, this::setMaxRefreshListeners);
         scopedSettings.addSettingsUpdateConsumer(MAX_ANALYZED_OFFSET_SETTING, this::setHighlightMaxAnalyzedOffset);
         scopedSettings.addSettingsUpdateConsumer(MAX_TERMS_COUNT_SETTING, this::setMaxTermsCount);
+        scopedSettings.addSettingsUpdateConsumer(MAX_NESTED_QUERY_DEPTH_SETTING, this::setMaxNestedQueryDepth);
         scopedSettings.addSettingsUpdateConsumer(MAX_SLICES_PER_SCROLL, this::setMaxSlicesPerScroll);
         scopedSettings.addSettingsUpdateConsumer(MAX_SLICES_PER_PIT, this::setMaxSlicesPerPit);
         scopedSettings.addSettingsUpdateConsumer(DEFAULT_FIELD_SETTING, this::setDefaultFields);
@@ -1517,6 +1532,17 @@ private void setMaxTermsCount(int maxTermsCount) {
         this.maxTermsCount = maxTermsCount;
     }
 
+    /**
+     * @return max level of nested queries and documents
+     */
+    public int getMaxNestedQueryDepth() {
+        return this.maxNestedQueryDepth;
+    }
+
+    private void setMaxNestedQueryDepth(int maxNestedQueryDepth) {
+        this.maxNestedQueryDepth = maxNestedQueryDepth;
+    }
+
     /**
      * Returns the maximum number of allowed script_fields to retrieve in a search request
      */

server/src/main/java/org/opensearch/index/query/NestedQueryBuilder.java

@@ -319,9 +319,9 @@ protected Query doToQuery(QueryShardContext context) throws IOException {
         }
 
         BitSetProducer previousParentFilter = context.getParentFilter();
+        context.nestedScope().nextLevel(nestedObjectMapper);
         try {
             context.setParentFilter(parentFilter);
-            context.nestedScope().nextLevel(nestedObjectMapper);
             innerQuery = this.query.toQuery(context);
         } finally {
             context.setParentFilter(previousParentFilter);

server/src/main/java/org/opensearch/index/query/QueryShardContext.java

@@ -256,7 +256,7 @@ private QueryShardContext(
         this.bitsetFilterCache = bitsetFilterCache;
         this.indexFieldDataService = indexFieldDataLookup;
         this.allowUnmappedFields = indexSettings.isDefaultAllowUnmappedFields();
-        this.nestedScope = new NestedScope();
+        this.nestedScope = new NestedScope(indexSettings);
         this.scriptService = scriptService;
         this.indexSettings = indexSettings;
         this.searcher = searcher;
@@ -270,7 +270,7 @@ private void reset() {
         allowUnmappedFields = indexSettings.isDefaultAllowUnmappedFields();
         this.lookup = null;
         this.namedQueries.clear();
-        this.nestedScope = new NestedScope();
+        this.nestedScope = new NestedScope(indexSettings);
     }
 
     public IndexAnalyzers getIndexAnalyzers() {

server/src/main/java/org/opensearch/index/query/support/NestedScope.java

@@ -33,6 +33,7 @@
 package org.opensearch.index.query.support;
 
 import org.opensearch.common.annotation.PublicApi;
+import org.opensearch.index.IndexSettings;
 import org.opensearch.index.mapper.ObjectMapper;
 
 import java.util.Deque;
@@ -47,6 +48,11 @@
 public final class NestedScope {
 
     private final Deque<ObjectMapper> levelStack = new LinkedList<>();
+    private final IndexSettings indexSettings;
+
+    public NestedScope(IndexSettings indexSettings) {
+        this.indexSettings = indexSettings;
+    }
 
     /**
      * @return For the current nested level returns the object mapper that belongs to that
@@ -60,7 +66,21 @@ public ObjectMapper getObjectMapper() {
      */
     public ObjectMapper nextLevel(ObjectMapper level) {
         ObjectMapper previous = levelStack.peek();
-        levelStack.push(level);
+        if (levelStack.size() < indexSettings.getMaxNestedQueryDepth()) {
+            levelStack.push(level);
+        } else {
+            throw new IllegalArgumentException(
+                "The depth of Nested Query is ["
+                    + (levelStack.size() + 1)
+                    + "] has exceeded "
+                    + "the allowed maximum of ["
+                    + indexSettings.getMaxNestedQueryDepth()
+                    + "]. "
+                    + "This maximum can be set by changing the ["
+                    + IndexSettings.MAX_NESTED_QUERY_DEPTH_SETTING.getKey()
+                    + "] index level setting."
+            );
+        }
         return previous;
     }
 

server/src/test/java/org/opensearch/index/query/NestedQueryBuilderTests.java

@@ -34,12 +34,16 @@
 
 import com.carrotsearch.randomizedtesting.generators.RandomPicks;
 
+import org.apache.lucene.search.BooleanClause;
+import org.apache.lucene.search.BooleanQuery;
+import org.apache.lucene.search.MatchAllDocsQuery;
 import org.apache.lucene.search.MatchNoDocsQuery;
 import org.apache.lucene.search.Query;
 import org.apache.lucene.search.join.ScoreMode;
 import org.opensearch.OpenSearchException;
 import org.opensearch.Version;
 import org.opensearch.action.admin.indices.mapping.put.PutMappingRequest;
+import org.opensearch.cluster.metadata.IndexMetadata;
 import org.opensearch.common.compress.CompressedXContent;
 import org.opensearch.common.settings.Settings;
 import org.opensearch.index.IndexSettings;
@@ -58,6 +62,7 @@
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.Map;
+import java.util.Optional;
 
 import static org.opensearch.index.IndexSettingsTests.newIndexMeta;
 import static org.opensearch.index.query.InnerHitBuilderTests.randomNestedInnerHits;
@@ -431,4 +436,88 @@ public void testSetParentFilterInContext() throws Exception {
         assertNull(queryShardContext.getParentFilter());
         verify(innerQueryBuilder).toQuery(queryShardContext);
     }
+
+    public void testNestedDepthProhibited() throws Exception {
+        doWithDepth(0, (context -> {
+            NestedQueryBuilder queryBuilder = new NestedQueryBuilder("nested1", new MatchAllQueryBuilder(), ScoreMode.None);
+            IllegalArgumentException e = expectThrows(IllegalArgumentException.class, () -> queryBuilder.toQuery(context));
+            assertEquals(
+                "The depth of Nested Query is [1] has exceeded the allowed maximum of [0]. This maximum can be set by changing the [index.query.max_nested_depth] index level setting.",
+                e.getMessage()
+            );
+        }));
+    }
+
+    public void testNestedDepthAllowed() throws Exception {
+        ThrowingConsumer<QueryShardContext> check = (context) -> {
+            NestedQueryBuilder queryBuilder = new NestedQueryBuilder("nested1", new MatchAllQueryBuilder(), ScoreMode.None);
+            OpenSearchToParentBlockJoinQuery blockJoinQuery = (OpenSearchToParentBlockJoinQuery) queryBuilder.toQuery(context);
+            Optional<BooleanClause> childLeg = ((BooleanQuery) blockJoinQuery.getChildQuery()).clauses()
+                .stream()
+                .filter(c -> c.getOccur() == BooleanClause.Occur.MUST)
+                .findFirst();
+            assertTrue(childLeg.isPresent());
+            assertEquals(new MatchAllDocsQuery(), childLeg.get().getQuery());
+        };
+        check.accept(createShardContext());
+        doWithDepth(randomIntBetween(1, 10), check);
+    }
+
+    public void testNestedDepthOnceOnly() throws Exception {
+        doWithDepth(1, (ctx) -> {
+            {
+                NestedQueryBuilder depth2 = new NestedQueryBuilder(
+                    "nested1",
+                    new NestedQueryBuilder("nested1", new MatchAllQueryBuilder(), ScoreMode.None),
+                    ScoreMode.None
+                );
+                IllegalArgumentException e = expectThrows(IllegalArgumentException.class, () -> depth2.toQuery(ctx));
+                assertEquals(
+                    "The depth of Nested Query is [2] has exceeded the allowed maximum of [1]. This maximum can be set by changing the [index.query.max_nested_depth] index level setting.",
+                    e.getMessage()
+                );
+            }
+            {
+                QueryBuilder mustBjqMustBjq = new BoolQueryBuilder().must(
+                    new NestedQueryBuilder("nested1", new MatchAllQueryBuilder(), ScoreMode.None)
+                ).must(new NestedQueryBuilder("nested1", new MatchAllQueryBuilder(), ScoreMode.None));
+                BooleanQuery bool = (BooleanQuery) mustBjqMustBjq.toQuery(ctx);
+                assertEquals(
+                    "Can parse joins one by one without breaching depth limit",
+                    2,
+                    bool.clauses().stream().filter(c -> c.getQuery() instanceof OpenSearchToParentBlockJoinQuery).count()
+                );
+            }
+        });
+    }
+
+    public void testUpdateMaxDepthSettings() throws Exception {
+        doWithDepth(2, (ctx) -> { assertEquals(ctx.getIndexSettings().getMaxNestedQueryDepth(), 2); });
+    }
+
+    void doWithDepth(int depth, ThrowingConsumer<QueryShardContext> test) throws Exception {
+        QueryShardContext context = createShardContext();
+        int defLimit = context.getIndexSettings().getMaxNestedQueryDepth();
+        assertTrue(defLimit > 0);
+        Settings updateSettings = Settings.builder()
+            .put(context.getIndexSettings().getSettings())
+            .put("index.query.max_nested_depth", depth)
+            .build();
+        context.getIndexSettings().updateIndexMetadata(IndexMetadata.builder("index").settings(updateSettings).build());
+        try {
+            test.accept(context);
+        } finally {
+            context.getIndexSettings()
+                .updateIndexMetadata(
+                    IndexMetadata.builder("index")
+                        .settings(
+                            Settings.builder()
+                                .put(context.getIndexSettings().getSettings())
+                                .put("index.query.max_nested_depth", defLimit)
+                                .build()
+                        )
+                        .build()
+                );
+        }
+    }
 }