From 48a68365bbb2ecfbc8b61c798eb312f99f0a5afc Mon Sep 17 00:00:00 2001
From: Zhao Yang <zhaoyangsingapore@gmail.com>
Date: Tue, 16 Apr 2019 09:32:40 +0800
Subject: [PATCH] Fix OneDimensionBKDWriter valueCount validation

Signed-off-by: Adrien Grand <jpountz@gmail.com>
---
 lucene/CHANGES.txt                            |  3 +
 .../org/apache/lucene/util/bkd/BKDWriter.java |  4 +-
 .../org/apache/lucene/util/bkd/TestBKD.java   | 91 +++++++++++++++++++
 3 files changed, 96 insertions(+), 2 deletions(-)

diff --git a/lucene/CHANGES.txt b/lucene/CHANGES.txt
index e80689d8bed..dae291c8c9d 100644
--- a/lucene/CHANGES.txt
+++ b/lucene/CHANGES.txt
@@ -90,6 +90,9 @@ Bug fixes
 * LUCENE-8754: Fix ConcurrentModificationException in SegmentInfo if
   attributes are accessed in MergePolicy while the merge is running (Simon Willnauer)
 
+* LUCENE-8765: Fixed validation of the number of added points in KD trees.
+  (Zhao Yang via Adrien Grand)
+
 Improvements
 
 * LUCENE-8673: Use radix partitioning when merging dimensional points instead
diff --git a/lucene/core/src/java/org/apache/lucene/util/bkd/BKDWriter.java b/lucene/core/src/java/org/apache/lucene/util/bkd/BKDWriter.java
index 7a1aec4d921..659935cc929 100644
--- a/lucene/core/src/java/org/apache/lucene/util/bkd/BKDWriter.java
+++ b/lucene/core/src/java/org/apache/lucene/util/bkd/BKDWriter.java
@@ -551,8 +551,8 @@ public class BKDWriter implements Closeable {
       docsSeen.set(docID);
       leafCount++;
 
-      if (valueCount > totalPointCount) {
-        throw new IllegalStateException("totalPointCount=" + totalPointCount + " was passed when we were created, but we just hit " + pointCount + " values");
+      if (valueCount + leafCount > totalPointCount) {
+        throw new IllegalStateException("totalPointCount=" + totalPointCount + " was passed when we were created, but we just hit " + pointCount + leafCount + " values");
       }
 
       if (leafCount == maxPointsInLeafNode) {
diff --git a/lucene/core/src/test/org/apache/lucene/util/bkd/TestBKD.java b/lucene/core/src/test/org/apache/lucene/util/bkd/TestBKD.java
index 98449adce58..ce76ca47072 100644
--- a/lucene/core/src/test/org/apache/lucene/util/bkd/TestBKD.java
+++ b/lucene/core/src/test/org/apache/lucene/util/bkd/TestBKD.java
@@ -24,6 +24,7 @@ import java.util.Arrays;
 import java.util.BitSet;
 import java.util.List;
 
+import org.apache.lucene.codecs.MutablePointValues;
 import org.apache.lucene.index.CorruptIndexException;
 import org.apache.lucene.index.MergeState;
 import org.apache.lucene.index.PointValues.IntersectVisitor;
@@ -1244,4 +1245,94 @@ public class TestBKD extends LuceneTestCase {
     pointsIn.close();
     dir.close();
   }
+
+  public void testTotalPointCountValidation() throws IOException {
+    Directory dir = newDirectory();
+    final int numValues = 10;
+    final int numPointsAdded = 50; // exceeds totalPointCount
+    final int numBytesPerDim = TestUtil.nextInt(random(), 1, 4);
+    final byte[] pointValue = new byte[numBytesPerDim];
+    random().nextBytes(pointValue);
+
+    MutablePointValues reader = new MutablePointValues() {
+
+      @Override
+      public void intersect(IntersectVisitor visitor) throws IOException {
+        for(int i=0;i<numPointsAdded;i++) {
+          visitor.visit(0, pointValue);
+        }
+      }
+
+      @Override
+      public long estimatePointCount(IntersectVisitor visitor) {
+        throw new UnsupportedOperationException();
+      }
+
+      @Override
+      public byte[] getMinPackedValue() {
+        throw new UnsupportedOperationException();
+      }
+
+      @Override
+      public byte[] getMaxPackedValue() {
+        throw new UnsupportedOperationException();
+      }
+
+      @Override
+      public int getNumDataDimensions() {
+        throw new UnsupportedOperationException();
+      }
+
+      @Override
+      public int getNumIndexDimensions() {
+        throw new UnsupportedOperationException();
+      }
+
+      @Override
+      public int getBytesPerDimension() {
+        throw new UnsupportedOperationException();
+      }
+
+      @Override
+      public long size() {
+        return numPointsAdded;
+      }
+
+      @Override
+      public int getDocCount() {
+        return numPointsAdded;
+      }
+
+      @Override
+      public void swap(int i, int j) {
+        // do nothing
+      }
+
+      @Override
+      public int getDocID(int i) {
+        return 0;
+      }
+
+      @Override
+      public void getValue(int i, BytesRef packedValue) {
+        packedValue.bytes = pointValue;
+      }
+
+      @Override
+      public byte getByteAt(int i, int k) {
+        throw new UnsupportedOperationException();
+      }
+    };
+
+    BKDWriter w = new BKDWriter(numValues, dir, "_temp", 1, 1, numBytesPerDim, BKDWriter.DEFAULT_MAX_POINTS_IN_LEAF_NODE,
+        BKDWriter.DEFAULT_MAX_MB_SORT_IN_HEAP, numValues);
+    expectThrows(IllegalStateException.class, () -> {
+      try (IndexOutput out = dir.createOutput("bkd", IOContext.DEFAULT)) {
+        w.writeField(out, "test_field_name", reader);
+      } finally {
+        w.close();
+        dir.close();
+      }
+    });
+  }
 }