Enforces limit of query tree size if configured

5 files changed, 71 insertions, 2 deletions

diff --git a/container-search/abi-spec.json b/container-search/abi-spec.json
index 183bb33b4f4..40783c3ce27 100644
--- a/container-search/abi-spec.json
+++ b/container-search/abi-spec.json
@@ -5285,6 +5285,7 @@
       "public boolean isEmpty()",
       "public com.yahoo.prelude.query.Item and(com.yahoo.prelude.query.Item)",
       "public static java.util.List getPositiveTerms(com.yahoo.prelude.query.Item)",
+      "public int getTreeSize()",
       "public bridge synthetic com.yahoo.prelude.query.CompositeItem clone()",
       "public bridge synthetic com.yahoo.prelude.query.Item clone()",
       "public bridge synthetic java.lang.Object clone()"
diff --git a/container-search/src/main/java/com/yahoo/prelude/query/QueryCanonicalizer.java b/container-search/src/main/java/com/yahoo/prelude/query/QueryCanonicalizer.java
index 1f30833b3db..14e316c1877 100644
--- a/container-search/src/main/java/com/yahoo/prelude/query/QueryCanonicalizer.java
+++ b/container-search/src/main/java/com/yahoo/prelude/query/QueryCanonicalizer.java
@@ -1,6 +1,7 @@
 // Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
 package com.yahoo.prelude.query;
 
+import com.yahoo.processing.request.CompoundName;
 import com.yahoo.search.Query;
 import com.yahoo.search.query.QueryTree;
 
@@ -19,6 +20,8 @@ public class QueryCanonicalizer {
     /** The name of the operation performed by this, for use in search chain ordering */
     public static final String queryCanonicalization = "queryCanonicalization";
 
+    private static final CompoundName MAX_QUERY_ITEMS = new CompoundName("maxQueryItems");
+
     /**
      * Validates this query and carries out possible operations on this query
      * which simplifies it without changing its semantics.
@@ -26,7 +29,17 @@ public class QueryCanonicalizer {
      * @return null if the query is valid, an error message if it is invalid
      */
     public static String canonicalize(Query query) {
-        return canonicalize(query.getModel().getQueryTree());
+        Integer maxQueryItems = query.properties().getInteger(MAX_QUERY_ITEMS, Integer.MAX_VALUE);
+        return canonicalize(query.getModel().getQueryTree(), maxQueryItems);
+    }
+
+    /**
+     * Canonicalizes this query, allowing any query tree size
+     *
+     * @return null if the query is valid, an error message if it is invalid
+     */
+    public static String canonicalize(QueryTree queryTree) {
+        return canonicalize(queryTree, Integer.MAX_VALUE);
     }
 
     /**
@@ -34,10 +47,11 @@ public class QueryCanonicalizer {
      * 
      * @return null if the query is valid, an error message if it is invalid
      */
-    public static String canonicalize(QueryTree query) {
+    private static String canonicalize(QueryTree query, Integer maxQueryItems) {
         ListIterator<Item> rootItemIterator = query.getItemIterator();
         CanonicalizationResult result = recursivelyCanonicalize(rootItemIterator.next(), rootItemIterator);
         if (query.isEmpty() && ! result.isError()) result = CanonicalizationResult.error("No query");
+        if (query.getTreeSize() > maxQueryItems) result = CanonicalizationResult.error("Query tree exceeds allowed item count");
         return result.error().orElse(null); // preserve old API, unfortunately
     }
 
diff --git a/container-search/src/main/java/com/yahoo/search/query/QueryTree.java b/container-search/src/main/java/com/yahoo/search/query/QueryTree.java
index 2fb16ef503f..dab481dc842 100644
--- a/container-search/src/main/java/com/yahoo/search/query/QueryTree.java
+++ b/container-search/src/main/java/com/yahoo/search/query/QueryTree.java
@@ -7,6 +7,7 @@ import java.nio.ByteBuffer;
 import java.util.ArrayList;
 import java.util.Iterator;
 import java.util.List;
+import java.util.ListIterator;
 
 /**
  * The root node of a query tree. This is always present above the actual semantic root to ease query manipulation,
@@ -179,4 +180,23 @@ public class QueryTree extends CompositeItem {
         }
     }
 
+    /**
+     * @return The total number of items in this query tree
+     */
+    public int getTreeSize() {
+        if (isEmpty()) return 0;
+        return(countItemsRecursively(getItemIterator().next()));
+    }
+
+    private int countItemsRecursively(Item item) {
+        int children = 0;
+        if (item instanceof CompositeItem) {
+            CompositeItem composite = (CompositeItem)item;
+            for (ListIterator<Item> i = composite.getItemIterator(); i.hasNext(); ) {
+                children += countItemsRecursively(i.next());
+            }
+        }
+        return children + 1;
+    }
+
 }
diff --git a/container-search/src/test/java/com/yahoo/prelude/query/test/QueryCanonicalizerTestCase.java b/container-search/src/test/java/com/yahoo/prelude/query/test/QueryCanonicalizerTestCase.java
index 1d2f92063fe..a5ff92294e8 100644
--- a/container-search/src/test/java/com/yahoo/prelude/query/test/QueryCanonicalizerTestCase.java
+++ b/container-search/src/test/java/com/yahoo/prelude/query/test/QueryCanonicalizerTestCase.java
@@ -490,6 +490,18 @@ public class QueryCanonicalizerTestCase {
         assertFalse(shoe.usePositionData());
     }
 
+    @Test
+    public void queryTreeExceedsAllowedSize() {
+        Query query = new Query();
+        QueryTree tree = query.getModel().getQueryTree();
+        tree.setRoot(new WordItem("A"));
+        tree.and(new WordItem("B"));
+
+        assertNull(QueryCanonicalizer.canonicalize(query));
+        query.properties().set("maxQueryItems", 2);
+        assertEquals("Query tree exceeds allowed item count", QueryCanonicalizer.canonicalize(query));
+    }
+
     private void assertCanonicalized(String canonicalForm, String expectedError, Item root) {
         Query query = new Query();
         query.getModel().getQueryTree().setRoot(root);
diff --git a/container-search/src/test/java/com/yahoo/search/query/QueryTreeTest.java b/container-search/src/test/java/com/yahoo/search/query/QueryTreeTest.java
index 054b752c067..2bff9ba1707 100644
--- a/container-search/src/test/java/com/yahoo/search/query/QueryTreeTest.java
+++ b/container-search/src/test/java/com/yahoo/search/query/QueryTreeTest.java
@@ -2,6 +2,7 @@
 package com.yahoo.search.query;
 
 import com.yahoo.prelude.query.NotItem;
+import com.yahoo.prelude.query.NullItem;
 import com.yahoo.prelude.query.WordItem;
 import org.junit.Assert;
 import org.junit.Test;
@@ -41,4 +42,25 @@ public class QueryTreeTest {
          assertEquals("+(AND p1 p2) -n1.1 -n1.2 -n2.1 -n2.2", tree.toString());
      }
 
+     @Test
+     public void getCorrectTreeSize() {
+        QueryTree nullTree = new QueryTree(new NullItem());
+        assertEquals(0, nullTree.getTreeSize());
+
+         NotItem not1 = new NotItem();
+         not1.addPositiveItem(new WordItem("p1"));
+         not1.addNegativeItem(new WordItem("n1.1"));
+         not1.addNegativeItem(new WordItem("n1.2"));
+
+         NotItem not2 = new NotItem();
+         not2.addPositiveItem(new WordItem("p2"));
+         not2.addNegativeItem(new WordItem("n2.1"));
+         not2.addNegativeItem(new WordItem("n2.2"));
+
+         QueryTree tree = new QueryTree(not1);
+         tree.and(not2);
+
+         assertEquals(8, tree.getTreeSize());
+     }
+
 }