GH-50304: [C++][IPC] Reject negative sparse tensor shape and non-zero length (#50305)

rok · web-flow · commit 72a970f713b3 · 2026-06-30T17:14:19.000+02:00
### Rationale for this change `GetSparseTensorMetadata` and the sparse-tensor readers in `cpp/src/arrow/ipc/` currently accept several sizes from an IPC SparseTensor message without validation. We Want to introduce validations. `GetSparseTensorMetadata` accepts the per-dimension `TensorDim.size` and the tensor's `non_zero_length` from an IPC `SparseTensor` message without validation. ### What changes are included in this PR? - Reject a negative `size` and `non_zero_length` in `GetSparseTensorMetadata`. - Remove the now-redundant `non_zero_length < 0` check in `ReadSparseCSXIndex` ### Are these changes tested? Yes, see `cpp/src/arrow/ipc/tensor_test.cc`. ### Are there any user-facing changes? IPC messages are more scrutinized. * GitHub Issue: #50304 Authored-by: Rok Mihevc <rok@mihevc.org> Signed-off-by: Rok Mihevc <rok@mihevc.org>
diff --git a/cpp/src/arrow/ipc/metadata_internal.cc b/cpp/src/arrow/ipc/metadata_internal.cc
@@ -1524,6 +1524,7 @@ Status GetSparseCSFIndexMetadata(const flatbuf::SparseTensorIndexCSF* sparse_ind
 
   auto* fb_axis_order = sparse_index->axisOrder();
   auto* fb_indices_buffers = sparse_index->indicesBuffers();
+  // ValidateSparseCSFIndexMetadata already checks this, keep this check defensively.
   if (fb_axis_order == nullptr || fb_indices_buffers == nullptr ||
       fb_axis_order->size() != fb_indices_buffers->size()) {
     return Status::Invalid(
@@ -1559,6 +1560,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());
       }
 
@@ -1569,6 +1573,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();
   }
 
diff --git a/cpp/src/arrow/ipc/reader.cc b/cpp/src/arrow/ipc/reader.cc
@@ -2293,6 +2293,25 @@ Result<std::shared_ptr<Tensor>> ReadTensor(const Message& message) {
 
 namespace {
 
+Status ValidateSparseCSFIndexMetadata(const flatbuf::SparseTensorIndexCSF* sparse_index,
+                                      int64_t ndim) {
+  if (sparse_index == nullptr) {
+    return Status::Invalid("Missing CSF sparse index metadata");
+  }
+  auto* indptr_buffers = sparse_index->indptrBuffers();
+  auto* indices_buffers = sparse_index->indicesBuffers();
+  auto* axis_order = sparse_index->axisOrder();
+  if (ndim < 2 || indptr_buffers == nullptr || indices_buffers == nullptr ||
+      axis_order == nullptr || static_cast<int64_t>(indptr_buffers->size()) != ndim - 1 ||
+      static_cast<int64_t>(indices_buffers->size()) != ndim ||
+      static_cast<int64_t>(axis_order->size()) != ndim) {
+    return Status::Invalid(
+        "Inconsistent CSF sparse index: a CSF tensor must have at least 2 dimensions "
+        "with indptr, indices and axis_order counts of ndim - 1, ndim and ndim");
+  }
+  return Status::OK();
+}
+
 Result<std::shared_ptr<SparseIndex>> ReadSparseCOOIndex(
     const flatbuf::SparseTensor* sparse_tensor, const std::vector<int64_t>& shape,
     int64_t non_zero_length, io::RandomAccessFile* file) {
@@ -2339,9 +2358,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();
 
@@ -2411,15 +2427,9 @@ 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());
+  RETURN_NOT_OK(ValidateSparseCSFIndexMetadata(sparse_index, ndim));
   auto* indptr_buffers = sparse_index->indptrBuffers();
   auto* indices_buffers = sparse_index->indicesBuffers();
-  if (ndim < 2 || indptr_buffers == nullptr || indices_buffers == nullptr ||
-      static_cast<int64_t>(indptr_buffers->size()) != ndim - 1 ||
-      static_cast<int64_t>(indices_buffers->size()) != ndim) {
-    return Status::Invalid(
-        "Inconsistent CSF sparse index: a CSF tensor must have at least 2 dimensions "
-        "with indptr and indices buffer counts of ndim - 1 and ndim");
-  }
   std::vector<std::shared_ptr<Buffer>> indptr_data(ndim - 1);
   std::vector<std::shared_ptr<Buffer>> indices_data(ndim);
 
@@ -2526,6 +2536,9 @@ Result<size_t> GetSparseTensorBodyBufferCount(SparseTensorFormat::type format_id
       return 3;
 
     case SparseTensorFormat::CSF:
+      if (ndim < 2) {
+        return Status::Invalid("Invalid shape length for a sparse CSF tensor");
+      }
       return 2 * ndim;
 
     default:
@@ -2621,9 +2634,11 @@ Result<std::shared_ptr<SparseTensor>> ReadSparseTensorPayload(const IpcPayload&
       std::shared_ptr<DataType> indptr_type, indices_type;
       std::vector<int64_t> axis_order, indices_size;
 
+      auto fb_sparse_index = sparse_tensor->sparseIndex_as_SparseTensorIndexCSF();
+      RETURN_NOT_OK(ValidateSparseCSFIndexMetadata(fb_sparse_index,
+                                                   static_cast<int64_t>(shape.size())));
       RETURN_NOT_OK(internal::GetSparseCSFIndexMetadata(
-          sparse_tensor->sparseIndex_as_SparseTensorIndexCSF(), &axis_order,
-          &indices_size, &indptr_type, &indices_type));
+          fb_sparse_index, &axis_order, &indices_size, &indptr_type, &indices_type));
       ARROW_CHECK_EQ(indptr_type, indices_type);
 
       const int64_t ndim = shape.size();
diff --git a/cpp/src/arrow/ipc/tensor_test.cc b/cpp/src/arrow/ipc/tensor_test.cc
@@ -510,17 +510,16 @@ namespace {
 // Build a SparseTensor IPC message carrying a CSF index with caller-controlled
 // buffer counts. This lets us exercise the reader's validation directly, since
 // such inconsistent indices can only be produced by hand-crafted flatbuffers.
-Result<std::shared_ptr<Message>> MakeCSFSparseTensorMessage(int num_dims,
-                                                            int num_indptr_buffers,
-                                                            int num_indices_buffers,
-                                                            int axis_order_size) {
+Result<std::shared_ptr<Message>> MakeCSFSparseTensorMessage(
+    const std::vector<int64_t>& shape, int num_indptr_buffers, int num_indices_buffers,
+    int axis_order_size, int64_t non_zero_length = 0) {
   flatbuffers::FlatBufferBuilder fbb;
 
   auto value_type = flatbuf::CreateInt(fbb, 64, /*is_signed=*/true);
 
   std::vector<flatbuffers::Offset<flatbuf::TensorDim>> dims;
-  for (int i = 0; i < num_dims; ++i) {
-    dims.push_back(flatbuf::CreateTensorDim(fbb, /*size=*/4, /*name=*/0));
+  for (int64_t dim_size : shape) {
+    dims.push_back(flatbuf::CreateTensorDim(fbb, dim_size, /*name=*/0));
   }
   auto fb_shape = fbb.CreateVector(dims);
 
@@ -540,7 +539,7 @@ Result<std::shared_ptr<Message>> MakeCSFSparseTensorMessage(int num_dims,
 
   flatbuf::Buffer data(0, 0);
   auto sparse_tensor = flatbuf::CreateSparseTensor(
-      fbb, flatbuf::Type_Int, value_type.Union(), fb_shape, /*non_zero_length=*/0,
+      fbb, flatbuf::Type_Int, value_type.Union(), fb_shape, non_zero_length,
       flatbuf::SparseTensorIndex::SparseTensorIndex_SparseTensorIndexCSF, csf.Union(),
       &data);
 
@@ -555,33 +554,85 @@ Result<std::shared_ptr<Message>> MakeCSFSparseTensorMessage(int num_dims,
   return std::shared_ptr<Message>(std::move(message));
 }
 
+IpcPayload MakeSparseTensorPayload(const std::shared_ptr<Message>& message,
+                                   int num_body_buffers) {
+  IpcPayload payload;
+  payload.metadata = message->metadata();
+  payload.body_buffers.assign(num_body_buffers, Buffer::FromString(""));
+  return payload;
+}
+
 }  // namespace
 
 TEST(TestSparseCSFIndex, RejectInconsistentBufferCounts) {
   // ndim == 1 is not a valid CSF index (it has no indptr buffers), and used to
   // reach SparseCSFIndex's constructor with an empty indptr vector.
-  ASSERT_OK_AND_ASSIGN(
-      auto message, MakeCSFSparseTensorMessage(/*num_dims=*/1, /*num_indptr_buffers=*/0,
-                                               /*num_indices_buffers=*/1,
-                                               /*axis_order_size=*/1));
+  ASSERT_OK_AND_ASSIGN(auto message,
+                       MakeCSFSparseTensorMessage(/*shape=*/{4}, /*num_indptr_buffers=*/0,
+                                                  /*num_indices_buffers=*/1,
+                                                  /*axis_order_size=*/1));
   ASSERT_RAISES(Invalid, ReadSparseTensor(*message));
 
   // Too many indices buffers for the declared number of dimensions, which used
   // to write past the end of the fixed-size index vectors.
   ASSERT_OK_AND_ASSIGN(
-      message, MakeCSFSparseTensorMessage(/*num_dims=*/2, /*num_indptr_buffers=*/1,
+      message, MakeCSFSparseTensorMessage(/*shape=*/{4, 4}, /*num_indptr_buffers=*/1,
                                           /*num_indices_buffers=*/3,
                                           /*axis_order_size=*/2));
   ASSERT_RAISES(Invalid, ReadSparseTensor(*message));
 
   // axisOrder and indicesBuffers lengths disagree (out-of-bounds read).
   ASSERT_OK_AND_ASSIGN(
-      message, MakeCSFSparseTensorMessage(/*num_dims=*/2, /*num_indptr_buffers=*/1,
+      message, MakeCSFSparseTensorMessage(/*shape=*/{4, 4}, /*num_indptr_buffers=*/1,
                                           /*num_indices_buffers=*/2,
                                           /*axis_order_size=*/3));
   ASSERT_RAISES(Invalid, ReadSparseTensor(*message));
 }
 
+TEST(TestSparseCSFIndex, RejectInconsistentPayloadBufferCounts) {
+  ASSERT_OK_AND_ASSIGN(auto message,
+                       MakeCSFSparseTensorMessage(/*shape=*/{4}, /*num_indptr_buffers=*/0,
+                                                  /*num_indices_buffers=*/1,
+                                                  /*axis_order_size=*/1));
+  ASSERT_RAISES(Invalid,
+                internal::ReadSparseTensorPayload(MakeSparseTensorPayload(message, 2)));
+
+  ASSERT_OK_AND_ASSIGN(
+      message, MakeCSFSparseTensorMessage(/*shape=*/{4, 4}, /*num_indptr_buffers=*/1,
+                                          /*num_indices_buffers=*/3,
+                                          /*axis_order_size=*/3));
+  ASSERT_RAISES(Invalid,
+                internal::ReadSparseTensorPayload(MakeSparseTensorPayload(message, 4)));
+}
+
+TEST(TestSparseCSXIndex, RejectIndptrLengthOverflow) {
+  auto empty = Buffer::FromString("");
+  ASSERT_RAISES(Invalid,
+                SparseCSRIndex::Make(int64(), {std::numeric_limits<int64_t>::max(), 1},
+                                     /*non_zero_length=*/0, empty, empty));
+  ASSERT_RAISES(Invalid,
+                SparseCSCIndex::Make(int64(), {1, std::numeric_limits<int64_t>::max()},
+                                     /*non_zero_length=*/0, empty, empty));
+}
+
+TEST(TestSparseTensor, RejectNegativeShapeAndNonZeroLength) {
+  // A negative non_zero_length must be rejected by GetSparseTensorMetadata,
+  // otherwise the negative size product bypasses the index buffer-size guards.
+  ASSERT_OK_AND_ASSIGN(
+      auto message, MakeCSFSparseTensorMessage(/*shape=*/{4, 4}, /*num_indptr_buffers=*/1,
+                                               /*num_indices_buffers=*/2,
+                                               /*axis_order_size=*/2,
+                                               /*non_zero_length=*/-1));
+  ASSERT_RAISES(Invalid, ReadSparseTensor(*message));
+
+  // A negative dimension size must likewise be rejected.
+  ASSERT_OK_AND_ASSIGN(
+      message, MakeCSFSparseTensorMessage(/*shape=*/{-1, 4}, /*num_indptr_buffers=*/1,
+                                          /*num_indices_buffers=*/2,
+                                          /*axis_order_size=*/2));
+  ASSERT_RAISES(Invalid, ReadSparseTensor(*message));
+}
+
 }  // namespace test
 }  // namespace ipc
 }  // namespace arrow
diff --git a/cpp/src/arrow/sparse_tensor.cc b/cpp/src/arrow/sparse_tensor.cc
@@ -26,6 +26,7 @@
 #include "arrow/compare.h"
 #include "arrow/type_traits.h"
 #include "arrow/util/checked_cast.h"
+#include "arrow/util/int_util_overflow.h"
 #include "arrow/util/logging_internal.h"
 #include "arrow/visit_type_inline.h"
 
@@ -297,6 +298,21 @@ std::string SparseCOOIndex::ToString() const { return std::string("SparseCOOInde
 
 namespace internal {
 
+Result<int64_t> ComputeSparseCSXIndptrLength(SparseMatrixCompressedAxis compressed_axis,
+                                             const std::vector<int64_t>& shape) {
+  if (shape.size() != 2) {
+    return Status::Invalid("Invalid shape length for a sparse matrix");
+  }
+  const int64_t compressed_axis_size =
+      compressed_axis == SparseMatrixCompressedAxis::ROW ? shape[0] : shape[1];
+  const auto indptr_length =
+      AddWithOverflow<int64_t>({compressed_axis_size, static_cast<int64_t>(1)});
+  if (!indptr_length.has_value()) {
+    return Status::Invalid("shape is inconsistent to the size of indptr buffer");
+  }
+  return indptr_length.value();
+}
+
 Status ValidateSparseCSXIndex(const std::shared_ptr<DataType>& indptr_type,
                               const std::shared_ptr<DataType>& indices_type,
                               const std::vector<int64_t>& indptr_shape,
diff --git a/cpp/src/arrow/sparse_tensor.h b/cpp/src/arrow/sparse_tensor.h
@@ -203,6 +203,10 @@ enum class SparseMatrixCompressedAxis : char {
   COLUMN
 };
 
+ARROW_EXPORT
+Result<int64_t> ComputeSparseCSXIndptrLength(SparseMatrixCompressedAxis compressed_axis,
+                                             const std::vector<int64_t>& shape);
+
 ARROW_EXPORT
 Status ValidateSparseCSXIndex(const std::shared_ptr<DataType>& indptr_type,
                               const std::shared_ptr<DataType>& indices_type,
@@ -252,7 +256,9 @@ class SparseCSXIndex : public SparseIndexBase<SparseIndexType> {
       const std::shared_ptr<DataType>& indices_type, const std::vector<int64_t>& shape,
       int64_t non_zero_length, std::shared_ptr<Buffer> indptr_data,
       std::shared_ptr<Buffer> indices_data) {
-    std::vector<int64_t> indptr_shape({shape[0] + 1});
+    ARROW_ASSIGN_OR_RAISE(auto indptr_length,
+                          ComputeSparseCSXIndptrLength(COMPRESSED_AXIS, shape));
+    std::vector<int64_t> indptr_shape({indptr_length});
     std::vector<int64_t> indices_shape({non_zero_length});
     return Make(indptr_type, indices_type, indptr_shape, indices_shape, indptr_data,
                 indices_data);
