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;
 
@@ -29,6 +28,7 @@
 import org.apache.arrow.util.Preconditions;
 import org.apache.arrow.vector.ExtensionTypeVector;
 import org.apache.arrow.vector.FieldVector;
+import org.apache.arrow.vector.TypeLayout;
 import org.apache.arrow.vector.complex.StructVector;
 import org.apache.arrow.vector.dictionary.Dictionary;
 import org.apache.arrow.vector.dictionary.DictionaryProvider;
@@ -104,6 +104,15 @@ private void importChild(ArrayImporter parent, ArrowArray.Snapshot snapshot) {
   private void doImport(ArrowArray.Snapshot snapshot) {
     checkState(snapshot.offset >= 0, "ArrowArray struct has negative offset (%s)", snapshot.offset);
     checkState(snapshot.length >= 0, "ArrowArray struct has negative length (%s)", snapshot.length);
+    // 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.
+    checkLengthAndNullCount(vector, snapshot);
+    if (snapshot.dictionary != NULL) {
+      checkLengthAndNullCount(
+          lookupDictionary().getVector(), ArrowArray.wrap(snapshot.dictionary).snapshot());
+    }
     if (snapshot.offset != 0 && snapshot.length != 0) {
       if (storageType().getTypeID() == ArrowTypeID.Struct) {
         importSlicedStruct(snapshot);
@@ -115,6 +124,22 @@ private void doImport(ArrowArray.Snapshot snapshot) {
     }
   }
 
+  /**
+   * Checks that the length and null count of {@code snapshot} fit an import into {@code vector}.
+   */
+  private static void checkLengthAndNullCount(FieldVector vector, ArrowArray.Snapshot snapshot) {
+    try {
+      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);
+    }
+  }
+
   /**
    * Imports a struct with a non-zero offset into the target itself, because transferring a struct
    * matches its children by name, which loses children that have the same name as another. The
@@ -127,7 +152,6 @@ private void importSlicedStruct(ArrowArray.Snapshot snapshot) {
     int offset = (int) (snapshot.offset - skippedRows);
     int length = (int) snapshot.length;
     importChildren(snapshot, snapshot.offset, /* childRowsPerRow= */ 1);
-    importDictionary(snapshot);
 
     snapshot.offset = 0;
     snapshot.length = unslicedLength;
@@ -142,6 +166,8 @@ private void importSlicedStruct(ArrowArray.Snapshot snapshot) {
             Collections.singletonList(sliced.getValidityBuffer()));
       }
     }
+    // Import the dictionary last, so that it is not cleared if this array's own data is invalid
+    importDictionary(snapshot);
   }
 
   /**
@@ -228,8 +254,9 @@ private void importUnsliced(
       ArrowArray.Snapshot snapshot, long skippedRows, boolean trimChildren) {
     // First import children (required for reconstituting parent array data)
     importChildren(snapshot, skippedRows, trimChildren ? childRowsPerRow(storageType()) : 0);
-    importDictionary(snapshot);
     importBuffers(vector, snapshot, skippedRows);
+    // Import the dictionary last, so that it is not cleared if this array's own data is invalid
+    importDictionary(snapshot);
   }
 
   /**
@@ -239,17 +266,18 @@ private void importUnsliced(
    */
   private void importChildren(
       ArrowArray.Snapshot snapshot, long skippedRows, long childRowsPerRow) {
-    long[] children =
-        NativeUtil.toJavaArray(snapshot.children, checkedCastToInt(snapshot.n_children));
-    if (children == null || children.length == 0) {
-      return;
-    }
+    // 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(
-        children.length == childVectors.size(),
+        snapshot.n_children == 0 || snapshot.n_children == childVectors.size(),
         "ArrowArray struct has %s children (expected %s)",
-        children.length,
+        snapshot.n_children,
         childVectors.size());
+    long[] children = NativeUtil.toJavaArray(snapshot.children, (int) snapshot.n_children);
+    if (children == null || children.length == 0) {
+      return;
+    }
     List<Field> childFields = field != null ? field.getChildren() : null;
     for (int i = 0; i < children.length; i++) {
       checkState(children[i] != NULL, "ArrowArray struct has NULL child at position %s", i);
@@ -278,11 +306,7 @@ private void importChildren(
   /** Imports the dictionary of a dictionary encoded array. */
   private void importDictionary(ArrowArray.Snapshot snapshot) {
     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");
+      Dictionary dictionary = lookupDictionary();
 
       // reset the dictionary vector to the initial state
       dictionary.getVector().clear();
@@ -293,25 +317,49 @@ private void importDictionary(ArrowArray.Snapshot snapshot) {
     }
   }
 
+  /** Looks up the dictionary of a dictionary encoded array in the dictionary provider. */
+  private Dictionary lookupDictionary() {
+    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");
+    return dictionary;
+  }
+
   /**
    * Imports the buffers of {@code snapshot} into {@code target}, skipping the first {@code
    * skippedRows} rows, which must be a multiple of 8.
    */
   private void importBuffers(FieldVector target, ArrowArray.Snapshot snapshot, long skippedRows) {
-    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, skippedRows)) {
-      final List<ArrowBuf> buffers;
-      if (bufferPointers == null || bufferPointers.length == 0) {
-        buffers = Collections.emptyList();
-      } else {
-        buffers = target.getField().getType().accept(visitor);
+    try {
+      ArrowFieldNode fieldNode = new ArrowFieldNode(snapshot.length, snapshot.null_count);
+      // 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(target.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, skippedRows)) {
+        final List<ArrowBuf> buffers;
+        if (bufferPointers == null || bufferPointers.length == 0) {
+          buffers = Collections.emptyList();
+        } else {
+          buffers = target.getField().getType().accept(visitor);
+        }
+        target.loadFieldBuffers(fieldNode, buffers);
       }
-      target.loadFieldBuffers(fieldNode, buffers);
     } catch (Exception e) {
       throw new IllegalArgumentException(
           "Could not load buffers for field "

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

@@ -104,6 +104,8 @@ ArrowBuf importBuffer(ArrowType type, int index, long capacity) {
   }
 
   private ArrowBuf importBuffer(ArrowType type, int index, long skippedBytes, 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",
@@ -143,7 +145,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, Math.multiplyExact(skippedRows, bytesPerSlot), capacity);
   }
 

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

@@ -18,6 +18,7 @@
 
 import java.util.List;
 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;
@@ -428,7 +429,12 @@ public static FieldVector importVector(
       boolean closeImportedStructs) {
     Field field = importField(allocator, schema, provider, closeImportedStructs);
     FieldVector vector = field.createVector(allocator);
-    importIntoVector(allocator, array, vector, field, provider, closeImportedStructs);
+    try {
+      importIntoVector(allocator, array, vector, field, 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,8 +27,10 @@
 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;
 import org.apache.arrow.vector.FieldVector;
 import org.apache.arrow.vector.IntVector;
 import org.apache.arrow.vector.ValueVector;
@@ -257,6 +260,105 @@ private VarCharVector newStringVector(String name, String... values) {
     return vector;
   }
 
+  @Test
+  public void testImportRejectsInvalidLengthWithoutReplacingDictionary() throws Exception {
+    assertRejectedImportKeepsDictionary(
+        arrowArray -> {
+          ArrowArray.Snapshot snapshot = arrowArray.snapshot();
+          snapshot.length = Integer.MAX_VALUE + 1L;
+          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();