diff --git a/vector/src/main/java/org/apache/arrow/vector/ipc/JsonFileReader.java b/vector/src/main/java/org/apache/arrow/vector/ipc/JsonFileReader.java index e4bab7eb80..9f8639f1ed 100644 --- a/vector/src/main/java/org/apache/arrow/vector/ipc/JsonFileReader.java +++ b/vector/src/main/java/org/apache/arrow/vector/ipc/JsonFileReader.java @@ -196,7 +196,7 @@ public boolean read(VectorSchemaRoot root) throws IOException { JsonToken t = parser.nextToken(); if (t == START_OBJECT) { { - int count = readNextField("count", Integer.class); + int count = readCount(); nextFieldIs("columns"); readToken(START_ARRAY); { @@ -224,7 +224,7 @@ public VectorSchemaRoot read() throws IOException { if (t == START_OBJECT) { VectorSchemaRoot recordBatch = VectorSchemaRoot.create(schema, allocator); { - int count = readNextField("count", Integer.class); + int count = readCount(); recordBatch.setRowCount(count); nextFieldIs("columns"); readToken(START_ARRAY); @@ -885,7 +885,7 @@ private void readFromJsonIntoVector(Field field, FieldVector vector) throws IOEx /* Initialize the vector with required capacity but don't allocateNew since we would * be doing loadFieldBuffers. */ - int valueCount = readNextField("count", Integer.class); + int valueCount = readCount(); vector.setInitialCapacity(valueCount); @@ -975,6 +975,14 @@ private T readNextField(String expectedFieldName, Class c) return parser.readValueAs(c); } + private int readCount() throws IOException { + int count = readNextField("count", Integer.class); + if (count < 0) { + throw new IllegalArgumentException("Expected count >= 0, got " + count); + } + return count; + } + private void nextFieldIs(String expectedFieldName) throws IOException, JsonParseException { String name = parser.nextFieldName(); if (name == null || !name.equals(expectedFieldName)) { diff --git a/vector/src/main/java/org/apache/arrow/vector/ipc/message/ArrowFieldNode.java b/vector/src/main/java/org/apache/arrow/vector/ipc/message/ArrowFieldNode.java index ad3f434e12..9f8f931247 100644 --- a/vector/src/main/java/org/apache/arrow/vector/ipc/message/ArrowFieldNode.java +++ b/vector/src/main/java/org/apache/arrow/vector/ipc/message/ArrowFieldNode.java @@ -16,10 +16,9 @@ */ package org.apache.arrow.vector.ipc.message; -import static org.apache.arrow.memory.util.LargeMemoryUtil.checkedCastToInt; - import com.google.flatbuffers.FlatBufferBuilder; import org.apache.arrow.flatbuf.FieldNode; +import org.apache.arrow.util.Preconditions; /** Metadata about Vectors/Arrays that is written to a channel. */ public class ArrowFieldNode implements FBSerializable { @@ -31,12 +30,14 @@ public class ArrowFieldNode implements FBSerializable { * Constructs a new instance. * * @param length The number of values written. - * @param nullCount The number of null values. + * @param nullCount The number of null values, or -1 if unknown. */ public ArrowFieldNode(long length, long nullCount) { super(); - this.length = checkedCastToInt(length); - this.nullCount = checkedCastToInt(nullCount); + Preconditions.checkArgument(length >= 0, "Expected length >= 0, got %s", length); + Preconditions.checkArgument(nullCount >= -1, "Expected nullCount >= -1, got %s", nullCount); + this.length = Math.toIntExact(length); + this.nullCount = Math.toIntExact(nullCount); } @Override diff --git a/vector/src/test/java/org/apache/arrow/vector/ipc/TestJSONFile.java b/vector/src/test/java/org/apache/arrow/vector/ipc/TestJSONFile.java index b394a667d1..a7f497a5b4 100644 --- a/vector/src/test/java/org/apache/arrow/vector/ipc/TestJSONFile.java +++ b/vector/src/test/java/org/apache/arrow/vector/ipc/TestJSONFile.java @@ -19,6 +19,7 @@ import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertNotNull; import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertThrows; import java.io.File; import java.io.IOException; @@ -48,6 +49,11 @@ public class TestJSONFile extends BaseFileTest { private static final Logger LOGGER = LoggerFactory.getLogger(TestJSONFile.class); + private static final String NEGATIVE_BATCH_COUNT_JSON = + "{\"schema\": {\"fields\": [{\"name\": \"i\", \"type\": {\"name\": \"int\"," + + " \"isSigned\": true, \"bitWidth\": 32}, \"nullable\": true, \"children\": []}]}," + + " \"batches\": [{\"count\": -100, \"columns\": [{\"name\": \"i\", \"count\": 0," + + " \"VALIDITY\": [], \"DATA\": []}]}]}"; @Test public void testNoBatches() throws IOException { @@ -521,4 +527,47 @@ public void testRoundtripEmptyVector() throws Exception { } } } + + @Test + public void testNegativeColumnCountRejected() throws IOException { + File file = + writeJson( + "{\"schema\": {\"fields\": [{\"name\": \"s\", \"type\": {\"name\": \"struct\"}," + + " \"nullable\": true, \"children\": []}]}," + + " \"batches\": [{\"count\": 0, \"columns\": [{\"name\": \"s\", \"count\": -100," + + " \"VALIDITY\": []}]}]}"); + // Small limit, so that sizing buffers from a negative count fails fast + try (BufferAllocator readerAllocator = allocator.newChildAllocator("reader", 0, 1 << 20); + JsonFileReader reader = new JsonFileReader(file, readerAllocator)) { + reader.start(); + assertThrows(IllegalArgumentException.class, () -> reader.read()); + } + } + + @Test + public void testNegativeBatchCountRejected() throws IOException { + File file = writeJson(NEGATIVE_BATCH_COUNT_JSON); + try (JsonFileReader reader = new JsonFileReader(file, allocator)) { + reader.start(); + assertThrows(IllegalArgumentException.class, () -> reader.read()); + } + } + + @Test + public void testNegativeBatchCountRejectedWhenReadingIntoRoot() throws IOException { + File file = writeJson(NEGATIVE_BATCH_COUNT_JSON); + try (JsonFileReader reader = new JsonFileReader(file, allocator)) { + Schema schema = reader.start(); + try (VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator)) { + assertThrows(IllegalArgumentException.class, () -> reader.read(root)); + } + } + } + + private static File writeJson(String json) throws IOException { + Path path = Files.createTempFile("arrow-", ".json"); + path.toFile().deleteOnExit(); + Files.writeString(path, json); + return path.toFile(); + } } diff --git a/vector/src/test/java/org/apache/arrow/vector/ipc/message/TestArrowFieldNode.java b/vector/src/test/java/org/apache/arrow/vector/ipc/message/TestArrowFieldNode.java new file mode 100644 index 0000000000..025217f03d --- /dev/null +++ b/vector/src/test/java/org/apache/arrow/vector/ipc/message/TestArrowFieldNode.java @@ -0,0 +1,65 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.arrow.vector.ipc.message; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; + +import org.junit.jupiter.api.Test; + +public class TestArrowFieldNode { + + @Test + public void testNegativeLengthRejected() { + IllegalArgumentException e = + assertThrows(IllegalArgumentException.class, () -> new ArrowFieldNode(-100, 0)); + assertEquals("Expected length >= 0, got -100", e.getMessage()); + } + + @Test + public void testLargeNegativeLengthRejected() { + assertThrows(IllegalArgumentException.class, () -> new ArrowFieldNode(-5_000_000_000L, 0)); + } + + @Test + public void testLengthAboveIntRangeRejected() { + assertThrows(ArithmeticException.class, () -> new ArrowFieldNode((1L << 32) - 100, 0)); + } + + @Test + public void testZeroAndPositiveLengthAccepted() { + assertEquals(0, new ArrowFieldNode(0, 0).getLength()); + assertEquals(42, new ArrowFieldNode(42, 0).getLength()); + } + + @Test + public void testNegativeNullCountSentinelPreserved() { + ArrowFieldNode node = new ArrowFieldNode(10, -1); + assertEquals(10, node.getLength()); + assertEquals(-1, node.getNullCount()); + } + + @Test + public void testNullCountBelowSentinelRejected() { + assertThrows(IllegalArgumentException.class, () -> new ArrowFieldNode(10, -2)); + } + + @Test + public void testNullCountAboveIntRangeRejected() { + assertThrows(ArithmeticException.class, () -> new ArrowFieldNode(10, (1L << 32) - 2)); + } +}