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 c3eafaacf3..6bc113e4ae 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 @@ -897,6 +897,12 @@ public void setBytes(long index, ByteBuffer src) { public void setBytes(long index, ByteBuffer src, int srcIndex, int length) { // bound check for this ArrowBuf where the data will be copied into checkIndex(index, length); + // bound check for src ByteBuffer + if (isOutOfBounds(srcIndex, length, src.capacity())) { + throw new IndexOutOfBoundsException( + String.format( + "index: %d, length: %d (expected: range(0, %d))", srcIndex, length, src.capacity())); + } if (src.isDirect()) { // copy length bytes of data from src ByteBuffer starting at address // srcAddress into this ArrowBuf at address dstAddress diff --git a/memory/memory-core/src/test/java/org/apache/arrow/memory/TestArrowBuf.java b/memory/memory-core/src/test/java/org/apache/arrow/memory/TestArrowBuf.java index 60e118d557..fdb77882f9 100644 --- a/memory/memory-core/src/test/java/org/apache/arrow/memory/TestArrowBuf.java +++ b/memory/memory-core/src/test/java/org/apache/arrow/memory/TestArrowBuf.java @@ -112,6 +112,17 @@ public void testSetBytesUnsliced() { } } + @Test + public void testSetBytesDirectOutOfBoundsSource_RaisesIndexOutOfBoundsException() { + ByteBuffer data = ByteBuffer.allocateDirect(8); + try (BufferAllocator allocator = new RootAllocator(128); + ArrowBuf buf = allocator.buffer(64)) { + assertThrows(IndexOutOfBoundsException.class, () -> buf.setBytes(0, data, 0, 64)); + assertThrows(IndexOutOfBoundsException.class, () -> buf.setBytes(0, data, 4, 8)); + assertThrows(IndexOutOfBoundsException.class, () -> buf.setBytes(0, data, -1, 4)); + } + } + /** ARROW-9221: guard against big-endian byte buffers. */ @Test public void testSetBytesBigEndian() {