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

Validate dimension sizes and non_zero_length in GetSparseTensorMetadata so
malformed IPC SparseTensor messages are rejected before the values flow into
buffer-size computations. The now-redundant non-zero-length check in
ReadSparseCSXIndex is removed in favour of this central guard.

Author: rok

cpp/src/arrow/ipc/metadata_internal.cc

@@ -1559,6 +1559,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 +1572,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

@@ -2339,9 +2339,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();
 

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);
 
@@ -560,28 +559,46 @@ Result<std::shared_ptr<Message>> MakeCSFSparseTensorMessage(int num_dims,
 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(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