GH-1421: reject negative FieldNode length during deserialization

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,25 @@ 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 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)
     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<FieldVector> childVectors = vector.getChildrenFromFields();
       checkState(
@@ -111,39 +126,47 @@ private void doImport(ArrowArray.Snapshot snapshot) {
       }
     }
 
-    // Handle import of a dictionary encoded vector
-    if (snapshot.dictionary != NULL) {
-      DictionaryEncoding encoding = vector.getField().getDictionary();
-      checkNotNull(encoding, "Missing encoding on import of ArrowArray with dictionary");
-
-      Dictionary dictionary = dictionaryProvider.lookup(encoding.getId());
-      checkNotNull(dictionary, "Dictionary lookup failed on import of ArrowArray with dictionary");
+    // Import main data
+    try {
+      long[] bufferPointers =
+          NativeUtil.toJavaArray(snapshot.buffers, Math.toIntExact(snapshot.n_buffers));
+
+      try (final BufferImportTypeVisitor visitor =
+          new BufferImportTypeVisitor(allocator, underlyingAllocation, fieldNode, bufferPointers)) {
+        final List<ArrowBuf> buffers;
+        if (bufferPointers == null || bufferPointers.length == 0) {
+          buffers = Collections.emptyList();
+        } else {
+          buffers = vector.getField().getType().accept(visitor);
+        }
+        vector.loadFieldBuffers(fieldNode, buffers);
+      }
+    } catch (Exception e) {
+      throw new IllegalArgumentException(
+          "Could not load buffers for field "
+              + vector.getField()
+              + ". error message: "
+              + e.getMessage(),
+          e);
+    }
 
+    // Import the dictionary last, so that it is not cleared if this array's own data is invalid
+    if (dictionary != null) {
       // reset the dictionary vector to the initial state
       dictionary.getVector().clear();
 
       ArrayImporter dictionaryImporter =
           new ArrayImporter(allocator, dictionary.getVector(), dictionaryProvider);
       dictionaryImporter.importChild(this, ArrowArray.wrap(snapshot.dictionary));
     }
+  }
 
-    // Import main data
-    ArrowFieldNode fieldNode = new ArrowFieldNode(snapshot.length, snapshot.null_count);
-    long[] bufferPointers =
-        NativeUtil.toJavaArray(snapshot.buffers, checkedCastToInt(snapshot.n_buffers));
-
-    try (final BufferImportTypeVisitor visitor =
-        new BufferImportTypeVisitor(allocator, underlyingAllocation, fieldNode, bufferPointers)) {
-      final List<ArrowBuf> buffers;
-      if (bufferPointers == null || bufferPointers.length == 0) {
-        buffers = Collections.emptyList();
-      } else {
-        buffers = vector.getField().getType().accept(visitor);
-      }
-      vector.loadFieldBuffers(fieldNode, buffers);
-    } catch (Exception e) {
+  private static ArrowFieldNode toFieldNode(FieldVector vector, ArrowArray.Snapshot snapshot) {
+    try {
+      return new ArrowFieldNode(snapshot.length, snapshot.null_count);
+    } catch (IllegalArgumentException | ArithmeticException e) {
       throw new IllegalArgumentException(
-          "Could not load buffers for field "
+          "Invalid length or null count for field "
               + vector.getField()
               + ". error message: "
               + e.getMessage(),

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);
   }
 

c/src/main/java/org/apache/arrow/c/Data.java

@@ -17,6 +17,7 @@
 package org.apache.arrow.c;
 
 import org.apache.arrow.memory.BufferAllocator;
+import org.apache.arrow.util.AutoCloseables;
 import org.apache.arrow.vector.FieldVector;
 import org.apache.arrow.vector.VectorLoader;
 import org.apache.arrow.vector.VectorSchemaRoot;
@@ -407,7 +408,12 @@ public static FieldVector importVector(
       boolean closeImportedStructs) {
     Field field = importField(allocator, schema, provider, closeImportedStructs);
     FieldVector vector = field.createVector(allocator);
-    importIntoVector(allocator, array, vector, provider, closeImportedStructs);
+    try {
+      importIntoVector(allocator, array, vector, provider, closeImportedStructs);
+    } catch (Throwable t) {
+      AutoCloseables.close(t, vector);
+      throw t;
+    }
     return vector;
   }
 

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

c/src/test/java/org/apache/arrow/c/DictionaryTest.java

@@ -18,6 +18,7 @@
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 
 import java.io.ByteArrayInputStream;
@@ -26,6 +27,7 @@
 import java.nio.channels.Channels;
 import java.nio.charset.StandardCharsets;
 import java.util.Collections;
+import java.util.function.Consumer;
 import org.apache.arrow.memory.BufferAllocator;
 import org.apache.arrow.memory.RootAllocator;
 import org.apache.arrow.util.AutoCloseables;
@@ -133,6 +135,105 @@ public void testWithDictionary() throws Exception {
     AutoCloseables.close((AutoCloseable) vector, encodedVector, dictVector);
   }
 
+  @Test
+  public void testImportRejectsInvalidLengthWithoutReplacingDictionary() throws Exception {
+    assertRejectedImportKeepsDictionary(
+        arrowArray -> {
+          ArrowArray.Snapshot snapshot = arrowArray.snapshot();
+          snapshot.length = -1;
+          arrowArray.save(snapshot);
+        });
+  }
+
+  @Test
+  public void testImportRejectsInvalidDictionaryLengthWithoutReplacingDictionary()
+      throws Exception {
+    assertRejectedImportKeepsDictionary(
+        arrowArray -> {
+          ArrowArray dictionaryArray = ArrowArray.wrap(arrowArray.snapshot().dictionary);
+          ArrowArray.Snapshot snapshot = dictionaryArray.snapshot();
+          snapshot.length = -1;
+          dictionaryArray.save(snapshot);
+        });
+  }
+
+  @Test
+  public void testImportRejectsMissingBufferWithoutReplacingDictionary() throws Exception {
+    assertRejectedImportKeepsDictionary(
+        arrowArray -> {
+          ArrowArray.Snapshot snapshot = arrowArray.snapshot();
+          snapshot.n_buffers = 1;
+          arrowArray.save(snapshot);
+        });
+  }
+
+  /**
+   * Imports a dictionary encoded batch, then checks that a second batch with a different
+   * dictionary, made invalid by {@code corrupt}, is rejected without replacing the first
+   * dictionary.
+   */
+  private void assertRejectedImportKeepsDictionary(Consumer<ArrowArray> corrupt) throws Exception {
+    DictionaryProvider.MapDictionaryProvider provider =
+        new DictionaryProvider.MapDictionaryProvider();
+    DictionaryProvider.MapDictionaryProvider otherProvider =
+        new DictionaryProvider.MapDictionaryProvider();
+    VarCharVector dictVector = newVarCharVector("dict", "aa", "bb", "cc");
+    VarCharVector otherDictVector = newVarCharVector("dict", "xx", "yy");
+    VarCharVector vector = newVarCharVector("vector", "bb", "cc");
+    VarCharVector otherVector = newVarCharVector("vector", "yy", "xx");
+    Dictionary dictionary = new Dictionary(dictVector, new DictionaryEncoding(0L, false, null));
+    Dictionary otherDictionary =
+        new Dictionary(otherDictVector, new DictionaryEncoding(0L, false, null));
+    provider.put(dictionary);
+    otherProvider.put(otherDictionary);
+    IntVector encoded = (IntVector) DictionaryEncoder.encode(vector, dictionary);
+    IntVector otherEncoded = (IntVector) DictionaryEncoder.encode(otherVector, otherDictionary);
+
+    try (CDataDictionaryProvider cDictionaryProvider = new CDataDictionaryProvider()) {
+      FieldVector imported;
+      try (ArrowSchema arrowSchema = ArrowSchema.allocateNew(allocator);
+          ArrowArray arrowArray = ArrowArray.allocateNew(allocator)) {
+        Data.exportVector(allocator, encoded, provider, arrowArray, arrowSchema);
+        imported = Data.importVector(allocator, arrowArray, arrowSchema, cDictionaryProvider);
+      }
+
+      try {
+        // Import a second, invalid batch with a different dictionary into the same vector and
+        // provider
+        try (ArrowSchema arrowSchema = ArrowSchema.allocateNew(allocator);
+            ArrowArray arrowArray = ArrowArray.allocateNew(allocator)) {
+          Data.exportVector(allocator, otherEncoded, otherProvider, arrowArray, arrowSchema);
+          // Only the array is imported, so release the exported schema here
+          arrowSchema.release();
+          corrupt.accept(arrowArray);
+          assertThrows(
+              IllegalArgumentException.class,
+              () -> Data.importIntoVector(allocator, arrowArray, imported, cDictionaryProvider));
+        }
+
+        assertTrue(
+            VectorEqualsVisitor.vectorEquals(
+                dictVector, cDictionaryProvider.lookup(0L).getVector()),
+            "dictionary was modified by the rejected import");
+      } finally {
+        imported.close();
+      }
+    } finally {
+      AutoCloseables.close(
+          (AutoCloseable) vector, otherVector, encoded, otherEncoded, dictVector, otherDictVector);
+    }
+  }
+
+  private VarCharVector newVarCharVector(String name, String... values) {
+    VarCharVector vector = new VarCharVector(name, allocator);
+    vector.allocateNewSafe();
+    for (int i = 0; i < values.length; i++) {
+      vector.setSafe(i, values[i].getBytes(StandardCharsets.UTF_8));
+    }
+    vector.setValueCount(values.length);
+    return vector;
+  }
+
   @Test
   public void testRoundtripMultipleBatches() throws IOException {
     try (ArrowStreamReader reader = createMultiBatchReader();

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<ArrowArray> 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);