[C++][IPC] Validate sparse tensor index sizes to prevent OOB access

Author: rok

cpp/src/arrow/ipc/metadata_internal.cc

@@ -1523,6 +1523,10 @@ Status GetSparseCSFIndexMetadata(const flatbuf::SparseTensorIndexCSF* sparse_ind
   RETURN_NOT_OK(IntFromFlatbuffer(sparse_index->indicesType(), indices_type));
 
   const int ndim = static_cast<int>(sparse_index->axisOrder()->size());
+  if (static_cast<int>(sparse_index->indicesBuffers()->size()) != ndim ||
+      static_cast<int>(sparse_index->indptrBuffers()->size()) != ndim - 1) {
+    return Status::Invalid("Inconsistent CSF index buffer counts");
+  }
   for (int i = 0; i < ndim; ++i) {
     axis_order->push_back(sparse_index->axisOrder()->Get(i));
     indices_size->push_back(sparse_index->indicesBuffers()->Get(i)->length());
@@ -1550,6 +1554,9 @@ Status GetSparseTensorMetadata(const Buffer& metadata, std::shared_ptr<DataType>
       auto dim = sparse_tensor->shape()->Get(i);
 
       if (shape) {
+        if (dim->size() < 0) {
+          return Status::Invalid("Invalid sparse tensor dimension size: ", dim->size());
+        }
         shape->push_back(dim->size());
       }
 
@@ -1560,6 +1567,10 @@ Status GetSparseTensorMetadata(const Buffer& metadata, std::shared_ptr<DataType>
   }
 
   if (non_zero_length) {
+    if (sparse_tensor->non_zero_length() < 0) {
+      return Status::Invalid("Invalid sparse tensor non-zero length: ",
+                             sparse_tensor->non_zero_length());
+    }
     *non_zero_length = sparse_tensor->non_zero_length();
   }
 

cpp/src/arrow/ipc/reader.cc

@@ -2332,9 +2332,6 @@ Result<std::shared_ptr<SparseIndex>> ReadSparseCSXIndex(
   if (shape.size() != 2) {
     return Status::Invalid("Invalid shape length for a sparse matrix");
   }
-  if (non_zero_length < 0) {
-    return Status::Invalid("Invalid non-zero length for a sparse matrix");
-  }
 
   auto* sparse_index = sparse_tensor->sparseIndex_as_SparseMatrixIndexCSX();
 
@@ -2404,8 +2401,16 @@ Result<std::shared_ptr<SparseIndex>> ReadSparseCSFIndex(
     io::RandomAccessFile* file) {
   auto* sparse_index = sparse_tensor->sparseIndex_as_SparseTensorIndexCSF();
   const auto ndim = static_cast<int64_t>(shape.size());
+  if (ndim < 1) {
+    return Status::Invalid("Invalid number of dimensions for a CSF sparse tensor");
+  }
   auto* indptr_buffers = sparse_index->indptrBuffers();
   auto* indices_buffers = sparse_index->indicesBuffers();
+  if (static_cast<int64_t>(indptr_buffers->size()) != ndim - 1 ||
+      static_cast<int64_t>(indices_buffers->size()) != ndim) {
+    return Status::Invalid(
+        "Sparse CSF index buffer counts are inconsistent with the tensor shape");
+  }
   std::vector<std::shared_ptr<Buffer>> indptr_data(ndim - 1);
   std::vector<std::shared_ptr<Buffer>> indices_data(ndim);
 
@@ -2414,6 +2419,10 @@ Result<std::shared_ptr<SparseIndex>> ReadSparseCSFIndex(
 
   RETURN_NOT_OK(internal::GetSparseCSFIndexMetadata(
       sparse_index, &axis_order, &indices_size, &indptr_type, &indices_type));
+  if (static_cast<int64_t>(axis_order.size()) != ndim) {
+    return Status::Invalid(
+        "Sparse CSF axisOrder length is inconsistent with the tensor shape");
+  }
   for (int i = 0; i < static_cast<int>(indptr_buffers->size()); ++i) {
     ARROW_ASSIGN_OR_RAISE(indptr_data[i], file->ReadAt(indptr_buffers->Get(i)->offset(),
                                                        indptr_buffers->Get(i)->length(),