diff --git a/c/src/main/java/org/apache/arrow/c/ArrayImporter.java b/c/src/main/java/org/apache/arrow/c/ArrayImporter.java index e92e832991..c7c410010b 100644 --- a/c/src/main/java/org/apache/arrow/c/ArrayImporter.java +++ b/c/src/main/java/org/apache/arrow/c/ArrayImporter.java @@ -17,7 +17,6 @@ package org.apache.arrow.c; import static org.apache.arrow.c.NativeUtil.NULL; -import static org.apache.arrow.memory.util.LargeMemoryUtil.checkedCastToInt; import static org.apache.arrow.util.Preconditions.checkNotNull; import static org.apache.arrow.util.Preconditions.checkState; @@ -27,6 +26,7 @@ import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.util.Preconditions; import org.apache.arrow.vector.FieldVector; +import org.apache.arrow.vector.TypeLayout; import org.apache.arrow.vector.dictionary.Dictionary; import org.apache.arrow.vector.dictionary.DictionaryProvider; import org.apache.arrow.vector.ipc.message.ArrowFieldNode; @@ -93,16 +93,33 @@ private void doImport(ArrowArray.Snapshot snapshot) { "ArrowArray struct has non-zero offset (%s), which is not supported", snapshot.offset); - // First import children (required for reconstituting parent array data) - long[] children = - NativeUtil.toJavaArray(snapshot.children, checkedCastToInt(snapshot.n_children)); + // Check the length and null count of this array and of its dictionary before importing + // anything, so that invalid values are rejected before the child vectors or the dictionary are + // modified. Errors found later, such as missing buffers or an invalid child, can still leave + // the child vectors or the dictionary partially imported. + final ArrowFieldNode fieldNode = toFieldNode(vector, snapshot); + Dictionary dictionary = null; + if (snapshot.dictionary != NULL) { + DictionaryEncoding encoding = vector.getField().getDictionary(); + checkNotNull(encoding, "Missing encoding on import of ArrowArray with dictionary"); + + dictionary = dictionaryProvider.lookup(encoding.getId()); + checkNotNull(dictionary, "Dictionary lookup failed on import of ArrowArray with dictionary"); + + toFieldNode(dictionary.getVector(), ArrowArray.wrap(snapshot.dictionary).snapshot()); + } + + // First import children (required for reconstituting parent array data). Check the number of + // children before reading the child pointers, so that an invalid count is not used to size + // the read. + List childVectors = vector.getChildrenFromFields(); + checkState( + snapshot.n_children == 0 || snapshot.n_children == childVectors.size(), + "ArrowArray struct has %s children (expected %s)", + snapshot.n_children, + childVectors.size()); + long[] children = NativeUtil.toJavaArray(snapshot.children, (int) snapshot.n_children); if (children != null && children.length > 0) { - List childVectors = vector.getChildrenFromFields(); - checkState( - children.length == childVectors.size(), - "ArrowArray struct has %s children (expected %s)", - children.length, - childVectors.size()); for (int i = 0; i < children.length; i++) { checkState(children[i] != NULL, "ArrowArray struct has NULL child at position %s", i); ArrayImporter childImporter = @@ -111,14 +128,44 @@ private void doImport(ArrowArray.Snapshot snapshot) { } } - // Handle import of a dictionary encoded vector - if (snapshot.dictionary != NULL) { - DictionaryEncoding encoding = vector.getField().getDictionary(); - checkNotNull(encoding, "Missing encoding on import of ArrowArray with dictionary"); - - Dictionary dictionary = dictionaryProvider.lookup(encoding.getId()); - checkNotNull(dictionary, "Dictionary lookup failed on import of ArrowArray with dictionary"); + // Import main data + try { + // Check the number of buffers before reading the buffer pointers. Types with variadic + // buffers, such as views, have no fixed upper bound. + TypeLayout layout = TypeLayout.getTypeLayout(vector.getField().getType()); + checkState( + snapshot.n_buffers >= 0, + "ArrowArray struct has a negative number of buffers (%s)", + snapshot.n_buffers); + checkState( + !layout.isFixedBufferCount() || snapshot.n_buffers <= layout.getBufferTypes().size(), + "ArrowArray struct has %s buffers (expected at most %s)", + snapshot.n_buffers, + layout.getBufferTypes().size()); + long[] bufferPointers = + NativeUtil.toJavaArray(snapshot.buffers, Math.toIntExact(snapshot.n_buffers)); + + try (final BufferImportTypeVisitor visitor = + new BufferImportTypeVisitor(allocator, underlyingAllocation, fieldNode, bufferPointers)) { + final List buffers; + if (bufferPointers == null || bufferPointers.length == 0) { + buffers = Collections.emptyList(); + } else { + buffers = vector.getField().getType().accept(visitor); + } + vector.loadFieldBuffers(fieldNode, buffers); + } + } catch (Exception e) { + throw new IllegalArgumentException( + "Could not load buffers for field " + + vector.getField() + + ". error message: " + + e.getMessage(), + e); + } + // Import the dictionary last, so that it is not cleared if this array's own data is invalid + if (dictionary != null) { // reset the dictionary vector to the initial state dictionary.getVector().clear(); @@ -126,24 +173,14 @@ private void doImport(ArrowArray.Snapshot snapshot) { new ArrayImporter(allocator, dictionary.getVector(), dictionaryProvider); dictionaryImporter.importChild(this, ArrowArray.wrap(snapshot.dictionary)); } + } - // Import main data - ArrowFieldNode fieldNode = new ArrowFieldNode(snapshot.length, snapshot.null_count); - long[] bufferPointers = - NativeUtil.toJavaArray(snapshot.buffers, checkedCastToInt(snapshot.n_buffers)); - - try (final BufferImportTypeVisitor visitor = - new BufferImportTypeVisitor(allocator, underlyingAllocation, fieldNode, bufferPointers)) { - final List buffers; - if (bufferPointers == null || bufferPointers.length == 0) { - buffers = Collections.emptyList(); - } else { - buffers = vector.getField().getType().accept(visitor); - } - vector.loadFieldBuffers(fieldNode, buffers); - } catch (Exception e) { + private static ArrowFieldNode toFieldNode(FieldVector vector, ArrowArray.Snapshot snapshot) { + try { + return new ArrowFieldNode(snapshot.length, snapshot.null_count); + } catch (IllegalArgumentException | ArithmeticException e) { throw new IllegalArgumentException( - "Could not load buffers for field " + "Invalid length or null count for field " + vector.getField() + ". error message: " + e.getMessage(), diff --git a/c/src/main/java/org/apache/arrow/c/BufferImportTypeVisitor.java b/c/src/main/java/org/apache/arrow/c/BufferImportTypeVisitor.java index 10f690fc87..5be28e3ce4 100644 --- a/c/src/main/java/org/apache/arrow/c/BufferImportTypeVisitor.java +++ b/c/src/main/java/org/apache/arrow/c/BufferImportTypeVisitor.java @@ -83,6 +83,8 @@ public void close() throws Exception { @VisibleForTesting ArrowBuf importBuffer(ArrowType type, int index, long capacity) { + checkState( + capacity >= 0, "Buffer %s for type %s has negative capacity %s", index, type, capacity); checkState( buffers.length > index, "Expected at least %s buffers for type %s, but found %s", @@ -118,7 +120,7 @@ private ArrowBuf importFixedBytes(ArrowType type, int index, long bytesPerSlot) } private ArrowBuf importOffsets(ArrowType type, long bytesPerSlot) { - final long capacity = bytesPerSlot * (fieldNode.getLength() + 1); + final long capacity = bytesPerSlot * ((long) fieldNode.getLength() + 1); return importBuffer(type, 1, capacity); } diff --git a/c/src/main/java/org/apache/arrow/c/Data.java b/c/src/main/java/org/apache/arrow/c/Data.java index d870b5884d..3e7b59e2bd 100644 --- a/c/src/main/java/org/apache/arrow/c/Data.java +++ b/c/src/main/java/org/apache/arrow/c/Data.java @@ -17,6 +17,7 @@ package org.apache.arrow.c; import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.util.AutoCloseables; import org.apache.arrow.vector.FieldVector; import org.apache.arrow.vector.VectorLoader; import org.apache.arrow.vector.VectorSchemaRoot; @@ -407,7 +408,12 @@ public static FieldVector importVector( boolean closeImportedStructs) { Field field = importField(allocator, schema, provider, closeImportedStructs); FieldVector vector = field.createVector(allocator); - importIntoVector(allocator, array, vector, provider, closeImportedStructs); + try { + importIntoVector(allocator, array, vector, provider, closeImportedStructs); + } catch (Throwable t) { + AutoCloseables.close(t, vector); + throw t; + } return vector; } diff --git a/c/src/test/java/org/apache/arrow/c/ArrowArrayUtilityTest.java b/c/src/test/java/org/apache/arrow/c/ArrowArrayUtilityTest.java index 511358a5e6..0e170cc62f 100644 --- a/c/src/test/java/org/apache/arrow/c/ArrowArrayUtilityTest.java +++ b/c/src/test/java/org/apache/arrow/c/ArrowArrayUtilityTest.java @@ -107,6 +107,20 @@ allocator, dummyHandle, new ArrowFieldNode(/* length= */ 0, 0), new long[] {0})) } } + @Test + void importBufferRejectsNegativeCapacity() throws Exception { + // Note values are all dummy values here + long address = MemoryUtil.allocateMemory(16); + try (BufferImportTypeVisitor visitor = + new BufferImportTypeVisitor( + allocator, dummyHandle, new ArrowFieldNode(1, 0), new long[] {address})) { + assertThrows( + IllegalStateException.class, () -> visitor.importBuffer(new ArrowType.Bool(), 0, -1)); + } finally { + MemoryUtil.freeMemory(address); + } + } + @Test void cleanupAfterFailure() throws Exception { // Note values are all dummy values here diff --git a/c/src/test/java/org/apache/arrow/c/DictionaryTest.java b/c/src/test/java/org/apache/arrow/c/DictionaryTest.java index 8cd4913f22..202fc9f167 100644 --- a/c/src/test/java/org/apache/arrow/c/DictionaryTest.java +++ b/c/src/test/java/org/apache/arrow/c/DictionaryTest.java @@ -18,6 +18,7 @@ import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; import java.io.ByteArrayInputStream; @@ -26,6 +27,7 @@ import java.nio.channels.Channels; import java.nio.charset.StandardCharsets; import java.util.Collections; +import java.util.function.Consumer; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.util.AutoCloseables; @@ -133,6 +135,105 @@ public void testWithDictionary() throws Exception { AutoCloseables.close((AutoCloseable) vector, encodedVector, dictVector); } + @Test + public void testImportRejectsInvalidLengthWithoutReplacingDictionary() throws Exception { + assertRejectedImportKeepsDictionary( + arrowArray -> { + ArrowArray.Snapshot snapshot = arrowArray.snapshot(); + snapshot.length = -1; + arrowArray.save(snapshot); + }); + } + + @Test + public void testImportRejectsInvalidDictionaryLengthWithoutReplacingDictionary() + throws Exception { + assertRejectedImportKeepsDictionary( + arrowArray -> { + ArrowArray dictionaryArray = ArrowArray.wrap(arrowArray.snapshot().dictionary); + ArrowArray.Snapshot snapshot = dictionaryArray.snapshot(); + snapshot.length = -1; + dictionaryArray.save(snapshot); + }); + } + + @Test + public void testImportRejectsMissingBufferWithoutReplacingDictionary() throws Exception { + assertRejectedImportKeepsDictionary( + arrowArray -> { + ArrowArray.Snapshot snapshot = arrowArray.snapshot(); + snapshot.n_buffers = 1; + arrowArray.save(snapshot); + }); + } + + /** + * Imports a dictionary encoded batch, then checks that a second batch with a different + * dictionary, made invalid by {@code corrupt}, is rejected without replacing the first + * dictionary. + */ + private void assertRejectedImportKeepsDictionary(Consumer corrupt) throws Exception { + DictionaryProvider.MapDictionaryProvider provider = + new DictionaryProvider.MapDictionaryProvider(); + DictionaryProvider.MapDictionaryProvider otherProvider = + new DictionaryProvider.MapDictionaryProvider(); + VarCharVector dictVector = newVarCharVector("dict", "aa", "bb", "cc"); + VarCharVector otherDictVector = newVarCharVector("dict", "xx", "yy"); + VarCharVector vector = newVarCharVector("vector", "bb", "cc"); + VarCharVector otherVector = newVarCharVector("vector", "yy", "xx"); + Dictionary dictionary = new Dictionary(dictVector, new DictionaryEncoding(0L, false, null)); + Dictionary otherDictionary = + new Dictionary(otherDictVector, new DictionaryEncoding(0L, false, null)); + provider.put(dictionary); + otherProvider.put(otherDictionary); + IntVector encoded = (IntVector) DictionaryEncoder.encode(vector, dictionary); + IntVector otherEncoded = (IntVector) DictionaryEncoder.encode(otherVector, otherDictionary); + + try (CDataDictionaryProvider cDictionaryProvider = new CDataDictionaryProvider()) { + FieldVector imported; + try (ArrowSchema arrowSchema = ArrowSchema.allocateNew(allocator); + ArrowArray arrowArray = ArrowArray.allocateNew(allocator)) { + Data.exportVector(allocator, encoded, provider, arrowArray, arrowSchema); + imported = Data.importVector(allocator, arrowArray, arrowSchema, cDictionaryProvider); + } + + try { + // Import a second, invalid batch with a different dictionary into the same vector and + // provider + try (ArrowSchema arrowSchema = ArrowSchema.allocateNew(allocator); + ArrowArray arrowArray = ArrowArray.allocateNew(allocator)) { + Data.exportVector(allocator, otherEncoded, otherProvider, arrowArray, arrowSchema); + // Only the array is imported, so release the exported schema here + arrowSchema.release(); + corrupt.accept(arrowArray); + assertThrows( + IllegalArgumentException.class, + () -> Data.importIntoVector(allocator, arrowArray, imported, cDictionaryProvider)); + } + + assertTrue( + VectorEqualsVisitor.vectorEquals( + dictVector, cDictionaryProvider.lookup(0L).getVector()), + "dictionary was modified by the rejected import"); + } finally { + imported.close(); + } + } finally { + AutoCloseables.close( + (AutoCloseable) vector, otherVector, encoded, otherEncoded, dictVector, otherDictVector); + } + } + + private VarCharVector newVarCharVector(String name, String... values) { + VarCharVector vector = new VarCharVector(name, allocator); + vector.allocateNewSafe(); + for (int i = 0; i < values.length; i++) { + vector.setSafe(i, values[i].getBytes(StandardCharsets.UTF_8)); + } + vector.setValueCount(values.length); + return vector; + } + @Test public void testRoundtripMultipleBatches() throws IOException { try (ArrowStreamReader reader = createMultiBatchReader(); diff --git a/c/src/test/java/org/apache/arrow/c/RoundtripTest.java b/c/src/test/java/org/apache/arrow/c/RoundtripTest.java index ed1690a9e7..4daa3bdd45 100644 --- a/c/src/test/java/org/apache/arrow/c/RoundtripTest.java +++ b/c/src/test/java/org/apache/arrow/c/RoundtripTest.java @@ -29,6 +29,7 @@ import java.util.List; import java.util.Map; import java.util.UUID; +import java.util.function.Consumer; import java.util.stream.Collectors; import java.util.stream.IntStream; import java.util.stream.Stream; @@ -115,6 +116,7 @@ import org.junit.jupiter.params.ParameterizedTest; import org.junit.jupiter.params.provider.Arguments; import org.junit.jupiter.params.provider.MethodSource; +import org.junit.jupiter.params.provider.ValueSource; public class RoundtripTest { private static final String EMPTY_SCHEMA_PATH = ""; @@ -1060,6 +1062,112 @@ public void testSchemaStructReuse() { } } + @ParameterizedTest + @ValueSource(longs = {-1L, Long.MIN_VALUE, Integer.MAX_VALUE + 1L, (1L << 32) - 100}) + public void testImportRejectsInvalidLength(long length) { + try (IntVector vector = new IntVector("ints", allocator)) { + setVector(vector, 1, 2, 3); + assertImportRejected(vector, "ints", array -> setLength(array, length)); + } + } + + @ParameterizedTest + @ValueSource(longs = {-1L, Long.MIN_VALUE, Integer.MAX_VALUE + 1L, (1L << 32) - 100}) + public void testImportRejectsInvalidStructLength(long length) { + try (StructVector vector = createStructWithTwoChildren()) { + assertImportRejected(vector, "struct", array -> setLength(array, length)); + } + } + + @ParameterizedTest + @ValueSource(longs = {-1L, Long.MIN_VALUE, Integer.MAX_VALUE + 1L, (1L << 32) - 100}) + public void testImportRejectsInvalidChildLength(long length) { + try (StructVector vector = createStructWithTwoChildren()) { + assertImportRejected( + vector, + "second", + array -> { + long[] children = NativeUtil.toJavaArray(array.snapshot().children, 2); + setLength(ArrowArray.wrap(children[1]), length); + }); + } + } + + @ParameterizedTest + @ValueSource(longs = {-1L, 3L, Integer.MAX_VALUE, Long.MAX_VALUE}) + public void testImportRejectsInvalidBufferCount(long bufferCount) { + try (IntVector vector = new IntVector("ints", allocator)) { + setVector(vector, 1, 2, 3); + assertImportRejected( + vector, + "ints", + array -> { + ArrowArray.Snapshot snapshot = array.snapshot(); + snapshot.n_buffers = bufferCount; + array.save(snapshot); + }); + } + } + + @Test + public void testImportRejectsInvalidChildCount() { + try (StructVector vector = createStructWithTwoChildren(); + ArrowSchema arrowSchema = ArrowSchema.allocateNew(allocator); + ArrowArray arrowArray = ArrowArray.allocateNew(allocator)) { + Data.exportVector(allocator, vector, null, arrowArray, arrowSchema); + ArrowArray.Snapshot snapshot = arrowArray.snapshot(); + // Releasing the array skips the children of a negative count, so release them first + for (long child : NativeUtil.toJavaArray(snapshot.children, 2)) { + ArrowArray.wrap(child).release(); + } + snapshot.n_children = -1; + arrowArray.save(snapshot); + IllegalStateException e = + assertThrows( + IllegalStateException.class, + () -> Data.importVector(childAllocator, arrowArray, arrowSchema, null)); + assertTrue(e.getMessage().contains("children"), e.getMessage()); + } + assertEquals(0, childAllocator.getAllocatedMemory()); + } + + private StructVector createStructWithTwoChildren() { + StructVector vector = StructVector.empty("struct", allocator); + FieldType type = FieldType.nullable(new ArrowType.Int(32, true)); + IntVector first = vector.addOrGet("first", type, IntVector.class); + IntVector second = vector.addOrGet("second", type, IntVector.class); + vector.allocateNew(); + for (int i = 0; i < 3; i++) { + first.setSafe(i, i); + second.setSafe(i, i); + vector.setIndexDefined(i); + } + vector.setValueCount(3); + return vector; + } + + private static void setLength(ArrowArray array, long length) { + ArrowArray.Snapshot snapshot = array.snapshot(); + snapshot.length = length; + array.save(snapshot); + } + + /** Exports the vector, corrupts the export and checks that the import fails cleanly. */ + private void assertImportRejected( + FieldVector vector, String expectedField, Consumer corrupt) { + try (ArrowSchema arrowSchema = ArrowSchema.allocateNew(allocator); + ArrowArray arrowArray = ArrowArray.allocateNew(allocator)) { + Data.exportVector(allocator, vector, null, arrowArray, arrowSchema); + corrupt.accept(arrowArray); + IllegalArgumentException e = + assertThrows( + IllegalArgumentException.class, + () -> Data.importVector(childAllocator, arrowArray, arrowSchema, null)); + assertTrue(e.getMessage().contains(expectedField), e.getMessage()); + } + assertEquals(0, childAllocator.getAllocatedMemory()); + } + private void exportSchema(Schema schema, ArrowSchema targetArrowSchema) { try (ArrowSchema arrowSchema = ArrowSchema.wrap(targetArrowSchema.memoryAddress())) { Data.exportSchema(allocator, schema, null, arrowSchema); diff --git a/vector/src/main/java/org/apache/arrow/vector/ipc/JsonFileReader.java b/vector/src/main/java/org/apache/arrow/vector/ipc/JsonFileReader.java index e4bab7eb80..85dd501a1b 100644 --- a/vector/src/main/java/org/apache/arrow/vector/ipc/JsonFileReader.java +++ b/vector/src/main/java/org/apache/arrow/vector/ipc/JsonFileReader.java @@ -52,6 +52,7 @@ import java.util.Set; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.util.AutoCloseables; import org.apache.arrow.util.Preconditions; import org.apache.arrow.vector.BaseVariableWidthViewVector; import org.apache.arrow.vector.BigIntVector; @@ -196,7 +197,7 @@ public boolean read(VectorSchemaRoot root) throws IOException { JsonToken t = parser.nextToken(); if (t == START_OBJECT) { { - int count = readNextField("count", Integer.class); + int count = readCount(); nextFieldIs("columns"); readToken(START_ARRAY); { @@ -223,9 +224,8 @@ public VectorSchemaRoot read() throws IOException { JsonToken t = parser.nextToken(); if (t == START_OBJECT) { VectorSchemaRoot recordBatch = VectorSchemaRoot.create(schema, allocator); - { - int count = readNextField("count", Integer.class); - recordBatch.setRowCount(count); + try { + int count = readCount(); nextFieldIs("columns"); readToken(START_ARRAY); { @@ -235,8 +235,14 @@ public VectorSchemaRoot read() throws IOException { } } readToken(END_ARRAY); + // Set the row count once the columns are read, as read(VectorSchemaRoot) does, so that + // the vectors are not first sized for a count the columns may not match + recordBatch.setRowCount(count); + readToken(END_OBJECT); + } catch (Throwable e) { + AutoCloseables.close(e, recordBatch); + throw e; } - readToken(END_OBJECT); return recordBatch; } else if (t == END_ARRAY) { return null; @@ -268,13 +274,29 @@ public int skip(int numBatches) throws IOException { } private abstract class BufferReader { + // Buffer allocated by the current read, released if the read fails before returning it + private ArrowBuf allocated; + protected abstract ArrowBuf read(BufferAllocator allocator, int count) throws IOException; + protected ArrowBuf allocate(BufferAllocator allocator, long size) { + allocated = allocator.buffer(size); + return allocated; + } + ArrowBuf readBuffer(BufferAllocator allocator, int count) throws IOException { readToken(START_ARRAY); - ArrowBuf buf = read(allocator, count); - readToken(END_ARRAY); - return buf; + ArrowBuf buf = null; + try { + buf = read(allocator, count); + readToken(END_ARRAY); + return buf; + } catch (Throwable e) { + AutoCloseables.close(e, buf != null ? buf : allocated); + throw e; + } finally { + allocated = null; + } } } @@ -290,22 +312,27 @@ List readVariadicBuffers(BufferAllocator allocator, int variadicBuffer throws IOException { readToken(START_ARRAY); ArrayList dataBuffers = new ArrayList<>(variadicBuffersCount); - for (int i = 0; i < variadicBuffersCount; i++) { - parser.nextToken(); - final byte[] value; + try { + for (int i = 0; i < variadicBuffersCount; i++) { + parser.nextToken(); + final byte[] value; - String variadicStr = parser.readValueAs(String.class); - if (variadicStr == null) { - value = new byte[0]; - } else { - value = decodeHexSafe(variadicStr); - } + String variadicStr = parser.readValueAs(String.class); + if (variadicStr == null) { + value = new byte[0]; + } else { + value = decodeHexSafe(variadicStr); + } - ArrowBuf buf = allocator.buffer(value.length); - buf.writeBytes(value); - dataBuffers.add(buf); + ArrowBuf buf = allocator.buffer(value.length); + dataBuffers.add(buf); + buf.writeBytes(value); + } + readToken(END_ARRAY); + } catch (Throwable e) { + AutoCloseables.close(e, dataBuffers); + throw e; } - readToken(END_ARRAY); return dataBuffers; } @@ -368,11 +395,15 @@ private ArrowBuf readViewBuffers( } ArrowBuf buf = allocator.buffer(bufferSize); - - for (byte[] value : values) { - buf.writeBytes(value); + try { + for (byte[] value : values) { + buf.writeBytes(value); + } + readToken(END_ARRAY); + } catch (Throwable e) { + AutoCloseables.close(e, buf); + throw e; } - readToken(END_ARRAY); return buf; } @@ -383,7 +414,7 @@ private class BufferHelper { @Override protected ArrowBuf read(BufferAllocator allocator, int count) throws IOException { final int bufferSize = getValidityBufferSizeFromCount(count); - ArrowBuf buf = allocator.buffer(bufferSize); + ArrowBuf buf = allocate(allocator, bufferSize); // C++ integration test fails without this. buf.setZero(0, bufferSize); @@ -403,7 +434,7 @@ protected ArrowBuf read(BufferAllocator allocator, int count) throws IOException @Override protected ArrowBuf read(BufferAllocator allocator, int count) throws IOException { final long size = (long) count * IntervalDayVector.TYPE_WIDTH; - ArrowBuf buf = allocator.buffer(size); + ArrowBuf buf = allocate(allocator, size); for (int i = 0; i < count; i++) { readToken(START_OBJECT); @@ -421,7 +452,7 @@ protected ArrowBuf read(BufferAllocator allocator, int count) throws IOException @Override protected ArrowBuf read(BufferAllocator allocator, int count) throws IOException { final long size = (long) count * IntervalMonthDayNanoVector.TYPE_WIDTH; - ArrowBuf buf = allocator.buffer(size); + ArrowBuf buf = allocate(allocator, size); for (int i = 0; i < count; i++) { readToken(START_OBJECT); @@ -440,7 +471,7 @@ protected ArrowBuf read(BufferAllocator allocator, int count) throws IOException @Override protected ArrowBuf read(BufferAllocator allocator, int count) throws IOException { final long size = (long) count * TinyIntVector.TYPE_WIDTH; - ArrowBuf buf = allocator.buffer(size); + ArrowBuf buf = allocate(allocator, size); for (int i = 0; i < count; i++) { parser.nextToken(); @@ -456,7 +487,7 @@ protected ArrowBuf read(BufferAllocator allocator, int count) throws IOException @Override protected ArrowBuf read(BufferAllocator allocator, int count) throws IOException { final long size = (long) count * SmallIntVector.TYPE_WIDTH; - ArrowBuf buf = allocator.buffer(size); + ArrowBuf buf = allocate(allocator, size); for (int i = 0; i < count; i++) { parser.nextToken(); @@ -472,7 +503,7 @@ protected ArrowBuf read(BufferAllocator allocator, int count) throws IOException @Override protected ArrowBuf read(BufferAllocator allocator, int count) throws IOException { final long size = (long) count * IntVector.TYPE_WIDTH; - ArrowBuf buf = allocator.buffer(size); + ArrowBuf buf = allocate(allocator, size); for (int i = 0; i < count; i++) { parser.nextToken(); @@ -488,7 +519,7 @@ protected ArrowBuf read(BufferAllocator allocator, int count) throws IOException @Override protected ArrowBuf read(BufferAllocator allocator, int count) throws IOException { final long size = (long) count * BigIntVector.TYPE_WIDTH; - ArrowBuf buf = allocator.buffer(size); + ArrowBuf buf = allocate(allocator, size); for (int i = 0; i < count; i++) { parser.nextToken(); @@ -505,7 +536,7 @@ protected ArrowBuf read(BufferAllocator allocator, int count) throws IOException @Override protected ArrowBuf read(BufferAllocator allocator, int count) throws IOException { final long size = (long) count * TinyIntVector.TYPE_WIDTH; - ArrowBuf buf = allocator.buffer(size); + ArrowBuf buf = allocate(allocator, size); for (int i = 0; i < count; i++) { parser.nextToken(); @@ -521,7 +552,7 @@ protected ArrowBuf read(BufferAllocator allocator, int count) throws IOException @Override protected ArrowBuf read(BufferAllocator allocator, int count) throws IOException { final long size = (long) count * SmallIntVector.TYPE_WIDTH; - ArrowBuf buf = allocator.buffer(size); + ArrowBuf buf = allocate(allocator, size); for (int i = 0; i < count; i++) { parser.nextToken(); @@ -537,7 +568,7 @@ protected ArrowBuf read(BufferAllocator allocator, int count) throws IOException @Override protected ArrowBuf read(BufferAllocator allocator, int count) throws IOException { final long size = (long) count * IntVector.TYPE_WIDTH; - ArrowBuf buf = allocator.buffer(size); + ArrowBuf buf = allocate(allocator, size); for (int i = 0; i < count; i++) { parser.nextToken(); @@ -553,7 +584,7 @@ protected ArrowBuf read(BufferAllocator allocator, int count) throws IOException @Override protected ArrowBuf read(BufferAllocator allocator, int count) throws IOException { final long size = (long) count * BigIntVector.TYPE_WIDTH; - ArrowBuf buf = allocator.buffer(size); + ArrowBuf buf = allocate(allocator, size); for (int i = 0; i < count; i++) { parser.nextToken(); @@ -570,7 +601,7 @@ protected ArrowBuf read(BufferAllocator allocator, int count) throws IOException @Override protected ArrowBuf read(BufferAllocator allocator, int count) throws IOException { final long size = (long) count * Float4Vector.TYPE_WIDTH; - ArrowBuf buf = allocator.buffer(size); + ArrowBuf buf = allocate(allocator, size); for (int i = 0; i < count; i++) { parser.nextToken(); @@ -586,7 +617,7 @@ protected ArrowBuf read(BufferAllocator allocator, int count) throws IOException @Override protected ArrowBuf read(BufferAllocator allocator, int count) throws IOException { final long size = (long) count * Float8Vector.TYPE_WIDTH; - ArrowBuf buf = allocator.buffer(size); + ArrowBuf buf = allocate(allocator, size); for (int i = 0; i < count; i++) { parser.nextToken(); @@ -602,7 +633,7 @@ protected ArrowBuf read(BufferAllocator allocator, int count) throws IOException @Override protected ArrowBuf read(BufferAllocator allocator, int count) throws IOException { final long size = (long) count * DecimalVector.TYPE_WIDTH; - ArrowBuf buf = allocator.buffer(size); + ArrowBuf buf = allocate(allocator, size); for (int i = 0; i < count; i++) { parser.nextToken(); @@ -621,7 +652,7 @@ protected ArrowBuf read(BufferAllocator allocator, int count) throws IOException @Override protected ArrowBuf read(BufferAllocator allocator, int count) throws IOException { final long size = (long) count * Decimal256Vector.TYPE_WIDTH; - ArrowBuf buf = allocator.buffer(size); + ArrowBuf buf = allocate(allocator, size); for (int i = 0; i < count; i++) { parser.nextToken(); @@ -874,7 +905,7 @@ private void readFromJsonIntoVector(Field field, FieldVector vector) throws IOEx * releasing the local buffers. */ readToken(START_OBJECT); - { + try { // If currently reading dictionaries, field name is not important so don't check String name = readNextField("name", String.class); if (started && !Objects.equals(field.getName(), name)) { @@ -885,7 +916,7 @@ private void readFromJsonIntoVector(Field field, FieldVector vector) throws IOEx /* Initialize the vector with required capacity but don't allocateNew since we would * be doing loadFieldBuffers. */ - int valueCount = readNextField("count", Integer.class); + int valueCount = readCount(); vector.setInitialCapacity(valueCount); @@ -898,7 +929,7 @@ private void readFromJsonIntoVector(Field field, FieldVector vector) throws IOEx && !(type instanceof ListView) && !(type instanceof LargeListView)) { /* offset buffer has 1 additional value capacity except for dense unions and ListView */ - innerBufferValueCount = valueCount + 1; + innerBufferValueCount = Math.addExact(valueCount, 1); } vectorBuffers.addAll( @@ -941,13 +972,13 @@ private void readFromJsonIntoVector(Field field, FieldVector vector) throws IOEx } readToken(END_ARRAY); } + } finally { + for (ArrowBuf buffer : vectorBuffers) { + buffer.getReferenceManager().release(); + } } readToken(END_OBJECT); - - for (ArrowBuf buffer : vectorBuffers) { - buffer.getReferenceManager().release(); - } } private byte[] decodeHexSafe(String hexString) throws IOException { @@ -975,6 +1006,14 @@ private T readNextField(String expectedFieldName, Class c) return parser.readValueAs(c); } + private int readCount() throws IOException { + Integer count = readNextField("count", Integer.class); + if (count == null || count < 0) { + throw new IllegalArgumentException("Expected count >= 0, got " + count); + } + return count; + } + private void nextFieldIs(String expectedFieldName) throws IOException, JsonParseException { String name = parser.nextFieldName(); if (name == null || !name.equals(expectedFieldName)) { diff --git a/vector/src/main/java/org/apache/arrow/vector/ipc/message/ArrowFieldNode.java b/vector/src/main/java/org/apache/arrow/vector/ipc/message/ArrowFieldNode.java index ad3f434e12..9f8f931247 100644 --- a/vector/src/main/java/org/apache/arrow/vector/ipc/message/ArrowFieldNode.java +++ b/vector/src/main/java/org/apache/arrow/vector/ipc/message/ArrowFieldNode.java @@ -16,10 +16,9 @@ */ package org.apache.arrow.vector.ipc.message; -import static org.apache.arrow.memory.util.LargeMemoryUtil.checkedCastToInt; - import com.google.flatbuffers.FlatBufferBuilder; import org.apache.arrow.flatbuf.FieldNode; +import org.apache.arrow.util.Preconditions; /** Metadata about Vectors/Arrays that is written to a channel. */ public class ArrowFieldNode implements FBSerializable { @@ -31,12 +30,14 @@ public class ArrowFieldNode implements FBSerializable { * Constructs a new instance. * * @param length The number of values written. - * @param nullCount The number of null values. + * @param nullCount The number of null values, or -1 if unknown. */ public ArrowFieldNode(long length, long nullCount) { super(); - this.length = checkedCastToInt(length); - this.nullCount = checkedCastToInt(nullCount); + Preconditions.checkArgument(length >= 0, "Expected length >= 0, got %s", length); + Preconditions.checkArgument(nullCount >= -1, "Expected nullCount >= -1, got %s", nullCount); + this.length = Math.toIntExact(length); + this.nullCount = Math.toIntExact(nullCount); } @Override diff --git a/vector/src/main/java/org/apache/arrow/vector/ipc/message/MessageSerializer.java b/vector/src/main/java/org/apache/arrow/vector/ipc/message/MessageSerializer.java index 36f6ea449b..589ddfdc56 100644 --- a/vector/src/main/java/org/apache/arrow/vector/ipc/message/MessageSerializer.java +++ b/vector/src/main/java/org/apache/arrow/vector/ipc/message/MessageSerializer.java @@ -383,21 +383,28 @@ public static ArrowRecordBatch deserializeRecordBatch( long totalLen = block.getMetadataLength() + block.getBodyLength(); ArrowBuf buffer = alloc.buffer(totalLen); - if (in.readFully(buffer, totalLen) != totalLen) { - throw new IOException("Unexpected end of input trying to read batch."); - } + final RecordBatch recordBatchFB; + final ArrowBuf body; + try { + if (in.readFully(buffer, totalLen) != totalLen) { + throw new IOException("Unexpected end of input trying to read batch."); + } - int prefixSize = buffer.getInt(0) == IPC_CONTINUATION_TOKEN ? 8 : 4; + int prefixSize = buffer.getInt(0) == IPC_CONTINUATION_TOKEN ? 8 : 4; - ArrowBuf metadataBuffer = buffer.slice(prefixSize, block.getMetadataLength() - prefixSize); + ArrowBuf metadataBuffer = buffer.slice(prefixSize, block.getMetadataLength() - prefixSize); - Message messageFB = Message.getRootAsMessage(metadataBuffer.nioBuffer().asReadOnlyBuffer()); + Message messageFB = Message.getRootAsMessage(metadataBuffer.nioBuffer().asReadOnlyBuffer()); - RecordBatch recordBatchFB = (RecordBatch) messageFB.header(new RecordBatch()); + recordBatchFB = (RecordBatch) messageFB.header(new RecordBatch()); - // Now read the body - final ArrowBuf body = - buffer.slice(block.getMetadataLength(), totalLen - block.getMetadataLength()); + // Now read the body + body = buffer.slice(block.getMetadataLength(), totalLen - block.getMetadataLength()); + } catch (RuntimeException | IOException e) { + buffer.close(); + throw e; + } + // Releases the body, and so the buffer, once the batch is deserialized return deserializeRecordBatch(recordBatchFB, body); } @@ -411,50 +418,89 @@ public static ArrowRecordBatch deserializeRecordBatch( */ public static ArrowRecordBatch deserializeRecordBatch(RecordBatch recordBatchFB, ArrowBuf body) throws IOException { - // Now read the body - int nodesLength = recordBatchFB.nodesLength(); - List nodes = new ArrayList<>(); - for (int i = 0; i < nodesLength; ++i) { - FieldNode node = recordBatchFB.nodes(i); - if ((int) node.length() != node.length() || (int) node.nullCount() != node.nullCount()) { - throw new IOException( - "Cannot currently deserialize record batches with " - + "node length larger than INT_MAX records."); + try { + // Now read the body + int nodesLength = recordBatchFB.nodesLength(); + List nodes = new ArrayList<>(); + for (int i = 0; i < nodesLength; ++i) { + FieldNode node = recordBatchFB.nodes(i); + if (node.length() < 0) { + throw new IOException("Record batch node " + i + " has negative length " + node.length()); + } + if (node.nullCount() < -1) { + throw new IOException( + "Record batch node " + i + " has invalid null count " + node.nullCount()); + } + if ((int) node.length() != node.length() || (int) node.nullCount() != node.nullCount()) { + throw new IOException( + "Cannot currently deserialize record batches with " + + "node length larger than INT_MAX records."); + } + nodes.add(new ArrowFieldNode(node.length(), node.nullCount())); + } + List buffers = new ArrayList<>(); + for (int i = 0; i < recordBatchFB.buffersLength(); ++i) { + Buffer bufferFB = recordBatchFB.buffers(i); + long offset = bufferFB.offset(); + long length = bufferFB.length(); + if (offset < 0 || length < 0 || offset > body.capacity() - length) { + throw new IOException( + "Record batch buffer " + + i + + " (offset " + + offset + + ", length " + + length + + ") does not fit in a body of " + + body.capacity() + + " bytes"); + } + ArrowBuf vectorBuffer = body.slice(offset, length); + buffers.add(vectorBuffer); } - nodes.add(new ArrowFieldNode(node.length(), node.nullCount())); - } - List buffers = new ArrayList<>(); - for (int i = 0; i < recordBatchFB.buffersLength(); ++i) { - Buffer bufferFB = recordBatchFB.buffers(i); - ArrowBuf vectorBuffer = body.slice(bufferFB.offset(), bufferFB.length()); - buffers.add(vectorBuffer); - } - - ArrowBodyCompression bodyCompression = - recordBatchFB.compression() == null - ? NoCompressionCodec.DEFAULT_BODY_COMPRESSION - : new ArrowBodyCompression( - recordBatchFB.compression().codec(), recordBatchFB.compression().method()); - List variadicBufferCounts = new ArrayList<>(); - for (int i = 0; i < recordBatchFB.variadicBufferCountsLength(); i++) { - variadicBufferCounts.add(recordBatchFB.variadicBufferCounts(i)); - } + ArrowBodyCompression bodyCompression = + recordBatchFB.compression() == null + ? NoCompressionCodec.DEFAULT_BODY_COMPRESSION + : new ArrowBodyCompression( + recordBatchFB.compression().codec(), recordBatchFB.compression().method()); + + List variadicBufferCounts = new ArrayList<>(); + long totalVariadicBufferCount = 0; + for (int i = 0; i < recordBatchFB.variadicBufferCountsLength(); i++) { + long count = recordBatchFB.variadicBufferCounts(i); + // Variadic buffers are also listed in the batch buffers, so they cannot outnumber them + if (count < 0 || count > recordBatchFB.buffersLength() - totalVariadicBufferCount) { + throw new IOException( + "Record batch variadic buffer count " + + i + + " is " + + count + + ", but the batch has " + + recordBatchFB.buffersLength() + + " buffers"); + } + totalVariadicBufferCount += count; + variadicBufferCounts.add(count); + } - if ((int) recordBatchFB.length() != recordBatchFB.length()) { - throw new IOException( - "Cannot currently deserialize record batches with more than INT_MAX records."); + if (recordBatchFB.length() < 0) { + throw new IOException("Record batch has negative length " + recordBatchFB.length()); + } + if ((int) recordBatchFB.length() != recordBatchFB.length()) { + throw new IOException( + "Cannot currently deserialize record batches with more than INT_MAX records."); + } + return new ArrowRecordBatch( + checkedCastToInt(recordBatchFB.length()), + nodes, + buffers, + bodyCompression, + variadicBufferCounts, + /*alignBuffers*/ true); + } finally { + body.getReferenceManager().release(); } - ArrowRecordBatch arrowRecordBatch = - new ArrowRecordBatch( - checkedCastToInt(recordBatchFB.length()), - nodes, - buffers, - bodyCompression, - variadicBufferCounts, - /*alignBuffers*/ true); - body.getReferenceManager().release(); - return arrowRecordBatch; } /** @@ -588,22 +634,31 @@ public static ArrowDictionaryBatch deserializeDictionaryBatch( long totalLen = block.getMetadataLength() + block.getBodyLength(); ArrowBuf buffer = alloc.buffer(totalLen); - if (in.readFully(buffer, totalLen) != totalLen) { - throw new IOException("Unexpected end of input trying to read batch."); - } + final DictionaryBatch dictionaryBatchFB; + final RecordBatch recordBatchFB; + final ArrowBuf body; + try { + if (in.readFully(buffer, totalLen) != totalLen) { + throw new IOException("Unexpected end of input trying to read batch."); + } - int prefixSize = buffer.getInt(0) == IPC_CONTINUATION_TOKEN ? 8 : 4; + int prefixSize = buffer.getInt(0) == IPC_CONTINUATION_TOKEN ? 8 : 4; - ArrowBuf metadataBuffer = buffer.slice(prefixSize, block.getMetadataLength() - prefixSize); + ArrowBuf metadataBuffer = buffer.slice(prefixSize, block.getMetadataLength() - prefixSize); - Message messageFB = Message.getRootAsMessage(metadataBuffer.nioBuffer().asReadOnlyBuffer()); + Message messageFB = Message.getRootAsMessage(metadataBuffer.nioBuffer().asReadOnlyBuffer()); - DictionaryBatch dictionaryBatchFB = (DictionaryBatch) messageFB.header(new DictionaryBatch()); + dictionaryBatchFB = (DictionaryBatch) messageFB.header(new DictionaryBatch()); + recordBatchFB = dictionaryBatchFB.data(); - // Now read the body - final ArrowBuf body = - buffer.slice(block.getMetadataLength(), totalLen - block.getMetadataLength()); - ArrowRecordBatch recordBatch = deserializeRecordBatch(dictionaryBatchFB.data(), body); + // Now read the body + body = buffer.slice(block.getMetadataLength(), totalLen - block.getMetadataLength()); + } catch (RuntimeException | IOException e) { + buffer.close(); + throw e; + } + // Releases the body, and so the buffer, once the batch is deserialized + ArrowRecordBatch recordBatch = deserializeRecordBatch(recordBatchFB, body); return new ArrowDictionaryBatch( dictionaryBatchFB.id(), recordBatch, dictionaryBatchFB.isDelta()); } @@ -620,25 +675,41 @@ public static ArrowMessage deserializeMessageBatch(MessageChannelReader reader) MessageResult result = reader.readNext(); if (result == null) { return null; - } else if (result.getMessage().bodyLength() > Integer.MAX_VALUE) { - throw new IOException("Cannot currently deserialize record batches over 2GB"); } - if (result.getMessage().version() != MetadataVersion.V4 - && result.getMessage().version() != MetadataVersion.V5) { - throw new IOException( - "Received metadata with an incompatible version number: " - + result.getMessage().version()); + ArrowBuf bodyBuffer = result.getBodyBuffer(); + // For zero-length batches, need an empty buffer to deserialize the batch + if (bodyBuffer == null) { + bodyBuffer = reader.allocator.getEmpty(); } - switch (result.getMessage().headerType()) { - case MessageHeader.RecordBatch: - return deserializeRecordBatch(result.getMessage(), result.getBodyBuffer()); - case MessageHeader.DictionaryBatch: - return deserializeDictionaryBatch(result.getMessage(), result.getBodyBuffer()); - default: + try { + if (result.getMessage().bodyLength() > Integer.MAX_VALUE) { + throw new IOException("Cannot currently deserialize record batches over 2GB"); + } + + if (result.getMessage().version() != MetadataVersion.V4 + && result.getMessage().version() != MetadataVersion.V5) { + throw new IOException( + "Received metadata with an incompatible version number: " + + result.getMessage().version()); + } + + if (result.getMessage().headerType() != MessageHeader.RecordBatch + && result.getMessage().headerType() != MessageHeader.DictionaryBatch) { throw new IOException("Unexpected message header type " + result.getMessage().headerType()); + } + } catch (IOException e) { + // The body was read with the message but will not be deserialized + bodyBuffer.close(); + throw e; } + + // Deserializing releases the body + if (result.getMessage().headerType() == MessageHeader.RecordBatch) { + return deserializeRecordBatch(result.getMessage(), bodyBuffer); + } + return deserializeDictionaryBatch(result.getMessage(), bodyBuffer); } /** diff --git a/vector/src/test/java/org/apache/arrow/vector/ipc/MessageSerializerTest.java b/vector/src/test/java/org/apache/arrow/vector/ipc/MessageSerializerTest.java index b529ca645a..910dd7d062 100644 --- a/vector/src/test/java/org/apache/arrow/vector/ipc/MessageSerializerTest.java +++ b/vector/src/test/java/org/apache/arrow/vector/ipc/MessageSerializerTest.java @@ -23,6 +23,7 @@ import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; +import com.google.flatbuffers.FlatBufferBuilder; import java.io.ByteArrayInputStream; import java.io.ByteArrayOutputStream; import java.io.IOException; @@ -32,9 +33,15 @@ import java.util.Arrays; import java.util.Collections; import java.util.List; +import java.util.stream.Stream; +import org.apache.arrow.flatbuf.Buffer; +import org.apache.arrow.flatbuf.FieldNode; +import org.apache.arrow.flatbuf.MessageHeader; +import org.apache.arrow.flatbuf.RecordBatch; 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.ArrowBlock; import org.apache.arrow.vector.ipc.message.ArrowFieldNode; import org.apache.arrow.vector.ipc.message.ArrowMessage; import org.apache.arrow.vector.ipc.message.ArrowRecordBatch; @@ -47,6 +54,10 @@ import org.apache.arrow.vector.types.pojo.FieldType; import org.apache.arrow.vector.types.pojo.Schema; import org.junit.jupiter.api.Test; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.Arguments; +import org.junit.jupiter.params.provider.MethodSource; +import org.junit.jupiter.params.provider.ValueSource; public class MessageSerializerTest { @@ -220,6 +231,234 @@ public void testSerializeRecordBatchV5() throws Exception { alloc.close(); } + @ParameterizedTest + @MethodSource("invalidFieldNodes") + public void testInvalidFieldNodeRejectedWithoutLeakingBody(long length, long nullCount) { + RecordBatch recordBatchFB = recordBatchWithNode(length, nullCount); + try (BufferAllocator allocator = new RootAllocator()) { + ArrowBuf body = allocator.buffer(8); + assertThrows( + IOException.class, () -> MessageSerializer.deserializeRecordBatch(recordBatchFB, body)); + assertEquals(0, allocator.getAllocatedMemory()); + } + } + + static Stream invalidFieldNodes() { + return Stream.of( + Arguments.of(-100L, 0L), + Arguments.of(Long.MIN_VALUE, 0L), + Arguments.of(Integer.MAX_VALUE + 1L, 0L), + Arguments.of(1L << 32, 0L), + Arguments.of(10L, -2L)); + } + + @ParameterizedTest + @MethodSource("validFieldNodes") + public void testValidFieldNodeDeserialized(long length, long nullCount) throws IOException { + RecordBatch recordBatchFB = recordBatchWithNode(length, nullCount); + try (BufferAllocator allocator = new RootAllocator()) { + ArrowBuf body = allocator.buffer(8); + try (ArrowRecordBatch batch = MessageSerializer.deserializeRecordBatch(recordBatchFB, body)) { + ArrowFieldNode node = batch.getNodes().get(0); + assertEquals(length, node.getLength()); + assertEquals(nullCount, node.getNullCount()); + } + assertEquals(0, allocator.getAllocatedMemory()); + } + } + + static Stream validFieldNodes() { + return Stream.of( + Arguments.of(0L, 0L), Arguments.of(10L, -1L), Arguments.of((long) Integer.MAX_VALUE, 0L)); + } + + @Test + public void testNegativeRecordBatchLengthRejectedWithoutLeakingBody() { + RecordBatch recordBatchFB = recordBatch(-100, 0, 0); + try (BufferAllocator allocator = new RootAllocator()) { + ArrowBuf body = allocator.buffer(8); + assertThrows( + IOException.class, () -> MessageSerializer.deserializeRecordBatch(recordBatchFB, body)); + assertEquals(0, allocator.getAllocatedMemory()); + } + } + + @ParameterizedTest + @MethodSource("invalidBuffers") + public void testInvalidBufferRejectedWithoutLeakingBody(long offset, long length) { + RecordBatch recordBatchFB = recordBatch(0, 0, 0, offset, length); + try (BufferAllocator allocator = new RootAllocator()) { + ArrowBuf body = allocator.buffer(8); + assertThrows( + IOException.class, () -> MessageSerializer.deserializeRecordBatch(recordBatchFB, body)); + assertEquals(0, allocator.getAllocatedMemory()); + } + } + + static Stream invalidBuffers() { + return Stream.of( + // Negative length, but offset + length is still within the body + Arguments.of(8L, -4L), + Arguments.of(-8L, 8L), + Arguments.of(4L, 8L), + Arguments.of(8L, Long.MAX_VALUE)); + } + + @Test + public void testDeserializeMessageBatchWithoutBody() throws Exception { + ByteArrayOutputStream out = new ByteArrayOutputStream(); + try (ArrowRecordBatch batch = + new ArrowRecordBatch(0, Collections.emptyList(), Collections.emptyList())) { + MessageSerializer.serialize(new WriteChannel(Channels.newChannel(out)), batch); + } + + try (BufferAllocator allocator = new RootAllocator()) { + ReadChannel channel = + new ReadChannel(Channels.newChannel(new ByteArrayInputStream(out.toByteArray()))); + try (ArrowMessage deserialized = + MessageSerializer.deserializeMessageBatch(channel, allocator)) { + assertEquals(ArrowRecordBatch.class, deserialized.getClass()); + assertEquals(0, ((ArrowRecordBatch) deserialized).getLength()); + } + assertEquals(0, allocator.getAllocatedMemory()); + } + } + + @ParameterizedTest + @ValueSource(longs = {-1L, 2L, Integer.MAX_VALUE + 1L, Long.MAX_VALUE}) + public void testInvalidVariadicBufferCountRejectedWithoutLeakingBody(long count) { + RecordBatch recordBatchFB = recordBatchWithVariadicBufferCount(count); + try (BufferAllocator allocator = new RootAllocator()) { + ArrowBuf body = allocator.buffer(8); + assertThrows( + IOException.class, () -> MessageSerializer.deserializeRecordBatch(recordBatchFB, body)); + assertEquals(0, allocator.getAllocatedMemory()); + } + } + + @Test + public void testVariadicBufferCountDeserialized() throws IOException { + RecordBatch recordBatchFB = recordBatchWithVariadicBufferCount(1); + try (BufferAllocator allocator = new RootAllocator()) { + ArrowBuf body = allocator.buffer(8); + try (ArrowRecordBatch batch = MessageSerializer.deserializeRecordBatch(recordBatchFB, body)) { + assertEquals(Collections.singletonList(1L), batch.getVariadicBufferCounts()); + } + assertEquals(0, allocator.getAllocatedMemory()); + } + } + + @ParameterizedTest + @MethodSource("unsupportedMessages") + public void testUnsupportedMessageRejectedWithoutLeakingBody( + byte headerType, MetadataVersion version, String expectedMessage) throws IOException { + byte[] message = messageWithBody(headerType, version); + try (BufferAllocator allocator = new RootAllocator()) { + ReadChannel channel = readChannel(message); + IOException e = + assertThrows( + IOException.class, + () -> MessageSerializer.deserializeMessageBatch(channel, allocator)); + assertTrue(e.getMessage().contains(expectedMessage), e.getMessage()); + assertEquals(0, allocator.getAllocatedMemory()); + } + } + + static Stream unsupportedMessages() { + return Stream.of( + Arguments.of(MessageHeader.RecordBatch, MetadataVersion.V3, "incompatible version"), + Arguments.of(MessageHeader.Tensor, MetadataVersion.V5, "Unexpected message header type")); + } + + @Test + public void testTruncatedBlockRejectedWithoutLeaking() { + ArrowBlock block = new ArrowBlock(0, 16, 8); + try (BufferAllocator allocator = new RootAllocator()) { + assertThrows( + IOException.class, + () -> + MessageSerializer.deserializeRecordBatch(readChannel(new byte[4]), block, allocator)); + assertThrows( + IOException.class, + () -> + MessageSerializer.deserializeDictionaryBatch( + readChannel(new byte[4]), block, allocator)); + assertEquals(0, allocator.getAllocatedMemory()); + } + } + + private static ReadChannel readChannel(byte[] bytes) { + return new ReadChannel(Channels.newChannel(new ByteArrayInputStream(bytes))); + } + + /** Writes a message of the given header type and version, with an 8 byte body. */ + private static byte[] messageWithBody(byte headerType, MetadataVersion version) + throws IOException { + FlatBufferBuilder builder = new FlatBufferBuilder(); + RecordBatch.startNodesVector(builder, 0); + int nodes = builder.endVector(); + RecordBatch.startBuffersVector(builder, 0); + int buffers = builder.endVector(); + RecordBatch.startRecordBatch(builder); + RecordBatch.addNodes(builder, nodes); + RecordBatch.addBuffers(builder, buffers); + int header = RecordBatch.endRecordBatch(builder); + IpcOption option = new IpcOption(false, version); + ByteBuffer metadata = + MessageSerializer.serializeMessage(builder, headerType, header, 8, option); + + ByteArrayOutputStream out = new ByteArrayOutputStream(); + WriteChannel channel = new WriteChannel(Channels.newChannel(out)); + MessageSerializer.writeMessageBuffer(channel, metadata.remaining(), metadata, option); + channel.write(new byte[8]); + return out.toByteArray(); + } + + /** Builds a record batch with one node, one 8 byte buffer and the given variadic count. */ + private static RecordBatch recordBatchWithVariadicBufferCount(long count) { + FlatBufferBuilder builder = new FlatBufferBuilder(); + RecordBatch.startNodesVector(builder, 1); + FieldNode.createFieldNode(builder, 0, 0); + int nodes = builder.endVector(); + RecordBatch.startBuffersVector(builder, 1); + Buffer.createBuffer(builder, 0, 8); + int buffers = builder.endVector(); + int variadicBufferCounts = + RecordBatch.createVariadicBufferCountsVector(builder, new long[] {count}); + RecordBatch.startRecordBatch(builder); + RecordBatch.addNodes(builder, nodes); + RecordBatch.addBuffers(builder, buffers); + RecordBatch.addVariadicBufferCounts(builder, variadicBufferCounts); + builder.finish(RecordBatch.endRecordBatch(builder)); + return RecordBatch.getRootAsRecordBatch(builder.dataBuffer()); + } + + private static RecordBatch recordBatchWithNode(long length, long nullCount) { + return recordBatch(0, length, nullCount); + } + + /** Builds a record batch with a single node and, if given, a single buffer. */ + private static RecordBatch recordBatch( + long batchLength, long nodeLength, long nullCount, long... bufferOffsetAndLength) { + FlatBufferBuilder builder = new FlatBufferBuilder(); + RecordBatch.startNodesVector(builder, 1); + FieldNode.createFieldNode(builder, nodeLength, nullCount); + int nodes = builder.endVector(); + if (bufferOffsetAndLength.length == 0) { + RecordBatch.startBuffersVector(builder, 0); + } else { + RecordBatch.startBuffersVector(builder, 1); + Buffer.createBuffer(builder, bufferOffsetAndLength[0], bufferOffsetAndLength[1]); + } + int buffers = builder.endVector(); + RecordBatch.startRecordBatch(builder); + RecordBatch.addLength(builder, batchLength); + RecordBatch.addNodes(builder, nodes); + RecordBatch.addBuffers(builder, buffers); + builder.finish(RecordBatch.endRecordBatch(builder)); + return RecordBatch.getRootAsRecordBatch(builder.dataBuffer()); + } + public static Schema testSchema() { return new Schema( asList( diff --git a/vector/src/test/java/org/apache/arrow/vector/ipc/TestJSONFile.java b/vector/src/test/java/org/apache/arrow/vector/ipc/TestJSONFile.java index b394a667d1..0a6fade9a1 100644 --- a/vector/src/test/java/org/apache/arrow/vector/ipc/TestJSONFile.java +++ b/vector/src/test/java/org/apache/arrow/vector/ipc/TestJSONFile.java @@ -19,6 +19,7 @@ import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertNotNull; import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertThrows; import java.io.File; import java.io.IOException; @@ -43,11 +44,22 @@ import org.apache.arrow.vector.types.pojo.Schema; import org.apache.arrow.vector.util.Validator; import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.io.TempDir; import org.slf4j.Logger; import org.slf4j.LoggerFactory; public class TestJSONFile extends BaseFileTest { private static final Logger LOGGER = LoggerFactory.getLogger(TestJSONFile.class); + private static final String NEGATIVE_BATCH_COUNT_JSON = + "{\"schema\": {\"fields\": [{\"name\": \"i\", \"type\": {\"name\": \"int\"," + + " \"isSigned\": true, \"bitWidth\": 32}, \"nullable\": true, \"children\": []}]}," + + " \"batches\": [{\"count\": -100, \"columns\": [{\"name\": \"i\", \"count\": 0," + + " \"VALIDITY\": [], \"DATA\": []}]}]}"; + private static final String INT_FIELD_JSON = + "{\"name\": \"NAME\", \"type\": {\"name\": \"int\", \"isSigned\": true," + + " \"bitWidth\": 32}, \"nullable\": true, \"children\": []}"; + + @TempDir Path tempDir; @Test public void testNoBatches() throws IOException { @@ -521,4 +533,127 @@ public void testRoundtripEmptyVector() throws Exception { } } } + + @Test + public void testNegativeColumnCountRejected() throws IOException { + File file = + writeJson( + "{\"schema\": {\"fields\": [{\"name\": \"s\", \"type\": {\"name\": \"struct\"}," + + " \"nullable\": true, \"children\": []}]}," + + " \"batches\": [{\"count\": 0, \"columns\": [{\"name\": \"s\", \"count\": -100," + + " \"VALIDITY\": []}]}]}"); + // Small limit, so that sizing buffers from a negative count fails fast + try (BufferAllocator readerAllocator = allocator.newChildAllocator("reader", 0, 1 << 20); + JsonFileReader reader = new JsonFileReader(file, readerAllocator)) { + reader.start(); + assertThrows(IllegalArgumentException.class, () -> reader.read()); + } + } + + @Test + public void testNegativeBatchCountRejected() throws IOException { + File file = writeJson(NEGATIVE_BATCH_COUNT_JSON); + try (JsonFileReader reader = new JsonFileReader(file, allocator)) { + reader.start(); + assertThrows(IllegalArgumentException.class, () -> reader.read()); + } + } + + @Test + public void testNegativeBatchCountRejectedWhenReadingIntoRoot() throws IOException { + File file = writeJson(NEGATIVE_BATCH_COUNT_JSON); + try (JsonFileReader reader = new JsonFileReader(file, allocator)) { + Schema schema = reader.start(); + try (VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator)) { + assertThrows(IllegalArgumentException.class, () -> reader.read(root)); + } + } + } + + @Test + public void testNegativeCountInLaterColumnDoesNotLeak() throws IOException { + File file = + writeJson( + "{\"schema\": {\"fields\": [" + + INT_FIELD_JSON.replace("NAME", "a") + + ", " + + INT_FIELD_JSON.replace("NAME", "b") + + "]}, \"batches\": [{\"count\": 1, \"columns\": [" + + "{\"name\": \"a\", \"count\": 1, \"VALIDITY\": [1], \"DATA\": [5]}," + + " {\"name\": \"b\", \"count\": -1, \"VALIDITY\": [], \"DATA\": []}]}]}"); + assertReadRejectedWithoutLeak(file); + } + + @Test + public void testNegativeChildCountDoesNotLeak() throws IOException { + File file = + writeJson( + "{\"schema\": {\"fields\": [{\"name\": \"s\", \"type\": {\"name\": \"struct\"}," + + " \"nullable\": true, \"children\": [" + + INT_FIELD_JSON.replace("NAME", "c") + + "]}]}, \"batches\": [{\"count\": 1, \"columns\": [" + + "{\"name\": \"s\", \"count\": 1, \"VALIDITY\": [1], \"children\": [" + + "{\"name\": \"c\", \"count\": -1, \"VALIDITY\": [], \"DATA\": []}]}]}]}"); + assertReadRejectedWithoutLeak(file); + } + + @Test + public void testNullColumnCountRejected() throws IOException { + File file = + writeJson( + "{\"schema\": {\"fields\": [" + + INT_FIELD_JSON.replace("NAME", "a") + + "]}, \"batches\": [{\"count\": 1, \"columns\": [" + + "{\"name\": \"a\", \"count\": null, \"VALIDITY\": [1], \"DATA\": [5]}]}]}"); + assertReadRejectedWithoutLeak(file); + } + + @Test + public void testExtraBufferValueDoesNotLeak() throws IOException { + File file = + writeJson( + "{\"schema\": {\"fields\": [" + + INT_FIELD_JSON.replace("NAME", "a") + + "]}, \"batches\": [{\"count\": 1, \"columns\": [" + + "{\"name\": \"a\", \"count\": 1, \"VALIDITY\": [1, 1], \"DATA\": [5]}]}]}"); + assertReadRejectedWithoutLeak(file, IllegalStateException.class); + } + + @Test + public void testInvalidBufferValueDoesNotLeak() throws IOException { + File file = + writeJson( + "{\"schema\": {\"fields\": [" + + INT_FIELD_JSON.replace("NAME", "a") + + "]}, \"batches\": [{\"count\": 1, \"columns\": [" + + "{\"name\": \"a\", \"count\": 1, \"VALIDITY\": [1], \"DATA\": [\"x\"]}]}]}"); + assertReadRejectedWithoutLeak(file, IOException.class); + } + + /** Checks that both read methods reject the file and release everything they allocated. */ + private void assertReadRejectedWithoutLeak(File file) throws IOException { + assertReadRejectedWithoutLeak(file, IllegalArgumentException.class); + } + + private void assertReadRejectedWithoutLeak(File file, Class expectedType) + throws IOException { + try (BufferAllocator readerAllocator = allocator.newChildAllocator("reader", 0, 1 << 20); + JsonFileReader reader = new JsonFileReader(file, readerAllocator)) { + reader.start(); + assertThrows(expectedType, () -> reader.read()); + } + try (BufferAllocator readerAllocator = allocator.newChildAllocator("reader", 0, 1 << 20); + JsonFileReader reader = new JsonFileReader(file, readerAllocator)) { + Schema schema = reader.start(); + try (VectorSchemaRoot root = VectorSchemaRoot.create(schema, readerAllocator)) { + assertThrows(expectedType, () -> reader.read(root)); + } + } + } + + private File writeJson(String json) throws IOException { + Path path = Files.createTempFile(tempDir, "arrow-", ".json"); + Files.writeString(path, json); + return path.toFile(); + } } diff --git a/vector/src/test/java/org/apache/arrow/vector/ipc/message/TestArrowFieldNode.java b/vector/src/test/java/org/apache/arrow/vector/ipc/message/TestArrowFieldNode.java new file mode 100644 index 0000000000..025217f03d --- /dev/null +++ b/vector/src/test/java/org/apache/arrow/vector/ipc/message/TestArrowFieldNode.java @@ -0,0 +1,65 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.arrow.vector.ipc.message; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; + +import org.junit.jupiter.api.Test; + +public class TestArrowFieldNode { + + @Test + public void testNegativeLengthRejected() { + IllegalArgumentException e = + assertThrows(IllegalArgumentException.class, () -> new ArrowFieldNode(-100, 0)); + assertEquals("Expected length >= 0, got -100", e.getMessage()); + } + + @Test + public void testLargeNegativeLengthRejected() { + assertThrows(IllegalArgumentException.class, () -> new ArrowFieldNode(-5_000_000_000L, 0)); + } + + @Test + public void testLengthAboveIntRangeRejected() { + assertThrows(ArithmeticException.class, () -> new ArrowFieldNode((1L << 32) - 100, 0)); + } + + @Test + public void testZeroAndPositiveLengthAccepted() { + assertEquals(0, new ArrowFieldNode(0, 0).getLength()); + assertEquals(42, new ArrowFieldNode(42, 0).getLength()); + } + + @Test + public void testNegativeNullCountSentinelPreserved() { + ArrowFieldNode node = new ArrowFieldNode(10, -1); + assertEquals(10, node.getLength()); + assertEquals(-1, node.getNullCount()); + } + + @Test + public void testNullCountBelowSentinelRejected() { + assertThrows(IllegalArgumentException.class, () -> new ArrowFieldNode(10, -2)); + } + + @Test + public void testNullCountAboveIntRangeRejected() { + assertThrows(ArithmeticException.class, () -> new ArrowFieldNode(10, (1L << 32) - 2)); + } +}