GH-1462: Check IPC and C Data counts before using them and release message bodies on errors

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;
 
@@ -27,6 +26,7 @@
 import org.apache.arrow.memory.BufferAllocator;
 import org.apache.arrow.util.Preconditions;
 import org.apache.arrow.vector.FieldVector;
+import org.apache.arrow.vector.TypeLayout;
 import org.apache.arrow.vector.dictionary.Dictionary;
 import org.apache.arrow.vector.dictionary.DictionaryProvider;
 import org.apache.arrow.vector.ipc.message.ArrowFieldNode;
@@ -93,16 +93,33 @@ private void doImport(ArrowArray.Snapshot snapshot) {
         "ArrowArray struct has non-zero offset (%s), which is not supported",
         snapshot.offset);
 
-    // First import children (required for reconstituting parent array data)
-    long[] children =
-        NativeUtil.toJavaArray(snapshot.children, checkedCastToInt(snapshot.n_children));
+    // 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). Check the number of
+    // children before reading the child pointers, so that an invalid count is not used to size
+    // the read.
+    List<FieldVector> childVectors = vector.getChildrenFromFields();
+    checkState(
+        snapshot.n_children == 0 || snapshot.n_children == childVectors.size(),
+        "ArrowArray struct has %s children (expected %s)",
+        snapshot.n_children,
+        childVectors.size());
+    long[] children = NativeUtil.toJavaArray(snapshot.children, (int) snapshot.n_children);
     if (children != null && children.length > 0) {
-      List<FieldVector> childVectors = vector.getChildrenFromFields();
-      checkState(
-          children.length == childVectors.size(),
-          "ArrowArray struct has %s children (expected %s)",
-          children.length,
-          childVectors.size());
       for (int i = 0; i < children.length; i++) {
         checkState(children[i] != NULL, "ArrowArray struct has NULL child at position %s", i);
         ArrayImporter childImporter =
@@ -111,39 +128,59 @@ 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 {
+      // Check the number of buffers before reading the buffer pointers. Types with variadic
+      // buffers, such as views, have no fixed upper bound.
+      TypeLayout layout = TypeLayout.getTypeLayout(vector.getField().getType());
+      checkState(
+          snapshot.n_buffers >= 0,
+          "ArrowArray struct has a negative number of buffers (%s)",
+          snapshot.n_buffers);
+      checkState(
+          !layout.isFixedBufferCount() || snapshot.n_buffers <= layout.getBufferTypes().size(),
+          "ArrowArray struct has %s buffers (expected at most %s)",
+          snapshot.n_buffers,
+          layout.getBufferTypes().size());
+      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();