From f92124f29cc41c078f46c9e2d32cfb9dd2f7a9a1 Mon Sep 17 00:00:00 2001 From: Rok Mihevc Date: Tue, 30 Jun 2026 14:05:51 +0200 Subject: [PATCH] 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. --- cpp/src/arrow/ipc/metadata_internal.cc | 8 +++ cpp/src/arrow/ipc/reader.cc | 39 +++++++++---- cpp/src/arrow/ipc/tensor_test.cc | 77 +++++++++++++++++++++----- cpp/src/arrow/sparse_tensor.cc | 16 ++++++ cpp/src/arrow/sparse_tensor.h | 8 ++- 5 files changed, 122 insertions(+), 26 deletions(-) diff --git a/cpp/src/arrow/ipc/metadata_internal.cc b/cpp/src/arrow/ipc/metadata_internal.cc index eef5cc221f5f..94a956916162 100644 --- 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 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 } 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 index 376d0fab4b51..8ea04aaba62f 100644 --- a/cpp/src/arrow/ipc/reader.cc +++ b/cpp/src/arrow/ipc/reader.cc @@ -2293,6 +2293,25 @@ Result> 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(indptr_buffers->size()) != ndim - 1 || + static_cast(indices_buffers->size()) != ndim || + static_cast(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> ReadSparseCOOIndex( const flatbuf::SparseTensor* sparse_tensor, const std::vector& shape, int64_t non_zero_length, io::RandomAccessFile* file) { @@ -2339,9 +2358,6 @@ Result> 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> ReadSparseCSFIndex( io::RandomAccessFile* file) { auto* sparse_index = sparse_tensor->sparseIndex_as_SparseTensorIndexCSF(); const auto ndim = static_cast(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(indptr_buffers->size()) != ndim - 1 || - static_cast(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> indptr_data(ndim - 1); std::vector> indices_data(ndim); @@ -2526,6 +2536,9 @@ Result 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> ReadSparseTensorPayload(const IpcPayload& std::shared_ptr indptr_type, indices_type; std::vector axis_order, indices_size; + auto fb_sparse_index = sparse_tensor->sparseIndex_as_SparseTensorIndexCSF(); + RETURN_NOT_OK(ValidateSparseCSFIndexMetadata(fb_sparse_index, + static_cast(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 index 4890dfc9926d..a9243e779932 100644 --- 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> MakeCSFSparseTensorMessage(int num_dims, - int num_indptr_buffers, - int num_indices_buffers, - int axis_order_size) { +Result> MakeCSFSparseTensorMessage( + const std::vector& 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> 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> 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> MakeCSFSparseTensorMessage(int num_dims, return std::shared_ptr(std::move(message)); } +IpcPayload MakeSparseTensorPayload(const std::shared_ptr& 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::max(), 1}, + /*non_zero_length=*/0, empty, empty)); + ASSERT_RAISES(Invalid, + SparseCSCIndex::Make(int64(), {1, std::numeric_limits::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 index 0852a0cdb898..c23612e25306 100644 --- 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 ComputeSparseCSXIndptrLength(SparseMatrixCompressedAxis compressed_axis, + const std::vector& 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({compressed_axis_size, static_cast(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& indptr_type, const std::shared_ptr& indices_type, const std::vector& indptr_shape, diff --git a/cpp/src/arrow/sparse_tensor.h b/cpp/src/arrow/sparse_tensor.h index 5faae16bb25c..a41024b74160 100644 --- 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 ComputeSparseCSXIndptrLength(SparseMatrixCompressedAxis compressed_axis, + const std::vector& shape); + ARROW_EXPORT Status ValidateSparseCSXIndex(const std::shared_ptr& indptr_type, const std::shared_ptr& indices_type, @@ -252,7 +256,9 @@ class SparseCSXIndex : public SparseIndexBase { const std::shared_ptr& indices_type, const std::vector& shape, int64_t non_zero_length, std::shared_ptr indptr_data, std::shared_ptr indices_data) { - std::vector indptr_shape({shape[0] + 1}); + ARROW_ASSIGN_OR_RAISE(auto indptr_length, + ComputeSparseCSXIndptrLength(COMPRESSED_AXIS, shape)); + std::vector indptr_shape({indptr_length}); std::vector indices_shape({non_zero_length}); return Make(indptr_type, indices_type, indptr_shape, indices_shape, indptr_data, indices_data);