diff --git a/memory/memory-core/src/main/java/org/apache/arrow/memory/ArrowBuf.java b/memory/memory-core/src/main/java/org/apache/arrow/memory/ArrowBuf.java index 9712be34d7..c3eafaacf3 100644 --- a/memory/memory-core/src/main/java/org/apache/arrow/memory/ArrowBuf.java +++ b/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. + * + *

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 diff --git a/memory/memory-core/src/main/java/org/apache/arrow/memory/BoundsChecking.java b/memory/memory-core/src/main/java/org/apache/arrow/memory/BoundsChecking.java index 50be9ad1fb..9b99edafa4 100644 --- a/memory/memory-core/src/main/java/org/apache/arrow/memory/BoundsChecking.java +++ b/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. + * + *

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 { diff --git a/vector/src/main/java/org/apache/arrow/vector/BaseVariableWidthViewVector.java b/vector/src/main/java/org/apache/arrow/vector/BaseVariableWidthViewVector.java index 8d2a2d7401..d9c48e6a85 100644 --- a/vector/src/main/java/org/apache/arrow/vector/BaseVariableWidthViewVector.java +++ b/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 consume(ArrowBuf buf, int offset, int length); + } + + /** Helper to get a single view value with sanity checking. */ + protected T getData(int index, ViewElementConsumer 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 diff --git a/vector/src/test/java/org/apache/arrow/vector/TestVariableWidthViewVector.java b/vector/src/test/java/org/apache/arrow/vector/TestVariableWidthViewVector.java index c4a1ae9b61..2d33f57e99 100644 --- a/vector/src/test/java/org/apache/arrow/vector/TestVariableWidthViewVector.java +++ b/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, Integer.MAX_VALUE); + 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")); + } + } }