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..49633f236a 100644 --- a/c/src/main/java/org/apache/arrow/c/ArrayImporter.java +++ b/c/src/main/java/org/apache/arrow/c/ArrayImporter.java @@ -128,19 +128,21 @@ private void doImport(ArrowArray.Snapshot snapshot) { } // 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); + try { + 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); } - vector.loadFieldBuffers(fieldNode, buffers); } catch (Exception e) { throw new IllegalArgumentException( "Could not load buffers for field " 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..37acaf521e 100644 --- a/c/src/main/java/org/apache/arrow/c/Data.java +++ b/c/src/main/java/org/apache/arrow/c/Data.java @@ -407,7 +407,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 (Exception e) { + vector.close(); + throw e; + } return vector; } 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..c51c7e90c9 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,74 @@ 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); + }); + } + } + + 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..9f8639f1ed 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 @@ -196,7 +196,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); { @@ -224,7 +224,7 @@ public VectorSchemaRoot read() throws IOException { if (t == START_OBJECT) { VectorSchemaRoot recordBatch = VectorSchemaRoot.create(schema, allocator); { - int count = readNextField("count", Integer.class); + int count = readCount(); recordBatch.setRowCount(count); nextFieldIs("columns"); readToken(START_ARRAY); @@ -885,7 +885,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); @@ -975,6 +975,14 @@ private T readNextField(String expectedFieldName, Class c) return parser.readValueAs(c); } + private int readCount() throws IOException { + int count = readNextField("count", Integer.class); + if (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..4bbcbe8a83 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 @@ -411,50 +411,58 @@ 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); + ArrowBuf vectorBuffer = body.slice(bufferFB.offset(), bufferFB.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()); + 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)); - } + List variadicBufferCounts = new ArrayList<>(); + for (int i = 0; i < recordBatchFB.variadicBufferCountsLength(); i++) { + variadicBufferCounts.add(recordBatchFB.variadicBufferCounts(i)); + } - if ((int) recordBatchFB.length() != recordBatchFB.length()) { - throw new IOException( - "Cannot currently deserialize record batches with more than INT_MAX records."); + 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; } /** 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..3eee24b7c1 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,6 +33,9 @@ import java.util.Arrays; import java.util.Collections; import java.util.List; +import java.util.stream.Stream; +import org.apache.arrow.flatbuf.FieldNode; +import org.apache.arrow.flatbuf.RecordBatch; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; @@ -47,6 +51,9 @@ 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; public class MessageSerializerTest { @@ -220,6 +227,62 @@ 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)); + } + + private static RecordBatch recordBatchWithNode(long length, long nullCount) { + FlatBufferBuilder builder = new FlatBufferBuilder(); + RecordBatch.startNodesVector(builder, 1); + FieldNode.createFieldNode(builder, length, nullCount); + int nodes = builder.endVector(); + RecordBatch.startBuffersVector(builder, 0); + int buffers = builder.endVector(); + RecordBatch.startRecordBatch(builder); + RecordBatch.addLength(builder, 0); + 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..a7f497a5b4 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; @@ -48,6 +49,11 @@ 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\": []}]}]}"; @Test public void testNoBatches() throws IOException { @@ -521,4 +527,47 @@ 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)); + } + } + } + + private static File writeJson(String json) throws IOException { + Path path = Files.createTempFile("arrow-", ".json"); + path.toFile().deleteOnExit(); + 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)); + } +}