diff --git a/vector/src/main/java/org/apache/arrow/vector/BaseFixedWidthVector.java b/vector/src/main/java/org/apache/arrow/vector/BaseFixedWidthVector.java index df1ac74f9b..478e1ae5a3 100644 --- a/vector/src/main/java/org/apache/arrow/vector/BaseFixedWidthVector.java +++ b/vector/src/main/java/org/apache/arrow/vector/BaseFixedWidthVector.java @@ -26,7 +26,6 @@ import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.util.ArrowBufPointer; import org.apache.arrow.memory.util.ByteFunctionHelpers; -import org.apache.arrow.memory.util.MemoryUtil; import org.apache.arrow.memory.util.hash.ArrowBufHasher; import org.apache.arrow.util.Preconditions; import org.apache.arrow.vector.compare.VectorVisitor; @@ -833,10 +832,12 @@ public void copyFrom(int fromIndex, int thisIndex, ValueVector from) { BitVectorHelper.unsetBit(this.getValidityBuffer(), thisIndex); } else { BitVectorHelper.setBit(this.getValidityBuffer(), thisIndex); - MemoryUtil.copyMemory( - from.getDataBuffer().memoryAddress() + (long) fromIndex * typeWidth, - this.getDataBuffer().memoryAddress() + (long) thisIndex * typeWidth, - typeWidth); + from.getDataBuffer() + .getBytes( + (long) fromIndex * typeWidth, + this.getDataBuffer(), + (long) thisIndex * typeWidth, + typeWidth); } } diff --git a/vector/src/main/java/org/apache/arrow/vector/BitVectorHelper.java b/vector/src/main/java/org/apache/arrow/vector/BitVectorHelper.java index bc2c3da98f..7fedd6933a 100644 --- a/vector/src/main/java/org/apache/arrow/vector/BitVectorHelper.java +++ b/vector/src/main/java/org/apache/arrow/vector/BitVectorHelper.java @@ -381,6 +381,8 @@ public static void concatBits( int numBytesOut = DataSizeRoundingUtil.divideBy8Ceil(numBits1 + numBits2); if (BoundsChecking.BOUNDS_CHECKING_ENABLED) { + input1.checkBytes(0, numBytes1); + input2.checkBytes(0, numBytes2); output.checkBytes(0, numBytesOut); } diff --git a/vector/src/main/java/org/apache/arrow/vector/util/VectorAppender.java b/vector/src/main/java/org/apache/arrow/vector/util/VectorAppender.java index f14158c9f9..9fbd7f7661 100644 --- a/vector/src/main/java/org/apache/arrow/vector/util/VectorAppender.java +++ b/vector/src/main/java/org/apache/arrow/vector/util/VectorAppender.java @@ -74,6 +74,18 @@ public ValueVector visit(BaseFixedWidthVector deltaVector, Void value) { return targetVector; // optimization, nothing to append, return } + // validate the delta before the target is touched, so a malformed delta leaves it unchanged + int deltaCount = deltaVector.getValueCount(); + checkCapacity( + deltaVector.getValidityBuffer(), + BitVectorHelper.getValidityBufferSizeFromCount(deltaCount), + "validity"); + long dataSize = + deltaVector instanceof BitVector + ? BitVectorHelper.getValidityBufferSizeFromCount(deltaCount) + : (long) deltaVector.getTypeWidth() * deltaCount; + checkCapacity(deltaVector.getDataBuffer(), dataSize, "data"); + int newValueCount = targetVector.getValueCount() + deltaVector.getValueCount(); // make sure there is enough capacity @@ -337,6 +349,16 @@ public ValueVector visit(BaseVariableWidthViewVector deltaVector, Void value) { return targetVector; // nothing to append, return } + // validate the delta before the target is touched, so a malformed delta leaves it unchanged + checkCapacity( + deltaVector.getValidityBuffer(), + BitVectorHelper.getValidityBufferSizeFromCount(deltaVector.getValueCount()), + "validity"); + checkCapacity( + deltaVector.getDataBuffer(), + (long) BaseVariableWidthViewVector.ELEMENT_SIZE * deltaVector.getValueCount(), + "view"); + int oldTargetValueCount = targetVector.getValueCount(); int newValueCount = oldTargetValueCount + deltaVector.getValueCount(); @@ -664,6 +686,11 @@ public ValueVector visit(UnionVector deltaVector, Void value) { return targetVector; // optimization, nothing to append, return } + checkCapacity( + deltaVector.getTypeBuffer(), + (long) deltaVector.getValueCount() * UnionVector.TYPE_WIDTH, + "type"); + UnionVector targetUnionVector = (UnionVector) targetVector; int newValueCount = targetVector.getValueCount() + deltaVector.getValueCount(); @@ -728,6 +755,11 @@ public ValueVector visit(DenseUnionVector deltaVector, Void value) { return targetVector; // optimization, nothing to append, return } + checkCapacity( + deltaVector.getTypeBuffer(), + (long) deltaVector.getValueCount() * DenseUnionVector.TYPE_WIDTH, + "type"); + DenseUnionVector targetDenseUnionVector = (DenseUnionVector) targetVector; int newValueCount = targetVector.getValueCount() + deltaVector.getValueCount(); diff --git a/vector/src/test/java/org/apache/arrow/vector/TestCopyFrom.java b/vector/src/test/java/org/apache/arrow/vector/TestCopyFrom.java index bd0b44a144..9a4513579b 100644 --- a/vector/src/test/java/org/apache/arrow/vector/TestCopyFrom.java +++ b/vector/src/test/java/org/apache/arrow/vector/TestCopyFrom.java @@ -20,15 +20,19 @@ import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; import java.math.BigDecimal; import java.nio.charset.StandardCharsets; import java.time.Duration; import java.time.Period; +import java.util.Arrays; import java.util.Objects; +import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; +import org.apache.arrow.vector.ipc.message.ArrowFieldNode; import org.apache.arrow.vector.types.Types.MinorType; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; @@ -1098,4 +1102,27 @@ public void testCopySafeArrow7837() { assertEquals("1234567890", Objects.requireNonNull(vc2.getObject(5)).toString()); } } + + @Test + public void testCopyFromFixedWidthSourceBufferTooShort() { + try (IntVector source = new IntVector("source", allocator); + IntVector target = new IntVector("target", allocator); + ArrowBuf body = allocator.buffer(64)) { + body.setZero(0, body.capacity()); + // all 4 values marked valid + body.setByte(0, 0x0F); + body.setInt(8, 1); + body.setInt(20, 0xBAD); + // Field node claims 4 values, but the data buffer only holds one, as a malformed IPC + // record batch can. + ArrowBuf validity = body.slice(0, 8); + ArrowBuf data = body.slice(8, IntVector.TYPE_WIDTH); + source.loadFieldBuffers(new ArrowFieldNode(4, 0), Arrays.asList(validity, data)); + + target.allocateNew(4); + target.copyFromSafe(0, 0, source); + assertEquals(1, target.get(0)); + assertThrows(IndexOutOfBoundsException.class, () -> target.copyFromSafe(3, 1, source)); + } + } } diff --git a/vector/src/test/java/org/apache/arrow/vector/util/TestVectorAppender.java b/vector/src/test/java/org/apache/arrow/vector/util/TestVectorAppender.java index 57dd852e31..5b0213ba0c 100644 --- a/vector/src/test/java/org/apache/arrow/vector/util/TestVectorAppender.java +++ b/vector/src/test/java/org/apache/arrow/vector/util/TestVectorAppender.java @@ -109,6 +109,65 @@ public void testAppendFixedWidthVector() { } } + @Test + public void testAppendFixedWidthVectorSourceBufferTooShort() { + try (IntVector target = new IntVector("", allocator); + IntVector delta = new IntVector("", allocator); + ArrowBuf body = allocator.buffer(64)) { + body.setZero(0, body.capacity()); + // all 4 values marked valid + body.setByte(0, 0x0F); + body.setInt(8, 1); + // Field node claims 4 values, but the data buffer only holds one, as a malformed IPC + // record batch can. + ArrowBuf validity = body.slice(0, 8); + ArrowBuf data = body.slice(8, IntVector.TYPE_WIDTH); + delta.loadFieldBuffers(new ArrowFieldNode(4, 0), Arrays.asList(validity, data)); + + target.allocateNew(4); + ValueVectorDataPopulator.setVector(target, 0, 1, 2, 3); + int capacity = target.getValueCapacity(); + + VectorAppender appender = new VectorAppender(target); + assertThrows(IndexOutOfBoundsException.class, () -> delta.accept(appender, null)); + + // the delta is rejected before the target is touched + assertEquals(4, target.getValueCount()); + assertEquals(capacity, target.getValueCapacity()); + for (int i = 0; i < 4; i++) { + assertEquals(i, target.get(i)); + } + } + } + + @Test + public void testAppendValidityBufferTooShort() { + try (IntVector target = new IntVector("", allocator); + IntVector delta = new IntVector("", allocator); + ArrowBuf body = allocator.buffer(128)) { + body.setZero(0, body.capacity()); + body.setByte(0, 0x0F); + // Field node claims 16 values, but the validity buffer only covers 8. + ArrowBuf validity = body.slice(0, 1); + ArrowBuf data = body.slice(64, 16 * IntVector.TYPE_WIDTH); + delta.loadFieldBuffers(new ArrowFieldNode(16, 4), Arrays.asList(validity, data)); + + target.allocateNew(8); + ValueVectorDataPopulator.setVector(target, 0, 1, 2, 3, 4, 5, 6, 7); + int capacity = target.getValueCapacity(); + + VectorAppender appender = new VectorAppender(target); + assertThrows(IndexOutOfBoundsException.class, () -> delta.accept(appender, null)); + + // the delta is rejected before the target is touched + assertEquals(8, target.getValueCount()); + assertEquals(capacity, target.getValueCapacity()); + for (int i = 0; i < 8; i++) { + assertEquals(i, target.get(i)); + } + } + } + @Test public void testAppendBitVector() { final int length1 = 10;