qzyu999 commented on code in PR #50121: URL: https://github.com/apache/arrow/pull/50121#discussion_r3605377762
########## cpp/src/arrow/extension/variant.cc: ########## @@ -0,0 +1,1314 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#include "arrow/extension/variant.h" + +#include <cstring> + +#include "arrow/extension/variant_internal_util.h" +#include "arrow/util/endian.h" +#include "arrow/util/logging_internal.h" + +namespace arrow::extension::variant { + +// Ensure view classes remain lightweight (stack-allocated, cache-friendly). +static_assert(sizeof(VariantView) <= 32, "VariantView should fit in 32 bytes"); +static_assert(sizeof(VariantObjectView) <= 80, + "VariantObjectView should fit in 80 bytes"); +static_assert(sizeof(VariantArrayView) <= 64, "VariantArrayView should fit in 64 bytes"); + +namespace { + +// --------------------------------------------------------------------------- +// Little-endian helpers (delegate to shared internal utility) +// --------------------------------------------------------------------------- + +using internal::ReadUnsignedLE; + +/// \brief Validate that offsets are monotonically non-decreasing and in bounds. +Status ValidateOffsets(const std::vector<uint32_t>& offsets, int64_t data_length) { + for (size_t i = 1; i < offsets.size(); ++i) { + if (offsets[i] < offsets[i - 1]) { + return Status::Invalid( + "Variant metadata: string offsets are not monotonically " + "non-decreasing at index ", + i); + } + } + if (!offsets.empty() && offsets.back() > static_cast<uint32_t>(data_length)) { + return Status::Invalid("Variant metadata: last string offset ", offsets.back(), + " exceeds data length ", data_length); + } + return Status::OK(); +} + +// --------------------------------------------------------------------------- +// Recursive visitor traversal (internal) +// --------------------------------------------------------------------------- + +Status VisitValueAt(const VariantMetadata& metadata, const uint8_t* data, int64_t length, + int64_t offset, VariantVisitor* visitor, int64_t* bytes_consumed, + int32_t depth); + +Status VisitPrimitive(const uint8_t* data, int64_t length, int64_t offset, uint8_t header, + VariantVisitor* visitor, int64_t* bytes_consumed) { + auto primitive_type = GetPrimitiveType(header); + int64_t pos = offset + 1; + + auto check_remaining = [&](int64_t needed) -> Status { + if (pos + needed > length) { + return Status::Invalid("Variant value: truncated primitive at offset ", offset, + ", need ", needed, " bytes but only ", length - pos, + " remaining"); + } + return Status::OK(); + }; + + switch (primitive_type) { + case PrimitiveType::kNull: + ARROW_RETURN_NOT_OK(visitor->Null()); + *bytes_consumed = 1; + return Status::OK(); + case PrimitiveType::kTrue: + ARROW_RETURN_NOT_OK(visitor->Bool(true)); + *bytes_consumed = 1; + return Status::OK(); + case PrimitiveType::kFalse: + ARROW_RETURN_NOT_OK(visitor->Bool(false)); + *bytes_consumed = 1; + return Status::OK(); + case PrimitiveType::kInt8: { + ARROW_RETURN_NOT_OK(check_remaining(1)); + ARROW_RETURN_NOT_OK(visitor->Int8(static_cast<int8_t>(data[pos]))); + *bytes_consumed = 2; + return Status::OK(); + } + case PrimitiveType::kInt16: { + ARROW_RETURN_NOT_OK(check_remaining(2)); + int16_t value; + std::memcpy(&value, data + pos, 2); + value = bit_util::FromLittleEndian(value); + ARROW_RETURN_NOT_OK(visitor->Int16(value)); + *bytes_consumed = 3; + return Status::OK(); + } + case PrimitiveType::kInt32: { + ARROW_RETURN_NOT_OK(check_remaining(4)); + int32_t value; + std::memcpy(&value, data + pos, 4); + value = bit_util::FromLittleEndian(value); + ARROW_RETURN_NOT_OK(visitor->Int32(value)); + *bytes_consumed = 5; + return Status::OK(); + } + case PrimitiveType::kInt64: { + ARROW_RETURN_NOT_OK(check_remaining(8)); + int64_t value; + std::memcpy(&value, data + pos, 8); + value = bit_util::FromLittleEndian(value); + ARROW_RETURN_NOT_OK(visitor->Int64(value)); + *bytes_consumed = 9; + return Status::OK(); + } + case PrimitiveType::kFloat: { + ARROW_RETURN_NOT_OK(check_remaining(4)); + float value; + std::memcpy(&value, data + pos, 4); + value = bit_util::FromLittleEndian(value); + ARROW_RETURN_NOT_OK(visitor->Float(value)); + *bytes_consumed = 5; + return Status::OK(); + } + case PrimitiveType::kDouble: { + ARROW_RETURN_NOT_OK(check_remaining(8)); + double value; + std::memcpy(&value, data + pos, 8); + value = bit_util::FromLittleEndian(value); + ARROW_RETURN_NOT_OK(visitor->Double(value)); + *bytes_consumed = 9; + return Status::OK(); + } + case PrimitiveType::kDecimal4: { + ARROW_RETURN_NOT_OK(check_remaining(5)); + auto scale = static_cast<int32_t>(data[pos]); + ARROW_RETURN_NOT_OK(visitor->Decimal4(data + pos + 1, scale)); + *bytes_consumed = 6; + return Status::OK(); + } + case PrimitiveType::kDecimal8: { + ARROW_RETURN_NOT_OK(check_remaining(9)); + auto scale = static_cast<int32_t>(data[pos]); + ARROW_RETURN_NOT_OK(visitor->Decimal8(data + pos + 1, scale)); + *bytes_consumed = 10; + return Status::OK(); + } + case PrimitiveType::kDecimal16: { + ARROW_RETURN_NOT_OK(check_remaining(17)); + auto scale = static_cast<int32_t>(data[pos]); + ARROW_RETURN_NOT_OK(visitor->Decimal16(data + pos + 1, scale)); + *bytes_consumed = 18; + return Status::OK(); + } + case PrimitiveType::kDate: { + ARROW_RETURN_NOT_OK(check_remaining(4)); + int32_t value; + std::memcpy(&value, data + pos, 4); + value = bit_util::FromLittleEndian(value); + ARROW_RETURN_NOT_OK(visitor->Date(value)); + *bytes_consumed = 5; + return Status::OK(); + } + case PrimitiveType::kTimestampMicros: { + ARROW_RETURN_NOT_OK(check_remaining(8)); + int64_t value; + std::memcpy(&value, data + pos, 8); + value = bit_util::FromLittleEndian(value); + ARROW_RETURN_NOT_OK(visitor->TimestampMicros(value)); + *bytes_consumed = 9; + return Status::OK(); + } + case PrimitiveType::kTimestampMicrosNTZ: { + ARROW_RETURN_NOT_OK(check_remaining(8)); + int64_t value; + std::memcpy(&value, data + pos, 8); + value = bit_util::FromLittleEndian(value); + ARROW_RETURN_NOT_OK(visitor->TimestampMicrosNTZ(value)); + *bytes_consumed = 9; + return Status::OK(); + } + case PrimitiveType::kBinary: { + ARROW_RETURN_NOT_OK(check_remaining(4)); + uint32_t bin_length; + std::memcpy(&bin_length, data + pos, 4); + bin_length = bit_util::FromLittleEndian(bin_length); + ARROW_RETURN_NOT_OK(check_remaining(4 + static_cast<int64_t>(bin_length))); + auto view = + std::string_view(reinterpret_cast<const char*>(data + pos + 4), bin_length); + ARROW_RETURN_NOT_OK(visitor->Binary(view)); + *bytes_consumed = 1 + 4 + static_cast<int64_t>(bin_length); + return Status::OK(); + } + case PrimitiveType::kString: { + ARROW_RETURN_NOT_OK(check_remaining(4)); + uint32_t str_length; + std::memcpy(&str_length, data + pos, 4); + str_length = bit_util::FromLittleEndian(str_length); + ARROW_RETURN_NOT_OK(check_remaining(4 + static_cast<int64_t>(str_length))); + auto view = + std::string_view(reinterpret_cast<const char*>(data + pos + 4), str_length); + ARROW_RETURN_NOT_OK(visitor->String(view)); + *bytes_consumed = 1 + 4 + static_cast<int64_t>(str_length); + return Status::OK(); + } + case PrimitiveType::kTimeNTZ: { + ARROW_RETURN_NOT_OK(check_remaining(8)); + int64_t value; + std::memcpy(&value, data + pos, 8); + value = bit_util::FromLittleEndian(value); + ARROW_RETURN_NOT_OK(visitor->TimeNTZ(value)); + *bytes_consumed = 9; + return Status::OK(); + } + case PrimitiveType::kTimestampNanos: { + ARROW_RETURN_NOT_OK(check_remaining(8)); + int64_t value; + std::memcpy(&value, data + pos, 8); + value = bit_util::FromLittleEndian(value); + ARROW_RETURN_NOT_OK(visitor->TimestampNanos(value)); + *bytes_consumed = 9; + return Status::OK(); + } + case PrimitiveType::kTimestampNanosNTZ: { + ARROW_RETURN_NOT_OK(check_remaining(8)); + int64_t value; + std::memcpy(&value, data + pos, 8); + value = bit_util::FromLittleEndian(value); + ARROW_RETURN_NOT_OK(visitor->TimestampNanosNTZ(value)); + *bytes_consumed = 9; + return Status::OK(); + } + case PrimitiveType::kUUID: { + ARROW_RETURN_NOT_OK(check_remaining(kUUIDByteLength)); + ARROW_RETURN_NOT_OK(visitor->UUID(data + pos)); + *bytes_consumed = kUUIDByteLength + 1; + return Status::OK(); + } + default: + return Status::Invalid("Variant value: unknown primitive type ", + static_cast<int>(primitive_type)); + } +} + +Status VisitShortString(const uint8_t* data, int64_t length, int64_t offset, + uint8_t header, VariantVisitor* visitor, + int64_t* bytes_consumed) { + int32_t str_len = (header >> 2) & 0x3F; + int64_t pos = offset + 1; + if (pos + str_len > length) { + return Status::Invalid("Variant value: truncated short string at offset ", offset); + } + auto view = std::string_view(reinterpret_cast<const char*>(data + pos), str_len); + ARROW_RETURN_NOT_OK(visitor->String(view)); + *bytes_consumed = 1 + str_len; + return Status::OK(); +} + +Status VisitObject(const VariantMetadata& metadata, const uint8_t* data, int64_t length, + int64_t offset, uint8_t header, VariantVisitor* visitor, + int64_t* bytes_consumed, int32_t depth) { + uint8_t type_info = (header >> 2) & 0x3F; + int32_t field_offset_size = (type_info & 0x03) + 1; + int32_t field_id_size = ((type_info >> 2) & 0x03) + 1; + bool is_large = ((type_info >> 4) & 0x01) != 0; + int32_t num_fields_size = is_large ? 4 : 1; + + int64_t pos = offset + 1; + if (pos + num_fields_size > length) { + return Status::Invalid("Variant value: truncated object num_fields at offset ", + offset); + } + auto num_fields = static_cast<int32_t>(ReadUnsignedLE(data + pos, num_fields_size)); + pos += num_fields_size; + + int64_t field_ids_size = static_cast<int64_t>(num_fields) * field_id_size; + if (pos + field_ids_size > length) { + return Status::Invalid("Variant value: truncated object field_ids at offset ", + offset); + } + std::vector<uint32_t> field_ids(num_fields); + for (int32_t i = 0; i < num_fields; ++i) { + field_ids[i] = ReadUnsignedLE(data + pos, field_id_size); + pos += field_id_size; + } + + int64_t offsets_size = (static_cast<int64_t>(num_fields) + 1) * field_offset_size; + if (pos + offsets_size > length) { + return Status::Invalid("Variant value: truncated object offsets at offset ", offset); + } + std::vector<uint32_t> value_offsets(num_fields + 1); + for (int32_t i = 0; i <= num_fields; ++i) { + value_offsets[i] = ReadUnsignedLE(data + pos, field_offset_size); + pos += field_offset_size; + } + + int64_t data_start = pos; + int64_t total_data_size = static_cast<int64_t>(value_offsets[num_fields]); + if (data_start + total_data_size > length) { + return Status::Invalid("Variant value: object data exceeds buffer at offset ", + offset); + } + + for (int32_t i = 0; i < num_fields; ++i) { + if (value_offsets[i] > static_cast<uint32_t>(total_data_size)) { + return Status::Invalid("Variant value: object field offset ", value_offsets[i], + " at index ", i, " exceeds data size ", total_data_size); + } + } + + ARROW_RETURN_NOT_OK(visitor->StartObject(num_fields)); + + for (int32_t i = 0; i < num_fields; ++i) { + auto field_id = field_ids[i]; + if (field_id >= metadata.strings.size()) { + return Status::Invalid("Variant value: field_id ", field_id, + " exceeds metadata dictionary size ", + metadata.strings.size()); + } + ARROW_RETURN_NOT_OK(visitor->FieldName(metadata.strings[field_id])); + + int64_t field_offset = data_start + value_offsets[i]; + int64_t consumed = 0; + ARROW_RETURN_NOT_OK(VisitValueAt(metadata, data, data_start + total_data_size, + field_offset, visitor, &consumed, depth)); + } + + ARROW_RETURN_NOT_OK(visitor->EndObject()); + *bytes_consumed = (data_start - offset) + total_data_size; + return Status::OK(); +} + +Status VisitArray(const VariantMetadata& metadata, const uint8_t* data, int64_t length, + int64_t offset, uint8_t header, VariantVisitor* visitor, + int64_t* bytes_consumed, int32_t depth) { + uint8_t type_info = (header >> 2) & 0x3F; + int32_t field_offset_size = (type_info & 0x03) + 1; + bool is_large = ((type_info >> 2) & 0x01) != 0; + int32_t num_elements_size = is_large ? 4 : 1; + + int64_t pos = offset + 1; + if (pos + num_elements_size > length) { + return Status::Invalid("Variant value: truncated array num_elements at offset ", + offset); + } + auto num_elements = static_cast<int32_t>(ReadUnsignedLE(data + pos, num_elements_size)); + pos += num_elements_size; + + int64_t offsets_size = (static_cast<int64_t>(num_elements) + 1) * field_offset_size; + if (pos + offsets_size > length) { + return Status::Invalid("Variant value: truncated array offsets at offset ", offset); + } + std::vector<uint32_t> value_offsets(num_elements + 1); + for (int32_t i = 0; i <= num_elements; ++i) { + value_offsets[i] = ReadUnsignedLE(data + pos, field_offset_size); + pos += field_offset_size; + } + + for (int32_t i = 1; i <= num_elements; ++i) { + if (value_offsets[i] < value_offsets[i - 1]) { + return Status::Invalid( + "Variant value: array value offsets are not monotonically " + "non-decreasing at index ", + i); + } + } + + int64_t data_start = pos; + int64_t total_data_size = static_cast<int64_t>(value_offsets[num_elements]); + if (data_start + total_data_size > length) { + return Status::Invalid("Variant value: array data exceeds buffer at offset ", offset); + } + + ARROW_RETURN_NOT_OK(visitor->StartArray(num_elements)); + + for (int32_t i = 0; i < num_elements; ++i) { + int64_t elem_offset = data_start + value_offsets[i]; + int64_t consumed = 0; + ARROW_RETURN_NOT_OK(VisitValueAt(metadata, data, data_start + total_data_size, + elem_offset, visitor, &consumed, depth)); + } + + ARROW_RETURN_NOT_OK(visitor->EndArray()); + *bytes_consumed = (data_start - offset) + total_data_size; + return Status::OK(); +} + +Status VisitValueAt(const VariantMetadata& metadata, const uint8_t* data, int64_t length, + int64_t offset, VariantVisitor* visitor, int64_t* bytes_consumed, + int32_t depth) { + if (offset >= length) { + return Status::Invalid("Variant value: offset ", offset, + " is at or beyond buffer length ", length); + } + if (depth > kMaxNestingDepth) { + return Status::Invalid("Variant value: nesting depth exceeds maximum of ", + kMaxNestingDepth); + } + + uint8_t header = data[offset]; + auto basic_type = GetBasicType(header); + + switch (basic_type) { + case BasicType::kPrimitive: + return VisitPrimitive(data, length, offset, header, visitor, bytes_consumed); + case BasicType::kShortString: + return VisitShortString(data, length, offset, header, visitor, bytes_consumed); + case BasicType::kObject: + return VisitObject(metadata, data, length, offset, header, visitor, bytes_consumed, + depth + 1); + case BasicType::kArray: + return VisitArray(metadata, data, length, offset, header, visitor, bytes_consumed, + depth + 1); + default: + return Status::Invalid("Variant value: unknown basic type ", + static_cast<int>(basic_type)); + } +} + +} // namespace + +// =========================================================================== +// Public API: Metadata +// =========================================================================== + +Result<VariantMetadata> DecodeMetadata(const uint8_t* data, int64_t length) { + if (data == nullptr || length < 1) { + return Status::Invalid("Variant metadata: buffer is null or empty"); + } + + uint8_t header = data[0]; + uint8_t version = header & 0x0F; + if (version != kVariantVersion) { + return Status::Invalid("Variant metadata: unsupported version ", + static_cast<int>(version), ", expected ", + static_cast<int>(kVariantVersion)); + } + + if ((header >> 5) & 0x01) { + return Status::Invalid("Variant metadata: reserved bit 5 is set in header"); + } + + bool is_sorted = ((header >> 4) & 0x01) != 0; + int32_t offset_size = ((header >> 6) & 0x03) + 1; + + int64_t pos = 1; + if (pos + offset_size > length) { + return Status::Invalid("Variant metadata: truncated dictionary size at byte ", pos); + } + auto dict_size = static_cast<int32_t>(ReadUnsignedLE(data + pos, offset_size)); + pos += offset_size; + + int64_t offsets_bytes = static_cast<int64_t>(dict_size + 1) * offset_size; Review Comment: Hi @Explooosion-code, thanks for catching this bug. `ReadUnsignedLE` returns `uint32_t`, and with `offset_size=4` it can return values `≥ 2^31` that wrap to negative when cast to `int32_t`. The negative `dict_size` makes `offsets_bytes` negative, causing `pos + offsets_bytes > length` to evaluate false (bypassing the bounds check), and then `std::vector<uint32_t>(dict_size + 1)` throws on the bogus size. I've pushed a fix (1eed28e) that changes all element-count reads `(dict_size, num_fields, num_elements)` from `static_cast<int32_t>` to `static_cast<int64_t>` across `DecodeMetadata`, `VisitObject`, `VisitArray`, `VariantObjectView::Make`, and `VariantArrayView::Make`. The narrowing cast to `int32_t` for view class storage now happens only after the bounds checks confirm the buffer is actually large enough. Added regression tests that craft buffers with counts `> INT32_MAX` and verify `Status::Invalid` is returned (not a crash). The same pattern already existed correctly in `ValueSize()`, this brings the other paths in line. Furthermore, #50122 and #50232 were rebased based on this fix. Reference: https://github.com/apache/parquet-format/issues/599 -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
