From ad2f9b9e77e07038d3dc342833d5849b26270499 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Sat, 31 Oct 2026 09:08:16 -0600 Subject: [PATCH 1/4] 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/4] 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/4] 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/4] 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);