From ad2f9b9e77e07038d3dc342833d5849b26270499 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Sat, 31 Oct 2026 09:08:16 -0600 Subject: [PATCH 1/9] fix(vector): reject negative FieldNode length during deserialization A negative ArrowFieldNode length flows unchecked into buffer-size math (getValidityBufferSizeFromCount uses an unsigned shift) and into ArrowBuf.setOne/setZero, risking an oversized allocation and, with bounds checking disabled, memory corruption via Unsafe.setMemory. The IPC, JSON, and C Data deserialization paths all construct ArrowFieldNode from external metadata, so validate length at that single chokepoint. Closes #1421 --- .../vector/ipc/message/ArrowFieldNode.java | 8 +++ .../ipc/message/TestArrowFieldNode.java | 55 +++++++++++++++++++ 2 files changed, 63 insertions(+) create mode 100644 vector/src/test/java/org/apache/arrow/vector/ipc/message/TestArrowFieldNode.java 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..6a1ffff117 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 @@ -20,6 +20,7 @@ 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 { @@ -35,6 +36,13 @@ public class ArrowFieldNode implements FBSerializable { */ public ArrowFieldNode(long length, long nullCount) { super(); + // A negative length is never valid and, if accepted, flows unchecked into buffer size + // computations (BitVectorHelper.getValidityBufferSizeFromCount performs an unsigned shift) + // and into ArrowBuf.setOne/setZero, which can corrupt memory when bounds checking is + // disabled. Reject it at this single chokepoint shared by every deserialization path + // (IPC, JSON, C Data). nullCount is intentionally not checked here: VectorUnloader uses -1 + // as a sentinel meaning "null count not included". + Preconditions.checkArgument(length >= 0, "Expected length >= 0, got %s", length); this.length = checkedCastToInt(length); this.nullCount = checkedCastToInt(nullCount); } 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..9071604a65 --- /dev/null +++ b/vector/src/test/java/org/apache/arrow/vector/ipc/message/TestArrowFieldNode.java @@ -0,0 +1,55 @@ +/* + * 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() { + // A negative length that still fits in an int must be rejected. Left unchecked it flows into + // BitVectorHelper.getValidityBufferSizeFromCount (unsigned shift -> ~512MB) and into + // ArrowBuf.setOne(0, length / 8) with a negative length, corrupting memory when bounds + // checking is disabled. See the deserialization paths in MessageSerializer and JsonFileReader. + 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 testZeroAndPositiveLengthAccepted() { + assertEquals(0, new ArrowFieldNode(0, 0).getLength()); + assertEquals(42, new ArrowFieldNode(42, 0).getLength()); + } + + @Test + public void testNegativeNullCountSentinelPreserved() { + // VectorUnloader uses nullCount == -1 to signal "null count not included"; it must survive. + ArrowFieldNode node = new ArrowFieldNode(10, -1); + assertEquals(10, node.getLength()); + assertEquals(-1, node.getNullCount()); + } +} From e593bb3d8df96ab9a63b1e98ce44cd83a12f02de Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Sun, 1 Nov 2026 07:46:31 -0700 Subject: [PATCH 2/9] fix(vector): reject truncated FieldNode values and negative JSON counts ArrowFieldNode narrowed its arguments with checkedCastToInt, a plain cast when bounds checking is disabled, so a C Data length such as 2^32 - 100 passed the length check and was stored as -100. Narrow with Math.toIntExact unconditionally, and reject a null count below the -1 "unknown" sentinel. JsonFileReader sized a column's buffers from its count before building the ArrowFieldNode, so a negative count still allocated about 512 MB, and batch counts went straight into setRowCount. Reject a negative count as soon as it is read. --- .../arrow/vector/ipc/JsonFileReader.java | 14 ++++-- .../vector/ipc/message/ArrowFieldNode.java | 15 ++---- .../apache/arrow/vector/ipc/TestJSONFile.java | 49 +++++++++++++++++++ .../ipc/message/TestArrowFieldNode.java | 20 ++++++-- 4 files changed, 79 insertions(+), 19 deletions(-) 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 6a1ffff117..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,8 +16,6 @@ */ 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; @@ -32,19 +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(); - // A negative length is never valid and, if accepted, flows unchecked into buffer size - // computations (BitVectorHelper.getValidityBufferSizeFromCount performs an unsigned shift) - // and into ArrowBuf.setOne/setZero, which can corrupt memory when bounds checking is - // disabled. Reject it at this single chokepoint shared by every deserialization path - // (IPC, JSON, C Data). nullCount is intentionally not checked here: VectorUnloader uses -1 - // as a sentinel meaning "null count not included". Preconditions.checkArgument(length >= 0, "Expected length >= 0, got %s", length); - this.length = checkedCastToInt(length); - this.nullCount = checkedCastToInt(nullCount); + 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/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 index 9071604a65..025217f03d 100644 --- 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 @@ -25,10 +25,6 @@ public class TestArrowFieldNode { @Test public void testNegativeLengthRejected() { - // A negative length that still fits in an int must be rejected. Left unchecked it flows into - // BitVectorHelper.getValidityBufferSizeFromCount (unsigned shift -> ~512MB) and into - // ArrowBuf.setOne(0, length / 8) with a negative length, corrupting memory when bounds - // checking is disabled. See the deserialization paths in MessageSerializer and JsonFileReader. IllegalArgumentException e = assertThrows(IllegalArgumentException.class, () -> new ArrowFieldNode(-100, 0)); assertEquals("Expected length >= 0, got -100", e.getMessage()); @@ -39,6 +35,11 @@ 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()); @@ -47,9 +48,18 @@ public void testZeroAndPositiveLengthAccepted() { @Test public void testNegativeNullCountSentinelPreserved() { - // VectorUnloader uses nullCount == -1 to signal "null count not included"; it must survive. 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)); + } } From 23609d6b94f3fb7af7debf63012702a6bedac710 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Mon, 2 Nov 2026 14:28:27 -0700 Subject: [PATCH 3/9] fix(vector): release the record batch body when deserialization fails MessageSerializer.deserializeRecordBatch(RecordBatch, ArrowBuf) released the body only after building the ArrowRecordBatch, so an exception before that point leaked it. The negative-length check in ArrowFieldNode made a negative node length one such case; node and batch lengths above INT_MAX already were. ArrowStreamReader, ArrowFileReader, dictionary batches, and Flight's ArrowMessage hand the body to this method and do not release it themselves. Release the body in a finally block. ArrowRecordBatch retains the slices it keeps, so the release is also correct when deserialization succeeds. --- .../vector/ipc/message/MessageSerializer.java | 79 ++++++++++--------- .../vector/ipc/MessageSerializerTest.java | 41 ++++++++++ 2 files changed, 81 insertions(+), 39 deletions(-) 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..e8e8f61c11 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,51 @@ 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 ((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..11057d3f59 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,8 @@ import java.util.Arrays; import java.util.Collections; import java.util.List; +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; @@ -220,6 +223,44 @@ public void testSerializeRecordBatchV5() throws Exception { alloc.close(); } + @Test + public void testNegativeNodeLengthDoesNotLeakBody() { + RecordBatch recordBatchFB = recordBatchWithNode(-100, 0); + try (BufferAllocator allocator = new RootAllocator()) { + ArrowBuf body = allocator.buffer(8); + assertThrows( + IllegalArgumentException.class, + () -> MessageSerializer.deserializeRecordBatch(recordBatchFB, body)); + assertEquals(0, allocator.getAllocatedMemory()); + } + } + + @Test + public void testNodeLengthAboveIntRangeDoesNotLeakBody() { + RecordBatch recordBatchFB = recordBatchWithNode(1L << 32, 0); + try (BufferAllocator allocator = new RootAllocator()) { + ArrowBuf body = allocator.buffer(8); + assertThrows( + IOException.class, () -> MessageSerializer.deserializeRecordBatch(recordBatchFB, body)); + assertEquals(0, allocator.getAllocatedMemory()); + } + } + + 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( From 6f2c26eda16181c12c200e3bc78193fb6ec88545 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Tue, 3 Nov 2026 13:23:11 -0700 Subject: [PATCH 4/9] fix(vector): throw IOException for invalid FieldNode values in IPC reads deserializeRecordBatch left a negative length or a null count below -1 to the ArrowFieldNode constructor, which throws IllegalArgumentException. Check both in the loop so the IPC readers report every malformed node with the IOException they declare, and test the boundary values through deserializeRecordBatch. --- .../vector/ipc/message/MessageSerializer.java | 7 ++++ .../vector/ipc/MessageSerializerTest.java | 42 ++++++++++++++----- 2 files changed, 39 insertions(+), 10 deletions(-) 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 e8e8f61c11..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 @@ -417,6 +417,13 @@ public static ArrowRecordBatch deserializeRecordBatch(RecordBatch recordBatchFB, 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 " 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 11057d3f59..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 @@ -33,6 +33,7 @@ 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; @@ -50,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 { @@ -223,29 +227,47 @@ public void testSerializeRecordBatchV5() throws Exception { alloc.close(); } - @Test - public void testNegativeNodeLengthDoesNotLeakBody() { - RecordBatch recordBatchFB = recordBatchWithNode(-100, 0); + @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( - IllegalArgumentException.class, - () -> MessageSerializer.deserializeRecordBatch(recordBatchFB, body)); + IOException.class, () -> MessageSerializer.deserializeRecordBatch(recordBatchFB, body)); assertEquals(0, allocator.getAllocatedMemory()); } } - @Test - public void testNodeLengthAboveIntRangeDoesNotLeakBody() { - RecordBatch recordBatchFB = recordBatchWithNode(1L << 32, 0); + 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); - assertThrows( - IOException.class, () -> MessageSerializer.deserializeRecordBatch(recordBatchFB, body)); + 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); From d5bc95194b1cbb61b5ef67d8d4bb8176d38a4ee5 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Mon, 9 Nov 2026 08:26:03 -0700 Subject: [PATCH 5/9] fix(c): name the field and avoid leaks on invalid import lengths Build the ArrowFieldNode inside the try block in ArrayImporter so an invalid length or null count reports the field it belongs to, and close the vector in Data.importVector when the import fails so children that were already loaded are released. --- .../org/apache/arrow/c/ArrayImporter.java | 26 +++---- c/src/main/java/org/apache/arrow/c/Data.java | 7 +- .../org/apache/arrow/c/RoundtripTest.java | 70 +++++++++++++++++++ 3 files changed, 90 insertions(+), 13 deletions(-) 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); From 613aba636317bccefc7157e5c9ba1f983dce023a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?JB=20Onofr=C3=A9?= Date: Thu, 12 Nov 2026 10:34:21 +0100 Subject: [PATCH 6/9] fix(vector): reject invalid IPC batch lengths and buffers, fix JSON reader leaks - Reject a negative record batch length, which the INT_MAX check let through. - Reject buffers with a negative offset or length, or that extend past the body, with an IOException. A negative length passed ArrowBuf.slice and made the ArrowRecordBatch constructor throw after retaining the body, leaking it. - Use an empty body for messages without one in deserializeMessageBatch, as ArrowStreamReader does, instead of failing with a NullPointerException. - Close the record batch in JsonFileReader.read() and release the local buffers in readFromJsonIntoVector when reading a column fails. - Detect overflow of the offset buffer count in JsonFileReader. --- .../arrow/vector/ipc/JsonFileReader.java | 20 +++--- .../vector/ipc/message/MessageSerializer.java | 29 +++++++- .../vector/ipc/MessageSerializerTest.java | 70 ++++++++++++++++++- .../apache/arrow/vector/ipc/TestJSONFile.java | 54 +++++++++++++- 4 files changed, 156 insertions(+), 17 deletions(-) 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 9f8639f1ed..3d74fe9688 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; @@ -223,7 +224,7 @@ public VectorSchemaRoot read() throws IOException { JsonToken t = parser.nextToken(); if (t == START_OBJECT) { VectorSchemaRoot recordBatch = VectorSchemaRoot.create(schema, allocator); - { + try { int count = readCount(); recordBatch.setRowCount(count); nextFieldIs("columns"); @@ -235,8 +236,11 @@ public VectorSchemaRoot read() throws IOException { } } readToken(END_ARRAY); + readToken(END_OBJECT); + } catch (Throwable e) { + AutoCloseables.close(e, recordBatch); + throw e; } - readToken(END_OBJECT); return recordBatch; } else if (t == END_ARRAY) { return null; @@ -874,7 +878,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)) { @@ -898,7 +902,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 +945,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 { 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 4bbcbe8a83..137f62e329 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 @@ -434,7 +434,21 @@ public static ArrowRecordBatch deserializeRecordBatch(RecordBatch recordBatchFB, 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()); + 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); } @@ -449,6 +463,9 @@ public static ArrowRecordBatch deserializeRecordBatch(RecordBatch recordBatchFB, variadicBufferCounts.add(recordBatchFB.variadicBufferCounts(i)); } + 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."); @@ -639,11 +656,17 @@ public static ArrowMessage deserializeMessageBatch(MessageChannelReader reader) + 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()); + return deserializeRecordBatch(result.getMessage(), bodyBuffer); case MessageHeader.DictionaryBatch: - return deserializeDictionaryBatch(result.getMessage(), result.getBodyBuffer()); + return deserializeDictionaryBatch(result.getMessage(), bodyBuffer); default: throw new IOException("Unexpected message header type " + result.getMessage().headerType()); } 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 3eee24b7c1..26d18b1926 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 @@ -34,6 +34,7 @@ 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.RecordBatch; import org.apache.arrow.memory.ArrowBuf; @@ -268,15 +269,78 @@ static Stream validFieldNodes() { 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()); + } + } + 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, length, nullCount); + FieldNode.createFieldNode(builder, nodeLength, nullCount); int nodes = builder.endVector(); - RecordBatch.startBuffersVector(builder, 0); + 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, 0); + RecordBatch.addLength(builder, batchLength); RecordBatch.addNodes(builder, nodes); RecordBatch.addBuffers(builder, buffers); builder.finish(RecordBatch.endRecordBatch(builder)); 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 a7f497a5b4..cb7ebbd204 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 @@ -44,6 +44,7 @@ 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; @@ -54,6 +55,11 @@ public class TestJSONFile extends BaseFileTest { + " \"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 { @@ -564,9 +570,51 @@ public void testNegativeBatchCountRejectedWhenReadingIntoRoot() throws IOExcepti } } - private static File writeJson(String json) throws IOException { - Path path = Files.createTempFile("arrow-", ".json"); - path.toFile().deleteOnExit(); + @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); + } + + /** Checks that both read methods reject the file and release everything they allocated. */ + private void assertReadRejectedWithoutLeak(File file) throws IOException { + try (BufferAllocator readerAllocator = allocator.newChildAllocator("reader", 0, 1 << 20); + JsonFileReader reader = new JsonFileReader(file, readerAllocator)) { + reader.start(); + assertThrows(IllegalArgumentException.class, () -> 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(IllegalArgumentException.class, () -> reader.read(root)); + } + } + } + + private File writeJson(String json) throws IOException { + Path path = Files.createTempFile(tempDir, "arrow-", ".json"); Files.writeString(path, json); return path.toFile(); } From d4c0d3f737cfa7c9defe624f9e07f61a6897439e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?JB=20Onofr=C3=A9?= Date: Thu, 12 Nov 2026 10:34:21 +0100 Subject: [PATCH 7/9] fix(c): check the array length before importing children and the dictionary - Build the ArrowFieldNode before importing children and the dictionary, so an array with an invalid length or null count no longer clears and reloads the provider's dictionary before it is rejected. - Use Math.toIntExact for n_children and n_buffers. - Compute the offsets buffer size as a long, so a length of Integer.MAX_VALUE no longer overflows, and reject a negative buffer capacity. - Keep the original exception if closing the vector fails after an import error, and close it on any Throwable. --- .../org/apache/arrow/c/ArrayImporter.java | 20 ++++-- .../arrow/c/BufferImportTypeVisitor.java | 4 +- c/src/main/java/org/apache/arrow/c/Data.java | 7 ++- .../apache/arrow/c/ArrowArrayUtilityTest.java | 14 +++++ .../org/apache/arrow/c/DictionaryTest.java | 62 +++++++++++++++++++ 5 files changed, 99 insertions(+), 8 deletions(-) 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 49633f236a..123fd7d72a 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; @@ -93,9 +92,23 @@ private void doImport(ArrowArray.Snapshot snapshot) { "ArrowArray struct has non-zero offset (%s), which is not supported", snapshot.offset); + // Check the length and null count before importing children or the dictionary, so that an + // invalid array is rejected before it modifies the child vectors or the dictionary provider + final ArrowFieldNode fieldNode; + try { + fieldNode = new ArrowFieldNode(snapshot.length, snapshot.null_count); + } catch (IllegalArgumentException | ArithmeticException e) { + throw new IllegalArgumentException( + "Invalid length or null count for field " + + vector.getField() + + ". error message: " + + e.getMessage(), + e); + } + // First import children (required for reconstituting parent array data) long[] children = - NativeUtil.toJavaArray(snapshot.children, checkedCastToInt(snapshot.n_children)); + NativeUtil.toJavaArray(snapshot.children, Math.toIntExact(snapshot.n_children)); if (children != null && children.length > 0) { List childVectors = vector.getChildrenFromFields(); checkState( @@ -129,9 +142,8 @@ private void doImport(ArrowArray.Snapshot snapshot) { // Import main data try { - ArrowFieldNode fieldNode = new ArrowFieldNode(snapshot.length, snapshot.null_count); long[] bufferPointers = - NativeUtil.toJavaArray(snapshot.buffers, checkedCastToInt(snapshot.n_buffers)); + NativeUtil.toJavaArray(snapshot.buffers, Math.toIntExact(snapshot.n_buffers)); try (final BufferImportTypeVisitor visitor = new BufferImportTypeVisitor(allocator, underlyingAllocation, fieldNode, bufferPointers)) { 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 37acaf521e..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; @@ -409,9 +410,9 @@ public static FieldVector importVector( FieldVector vector = field.createVector(allocator); try { importIntoVector(allocator, array, vector, provider, closeImportedStructs); - } catch (Exception e) { - vector.close(); - throw e; + } 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..856a33134a 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; @@ -133,6 +134,67 @@ public void testWithDictionary() throws Exception { AutoCloseables.close((AutoCloseable) vector, encodedVector, dictVector); } + @Test + public void testImportRejectsInvalidLengthWithoutReplacingDictionary() 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); + } + + // Import a second batch with a different dictionary and an invalid length 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(); + ArrowArray.Snapshot snapshot = arrowArray.snapshot(); + snapshot.length = -1; + arrowArray.save(snapshot); + assertThrows( + IllegalArgumentException.class, + () -> Data.importIntoVector(allocator, arrowArray, imported, cDictionaryProvider)); + } + + assertTrue( + VectorEqualsVisitor.vectorEquals(dictVector, cDictionaryProvider.lookup(0L).getVector()), + "dictionary was modified by the rejected import"); + imported.close(); + } + + 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(); From 11adc4aa963875da9c949d99b0194db21cee0a6a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?JB=20Onofr=C3=A9?= Date: Thu, 12 Nov 2026 11:11:33 +0100 Subject: [PATCH 8/9] fix(c): keep the stored dictionary when an import is rejected Check the dictionary's length and null count together with the array's before importing anything, and import the dictionary after the array's own buffers. An array with an invalid length, null count or missing buffer, or a dictionary with an invalid length or null count, no longer clears the dictionary in the provider. Reword the comment to say only what is checked up front, and release the test vectors when an assertion fails. --- .../org/apache/arrow/c/ArrayImporter.java | 65 ++++++++------- .../org/apache/arrow/c/DictionaryTest.java | 81 ++++++++++++++----- 2 files changed, 97 insertions(+), 49 deletions(-) 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 123fd7d72a..4a955e0e8e 100644 --- a/c/src/main/java/org/apache/arrow/c/ArrayImporter.java +++ b/c/src/main/java/org/apache/arrow/c/ArrayImporter.java @@ -92,18 +92,20 @@ private void doImport(ArrowArray.Snapshot snapshot) { "ArrowArray struct has non-zero offset (%s), which is not supported", snapshot.offset); - // Check the length and null count before importing children or the dictionary, so that an - // invalid array is rejected before it modifies the child vectors or the dictionary provider - final ArrowFieldNode fieldNode; - try { - fieldNode = new ArrowFieldNode(snapshot.length, snapshot.null_count); - } catch (IllegalArgumentException | ArithmeticException e) { - throw new IllegalArgumentException( - "Invalid length or null count for field " - + vector.getField() - + ". error message: " - + e.getMessage(), - e); + // 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) @@ -124,22 +126,6 @@ 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"); - - // reset the dictionary vector to the initial state - dictionary.getVector().clear(); - - ArrayImporter dictionaryImporter = - new ArrayImporter(allocator, dictionary.getVector(), dictionaryProvider); - dictionaryImporter.importChild(this, ArrowArray.wrap(snapshot.dictionary)); - } - // Import main data try { long[] bufferPointers = @@ -163,5 +149,28 @@ private void doImport(ArrowArray.Snapshot snapshot) { + 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(); + + ArrayImporter dictionaryImporter = + new ArrayImporter(allocator, dictionary.getVector(), dictionaryProvider); + dictionaryImporter.importChild(this, ArrowArray.wrap(snapshot.dictionary)); + } + } + + 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( + "Invalid length or null count for field " + + vector.getField() + + ". error message: " + + e.getMessage(), + e); + } } } 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 856a33134a..202fc9f167 100644 --- a/c/src/test/java/org/apache/arrow/c/DictionaryTest.java +++ b/c/src/test/java/org/apache/arrow/c/DictionaryTest.java @@ -27,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; @@ -136,6 +137,42 @@ public void testWithDictionary() throws Exception { @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 = @@ -160,29 +197,31 @@ public void testImportRejectsInvalidLengthWithoutReplacingDictionary() throws Ex imported = Data.importVector(allocator, arrowArray, arrowSchema, cDictionaryProvider); } - // Import a second batch with a different dictionary and an invalid length 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(); - ArrowArray.Snapshot snapshot = arrowArray.snapshot(); - snapshot.length = -1; - arrowArray.save(snapshot); - assertThrows( - IllegalArgumentException.class, - () -> Data.importIntoVector(allocator, arrowArray, imported, 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"); - imported.close(); + 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); } - - AutoCloseables.close( - (AutoCloseable) vector, otherVector, encoded, otherEncoded, dictVector, otherDictVector); } private VarCharVector newVarCharVector(String name, String... values) { From b64073db58e0a2b699cd6694687a43dfd776a3fc Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?JB=20Onofr=C3=A9?= Date: Thu, 12 Nov 2026 11:11:33 +0100 Subject: [PATCH 9/9] fix(vector): release JSON reader buffers when parsing a buffer fails A buffer allocated while reading a JSON buffer leaked if parsing failed before the buffer was returned, for example on an invalid value or an extra value before the closing bracket. Release it on failure, including for view and variadic buffers. Also reject a null count with IllegalArgumentException instead of a NullPointerException. --- .../arrow/vector/ipc/JsonFileReader.java | 99 ++++++++++++------- .../apache/arrow/vector/ipc/TestJSONFile.java | 42 +++++++- 2 files changed, 102 insertions(+), 39 deletions(-) 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 3d74fe9688..8be0142e84 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 @@ -272,13 +272,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; + } } } @@ -294,22 +310,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; } @@ -372,11 +393,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; } @@ -387,7 +412,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); @@ -407,7 +432,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); @@ -425,7 +450,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); @@ -444,7 +469,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(); @@ -460,7 +485,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(); @@ -476,7 +501,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(); @@ -492,7 +517,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(); @@ -509,7 +534,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(); @@ -525,7 +550,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(); @@ -541,7 +566,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(); @@ -557,7 +582,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(); @@ -574,7 +599,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(); @@ -590,7 +615,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(); @@ -606,7 +631,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(); @@ -625,7 +650,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(); @@ -980,8 +1005,8 @@ private T readNextField(String expectedFieldName, Class c) } private int readCount() throws IOException { - int count = readNextField("count", Integer.class); - if (count < 0) { + Integer count = readNextField("count", Integer.class); + if (count == null || count < 0) { throw new IllegalArgumentException("Expected count >= 0, got " + count); } return count; 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 cb7ebbd204..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 @@ -597,18 +597,56 @@ public void testNegativeChildCountDoesNotLeak() throws IOException { 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(IllegalArgumentException.class, () -> reader.read()); + 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(IllegalArgumentException.class, () -> reader.read(root)); + assertThrows(expectedType, () -> reader.read(root)); } } }