[SPARK-9364] Fix array out of bounds and use-after-free bugs in UnsafeExternalSorter

core/src/main/java/org/apache/spark/util/collection/unsafe/sort/UnsafeExternalSorter.java

@@ -150,6 +150,11 @@ private long getMemoryUsage() {
     return sorter.getMemoryUsage() + (allocatedPages.size() * (long) PAGE_SIZE);
   }
 
+  @VisibleForTesting
+  public int getNumberOfAllocatedPages() {
+    return allocatedPages.size();
+  }
+
   public long freeMemory() {
     long memoryFreed = 0;
     for (MemoryBlock block : allocatedPages) {
@@ -257,7 +262,7 @@ public void insertRecord(
       currentPagePosition,
       lengthInBytes);
     currentPagePosition += lengthInBytes;
-
+    freeSpaceInCurrentPage -= totalSpaceRequired;
     sorter.insertRecord(recordAddress, prefix);
   }
 

core/src/test/java/org/apache/spark/util/collection/unsafe/sort/UnsafeExternalSorterSuite.java

@@ -199,4 +199,23 @@ public void testSortingEmptyArrays() throws Exception {
     }
   }
 
+  @Test
+  public void testFillingPage() throws Exception {
+    final UnsafeExternalSorter sorter = new UnsafeExternalSorter(
+      memoryManager,
+      shuffleMemoryManager,
+      blockManager,
+      taskContext,
+      recordComparator,
+      prefixComparator,
+      1024,
+      new SparkConf());
+
+    byte[] record = new byte[16];
+    while (sorter.getNumberOfAllocatedPages() < 2) {
+      sorter.insertRecord(record, PlatformDependent.BYTE_ARRAY_OFFSET, record.length, 0);
+    }
+    sorter.freeMemory();
+  }
+
 }

sql/catalyst/src/main/java/org/apache/spark/sql/execution/UnsafeExternalRowSorter.java

@@ -124,7 +124,7 @@ Iterator<InternalRow> sort() throws IOException {
       return new AbstractScalaRowIterator() {
 
         private final int numFields = schema.length();
-        private final UnsafeRow row = new UnsafeRow();
+        private UnsafeRow row = new UnsafeRow();
 
         @Override
         public boolean hasNext() {
@@ -141,10 +141,13 @@ public InternalRow next() {
               numFields,
               sortedIterator.getRecordLength());
             if (!hasNext()) {
-              row.copy(); // so that we don't have dangling pointers to freed page
+              UnsafeRow copy = row.copy(); // so that we don't have dangling pointers to freed page
+              row = null; // so that we don't keep references to the base object
               cleanupResources();
+              return copy;
+            } else {
+              return row;
             }
-            return row;
           } catch (IOException e) {
             cleanupResources();
             // Scala iterators don't declare any checked exceptions, so we need to use this hack