From c7d5a54b42a88fc89dff6942dde99b79c3e36cdf Mon Sep 17 00:00:00 2001
From: Jim Ferenczi <jim.ferenczi@elastic.co>
Date: Thu, 1 Feb 2018 09:01:48 +0100
Subject: [PATCH] Fix AIOOB on indexed geo_shape query (#28458)

This change fixes a possible AIOOB during the parsing of the document that contains the indexed shape.
This change ensures that the parsing does not continue when the field that contains the shape has been found.

Closes #28456
---
 .../index/query/GeoShapeQueryBuilder.java     |  1 +
 .../query/GeoShapeQueryBuilderTests.java      | 26 ++++++++++++++-----
 2 files changed, 20 insertions(+), 7 deletions(-)

diff --git a/server/src/main/java/org/elasticsearch/index/query/GeoShapeQueryBuilder.java b/server/src/main/java/org/elasticsearch/index/query/GeoShapeQueryBuilder.java
index c4a6f7fa6b7..5fb97769463 100644
--- a/server/src/main/java/org/elasticsearch/index/query/GeoShapeQueryBuilder.java
+++ b/server/src/main/java/org/elasticsearch/index/query/GeoShapeQueryBuilder.java
@@ -409,6 +409,7 @@ public class GeoShapeQueryBuilder extends AbstractQueryBuilder<GeoShapeQueryBuil
                                     parser.nextToken();
                                     if (++currentPathSlot == pathElements.length) {
                                         listener.onResponse(ShapeParser.parse(parser));
+                                        return;
                                     }
                                 } else {
                                     parser.nextToken();
diff --git a/server/src/test/java/org/elasticsearch/index/query/GeoShapeQueryBuilderTests.java b/server/src/test/java/org/elasticsearch/index/query/GeoShapeQueryBuilderTests.java
index 6ddcd671a40..96d756f0083 100644
--- a/server/src/test/java/org/elasticsearch/index/query/GeoShapeQueryBuilderTests.java
+++ b/server/src/test/java/org/elasticsearch/index/query/GeoShapeQueryBuilderTests.java
@@ -119,6 +119,7 @@ public class GeoShapeQueryBuilderTests extends AbstractQueryTestCase<GeoShapeQue
             XContentBuilder builder = XContentFactory.jsonBuilder().prettyPrint();
             builder.startObject();
             builder.field(expectedShapePath, indexedShapeToReturn);
+            builder.field(randomAlphaOfLengthBetween(10, 20), "something");
             builder.endObject();
             json = builder.string();
         } catch (IOException ex) {
@@ -227,13 +228,7 @@ public class GeoShapeQueryBuilderTests extends AbstractQueryTestCase<GeoShapeQue
 
     @Override
     public void testMustRewrite() throws IOException {
-        GeoShapeQueryBuilder sqb;
-        do {
-            sqb = doCreateTestQueryBuilder();
-            // do this until we get one without a shape
-        } while (sqb.shape() != null);
-
-        GeoShapeQueryBuilder query = sqb;
+        GeoShapeQueryBuilder query = doCreateTestQueryBuilder(true);
 
         UnsupportedOperationException e = expectThrows(UnsupportedOperationException.class, () -> query.toQuery(createShardContext()));
         assertEquals("query must be rewritten first", e.getMessage());
@@ -244,6 +239,23 @@ public class GeoShapeQueryBuilderTests extends AbstractQueryTestCase<GeoShapeQue
         assertEquals(geoShapeQueryBuilder, rewrite);
     }
 
+    public void testMultipleRewrite() throws IOException {
+        GeoShapeQueryBuilder shape = doCreateTestQueryBuilder(true);
+        QueryBuilder builder = new BoolQueryBuilder()
+            .should(shape)
+            .should(shape);
+
+        builder = rewriteAndFetch(builder, createShardContext());
+
+        GeoShapeQueryBuilder expectedShape = new GeoShapeQueryBuilder(GEO_SHAPE_FIELD_NAME, indexedShapeToReturn);
+        expectedShape.strategy(shape.strategy());
+        expectedShape.relation(shape.relation());
+        QueryBuilder expected = new BoolQueryBuilder()
+            .should(expectedShape)
+            .should(expectedShape);
+        assertEquals(expected, builder);
+    }
+
     public void testIgnoreUnmapped() throws IOException {
         ShapeType shapeType = ShapeType.randomType(random());
         ShapeBuilder shape = RandomShapeGenerator.createShapeWithin(random(), null, shapeType);