GH-1217: Bounds check before alloc in var-width view vectors (#1291)

Also, document that when enable_unsafe_memory_access is enabled, all
bets are off.

Reported by n0mi1k.

Author: lidavidm

memory/memory-core/src/main/java/org/apache/arrow/memory/ArrowBuf.java

@@ -731,6 +731,24 @@ public void getBytes(long index, byte[] dst, int dstIndex, int length) {
     }
   }
 
+  /**
+   * Copy data from this ArrowBuf into a newly allocated array.
+   *
+   * <p>This method is more resilient to invalid data inadvertently causing large allocations, as
+   * the byte[] will not be allocated until we check the length.
+   *
+   * @param index index (0 based relative to the portion of memory this ArrowBuf has access to)
+   * @param length length of data to copy from this ArrowBuf
+   */
+  public byte[] getBytesAsArray(long index, int length) {
+    checkIndex(index, length);
+    byte[] dst = new byte[length];
+    if (length != 0) {
+      MemoryUtil.copyFromMemory(addr(index), dst, 0, length);
+    }
+    return dst;
+  }
+
   /**
    * Copy data from a given byte array into this ArrowBuf starting at a given index.
    *
@@ -1008,8 +1026,7 @@ public int setBytes(long index, InputStream in, int length) throws IOException {
   /**
    * Copy a certain length of bytes from this ArrowBuf at a given index into the given OutputStream.
    *
-   * @param index index index (0 based relative to the portion of memory this ArrowBuf has access
-   *     to)
+   * @param index index (0 based relative to the portion of memory this ArrowBuf has access to)
    * @param out dst stream to copy data into
    * @param length length of data to copy
    * @throws IOException on failing to write to stream

memory/memory-core/src/main/java/org/apache/arrow/memory/BoundsChecking.java

@@ -24,6 +24,10 @@
  * "arrow.enable_unsafe_memory_access" or "drill.enable_unsafe_memory_access". The latter is
  * deprecated. The environmental variable is named "ARROW_ENABLE_UNSAFE_MEMORY_ACCESS". When both
  * the system property and the environmental variable are set, the system property takes precedence.
+ *
+ * <p>WARNING: disabling bounds checking means that out-of-bounds memory access is possible! This
+ * can lead to security vulnerabilities. You should not read or write untrusted data when bounds
+ * checking is disabled.
  */
 public class BoundsChecking {
 

vector/src/main/java/org/apache/arrow/vector/BaseVariableWidthViewVector.java

@@ -912,6 +912,13 @@ private void splitAndTransferViewBufferAndDataBuffer(
             viewBuffer.getInt(
                 ((long) i * ELEMENT_SIZE) + LENGTH_WIDTH + PREFIX_WIDTH + BUF_INDEX_WIDTH);
         final ArrowBuf dataBuf = dataBuffers.get(readBufIndex);
+        if (readBufOffset < 0
+            || ((long) readBufOffset + (long) stringLength) > dataBuf.capacity()) {
+          throw new IndexOutOfBoundsException(
+              String.format(
+                  "index: %d, length: %d (expected: range(0, %d))",
+                  readBufOffset, stringLength, dataBuf.capacity()));
+        }
 
         // allocate data buffer
         ArrowBuf currentDataBuf = target.allocateOrGetLastDataBuffer(stringLength);
@@ -1432,7 +1439,7 @@ public void copyFrom(int fromIndex, int thisIndex, ValueVector from) {
       BitVectorHelper.unsetBit(validityBuffer, thisIndex);
     } else {
       final int viewLength = from.getDataBuffer().getInt((long) fromIndex * ELEMENT_SIZE);
-      copyFromNotNull(fromIndex, thisIndex, from, viewLength);
+      copyFromNotNull(from, fromIndex, thisIndex, viewLength);
     }
     lastSet = thisIndex;
   }
@@ -1454,39 +1461,35 @@ public void copyFromSafe(int fromIndex, int thisIndex, ValueVector from) {
     } else {
       final int viewLength = from.getDataBuffer().getInt((long) fromIndex * ELEMENT_SIZE);
       handleSafe(thisIndex, viewLength);
-      copyFromNotNull(fromIndex, thisIndex, from, viewLength);
+      copyFromNotNull(from, fromIndex, thisIndex, viewLength);
     }
     lastSet = thisIndex;
   }
 
-  private void copyFromNotNull(int fromIndex, int thisIndex, ValueVector from, int viewLength) {
+  private void copyFromNotNull(ValueVector from, int fromIndex, int thisIndex, int viewLength) {
     BitVectorHelper.setBit(validityBuffer, thisIndex);
     final int start = thisIndex * ELEMENT_SIZE;
     final int copyStart = fromIndex * ELEMENT_SIZE;
     if (viewLength > INLINE_SIZE) {
-      final int bufIndex =
-          from.getDataBuffer()
-              .getInt(((long) fromIndex * ELEMENT_SIZE) + LENGTH_WIDTH + PREFIX_WIDTH);
-      final int dataOffset =
-          from.getDataBuffer()
-              .getInt(
-                  ((long) fromIndex * ELEMENT_SIZE)
-                      + LENGTH_WIDTH
-                      + PREFIX_WIDTH
-                      + BUF_INDEX_WIDTH);
-      final ArrowBuf dataBuf = ((BaseVariableWidthViewVector) from).dataBuffers.get(bufIndex);
-      final ArrowBuf thisDataBuf = allocateOrGetLastDataBuffer(viewLength);
-
-      viewBuffer.setBytes(start, from.getDataBuffer(), copyStart, LENGTH_WIDTH + PREFIX_WIDTH);
-      int writePosition = start + LENGTH_WIDTH + PREFIX_WIDTH;
-      // set buf id
-      viewBuffer.setInt(writePosition, dataBuffers.size() - 1);
-      writePosition += BUF_INDEX_WIDTH;
-      // set offset
-      viewBuffer.setInt(writePosition, (int) thisDataBuf.writerIndex());
-
-      thisDataBuf.setBytes(thisDataBuf.writerIndex(), dataBuf, dataOffset, viewLength);
-      thisDataBuf.writerIndex(thisDataBuf.writerIndex() + viewLength);
+      BaseVariableWidthViewVector fromVector = (BaseVariableWidthViewVector) from;
+      fromVector.getData(
+          fromIndex,
+          (dataBuf, dataOffset, dataLength) -> {
+            assert dataLength == viewLength;
+            viewBuffer.setBytes(
+                start, fromVector.getDataBuffer(), copyStart, LENGTH_WIDTH + PREFIX_WIDTH);
+            //noinspection resource
+            final ArrowBuf thisDataBuf = allocateOrGetLastDataBuffer(viewLength);
+            int writePosition = start + LENGTH_WIDTH + PREFIX_WIDTH;
+            // set buf id
+            viewBuffer.setInt(writePosition, dataBuffers.size() - 1);
+            writePosition += BUF_INDEX_WIDTH;
+            // set offset
+            viewBuffer.setInt(writePosition, (int) thisDataBuf.writerIndex());
+            thisDataBuf.setBytes(thisDataBuf.writerIndex(), dataBuf, dataOffset, viewLength);
+            thisDataBuf.writerIndex(thisDataBuf.writerIndex() + viewLength);
+            return null;
+          });
     } else {
       from.getDataBuffer().getBytes(copyStart, viewBuffer, start, ELEMENT_SIZE);
     }
@@ -1502,16 +1505,12 @@ public ArrowBufPointer getDataPointer(int index, ArrowBufPointer reuse) {
     if (isNull(index)) {
       reuse.set(null, 0, 0);
     } else {
-      int length = getValueLength(index);
-      if (length < INLINE_SIZE) {
-        int start = index * ELEMENT_SIZE + LENGTH_WIDTH;
-        reuse.set(viewBuffer, start, length);
-      } else {
-        final int bufIndex =
-            viewBuffer.getInt(((long) index * ELEMENT_SIZE) + LENGTH_WIDTH + PREFIX_WIDTH);
-        ArrowBuf dataBuf = dataBuffers.get(bufIndex);
-        reuse.set(dataBuf, 0, length);
-      }
+      getData(
+          index,
+          (buf, offset, length) -> {
+            reuse.set(buf, offset, length);
+            return null;
+          });
     }
     return reuse;
   }
@@ -1526,19 +1525,45 @@ public int hashCode(int index, ArrowBufHasher hasher) {
     if (isNull(index)) {
       return ArrowBufPointer.NULL_HASH_CODE;
     }
-    final int length = getValueLength(index);
-    if (length < INLINE_SIZE) {
-      int start = index * ELEMENT_SIZE + LENGTH_WIDTH;
-      return ByteFunctionHelpers.hash(hasher, this.getDataBuffer(), start, start + length);
-    } else {
-      final int bufIndex =
+    return getData(
+        index,
+        (buf, offset, length) -> ByteFunctionHelpers.hash(hasher, buf, offset, offset + length));
+  }
+
+  @FunctionalInterface
+  protected interface ViewElementConsumer<T> {
+    T consume(ArrowBuf buf, int offset, int length);
+  }
+
+  /** Helper to get a single view value with sanity checking. */
+  protected <T> T getData(int index, ViewElementConsumer<T> consumer) {
+    final int dataLength = getValueLength(index);
+    final ArrowBuf dataBuffer;
+    final int dataOffset;
+    if (dataLength > INLINE_SIZE) {
+      final int bufferIndex =
           viewBuffer.getInt(((long) index * ELEMENT_SIZE) + LENGTH_WIDTH + PREFIX_WIDTH);
-      final int dataOffset =
+      dataOffset =
           viewBuffer.getInt(
               ((long) index * ELEMENT_SIZE) + LENGTH_WIDTH + PREFIX_WIDTH + BUF_INDEX_WIDTH);
-      ArrowBuf dataBuf = dataBuffers.get(bufIndex);
-      return ByteFunctionHelpers.hash(hasher, dataBuf, dataOffset, dataOffset + length);
+      dataBuffer = dataBuffers.get(bufferIndex);
+    } else {
+      dataBuffer = viewBuffer;
+      dataOffset = index * ELEMENT_SIZE + BUF_INDEX_WIDTH;
+    }
+    if (dataOffset < 0
+        || dataLength < 0
+        || ((long) dataOffset + (long) dataLength) > dataBuffer.capacity()) {
+      // In this case we don't check BOUNDS_CHECKING_ENABLED
+      // Likely this check is redundant, but we are trying to check eagerly before downstream code
+      // potentially
+      // tries to allocate based on the given dataLength
+      throw new IndexOutOfBoundsException(
+          String.format(
+              "index: %d, length: %d (expected: range(0, %d))",
+              dataOffset, dataLength, dataBuffer.capacity()));
     }
+    return consumer.consume(dataBuffer, dataOffset, dataLength);
   }
 
   /**
@@ -1555,42 +1580,16 @@ public int hashCode(int index, ArrowBufHasher hasher) {
    * @return byte array containing the data of the element
    */
   protected byte[] getData(int index) {
-    final int dataLength = getValueLength(index);
-    byte[] result = new byte[dataLength];
-    if (dataLength > INLINE_SIZE) {
-      // data is in the data buffer
-      // get buffer index
-      final int bufferIndex =
-          viewBuffer.getInt(((long) index * ELEMENT_SIZE) + LENGTH_WIDTH + PREFIX_WIDTH);
-      // get data offset
-      final int dataOffset =
-          viewBuffer.getInt(
-              ((long) index * ELEMENT_SIZE) + LENGTH_WIDTH + PREFIX_WIDTH + BUF_INDEX_WIDTH);
-      dataBuffers.get(bufferIndex).getBytes(dataOffset, result, 0, dataLength);
-    } else {
-      // data is in the view buffer
-      viewBuffer.getBytes((long) index * ELEMENT_SIZE + BUF_INDEX_WIDTH, result, 0, dataLength);
-    }
-    return result;
+    return getData(index, ArrowBuf::getBytesAsArray);
   }
 
   protected void getData(int index, ReusableBuffer<?> buffer) {
-    final int dataLength = getValueLength(index);
-    if (dataLength > INLINE_SIZE) {
-      // data is in the data buffer
-      // get buffer index
-      final int bufferIndex =
-          viewBuffer.getInt(((long) index * ELEMENT_SIZE) + LENGTH_WIDTH + PREFIX_WIDTH);
-      // get data offset
-      final int dataOffset =
-          viewBuffer.getInt(
-              ((long) index * ELEMENT_SIZE) + LENGTH_WIDTH + PREFIX_WIDTH + BUF_INDEX_WIDTH);
-      ArrowBuf dataBuf = dataBuffers.get(bufferIndex);
-      buffer.set(dataBuf, dataOffset, dataLength);
-    } else {
-      // data is in the value buffer
-      buffer.set(viewBuffer, ((long) index * ELEMENT_SIZE) + BUF_INDEX_WIDTH, dataLength);
-    }
+    getData(
+        index,
+        (buf, offset, length) -> {
+          buffer.set(buf, offset, length);
+          return null;
+        });
   }
 
   @Override

vector/src/test/java/org/apache/arrow/vector/TestVariableWidthViewVector.java

@@ -2921,4 +2921,21 @@ public void testValidate() {
       assertTrue(e.getMessage().contains("Not enough capacity for data buffer"));
     }
   }
+
+  @Test
+  public void testValidateInvalidOffsets() {
+    try (final ViewVarCharVector vector = new ViewVarCharVector("v", allocator)) {
+      vector.allocateNew(16, 1);
+      vector.allocateOrGetLastDataBuffer(8);
+      var offsets = vector.getDataBuffer();
+      offsets.setInt(0, 64);
+      offsets.setInt(4, 0);
+      offsets.setInt(8, 0);
+      offsets.setInt(12, 1024);
+      vector.setValueCount(1);
+      vector.setIndexDefined(0);
+      var e = assertThrows(IndexOutOfBoundsException.class, vector::validateFull);
+      assertTrue(e.getMessage().contains("index: 1024"));
+    }
+  }
 }