This is an automated email from the ASF dual-hosted git repository.
rok pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow.git
The following commit(s) were added to refs/heads/main by this push:
new 72a970f713 GH-50304: [C++][IPC] Reject negative sparse tensor shape
and non-zero length (#50305)
72a970f713 is described below
commit 72a970f713b334e9faa77113daf4dcf77c1bcdc2
Author: Rok Mihevc <[email protected]>
AuthorDate: Tue Jun 30 17:14:19 2026 +0200
GH-50304: [C++][IPC] Reject negative sparse tensor shape and non-zero
length (#50305)
### 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 <[email protected]>
Signed-off-by: Rok Mihevc <[email protected]>
---
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 eef5cc221f..94a9569161 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<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
index 376d0fab4b..8ea04aaba6 100644
--- 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
index 4890dfc992..a9243e7799 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<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
index 0852a0cdb8..c23612e253 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<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
index 5faae16bb2..a41024b741 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<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);
