diff --git a/be/src/storage/segment/binary_dict_page.cpp b/be/src/storage/segment/binary_dict_page.cpp index fb85020bb15f13..eb7b53f9d3832a 100644 --- a/be/src/storage/segment/binary_dict_page.cpp +++ b/be/src/storage/segment/binary_dict_page.cpp @@ -265,6 +265,18 @@ void BinaryDictPageDecoder::set_dict_decoder(uint32_t num_dict_items, StringRef* _dict_word_info = dict_word_info; }; +Status BinaryDictPageDecoder::_validate_dict_codes(const int32_t* codes, size_t start_index, + size_t count) const { + for (size_t i = 0; i < count; ++i) { + int32_t code = codes[start_index + i]; + if (UNLIKELY(code < 0 || std::cmp_greater_equal(code, _num_dict_items))) { + return Status::Corruption("Invalid dictionary code {} at index {}, dictionary size {}", + code, start_index + i, _num_dict_items); + } + } + return Status::OK(); +} + Status BinaryDictPageDecoder::next_batch(size_t* n, MutableColumnPtr& dst) { if (!is_dict_encoding()) { dst = dst->convert_to_predicate_column_if_dictionary(); @@ -281,7 +293,9 @@ Status BinaryDictPageDecoder::next_batch(size_t* n, MutableColumnPtr& dst) { size_t max_fetch = std::min(*n, static_cast(_bit_shuffle_ptr->_num_elements - _bit_shuffle_ptr->_cur_index)); - *n = max_fetch; + const auto* data_array = reinterpret_cast(_bit_shuffle_ptr->get_data(0)); + size_t start_index = _bit_shuffle_ptr->_cur_index; + RETURN_IF_ERROR(_validate_dict_codes(data_array, start_index, max_fetch)); if (_options.only_read_offsets) { // OFFSET_ONLY mode: resolve dict codes to get real string lengths @@ -290,8 +304,6 @@ Status BinaryDictPageDecoder::next_batch(size_t* n, MutableColumnPtr& dst) { // it to a predicate column (ColumnString) first. This is a no-op for // non-dictionary columns and for ColumnNullable it converts the nested column. dst = dst->convert_to_predicate_column_if_dictionary(); - const auto* data_array = reinterpret_cast(_bit_shuffle_ptr->get_data(0)); - size_t start_index = _bit_shuffle_ptr->_cur_index; // Reuse _buffer (int32_t vector) to store uint32_t lengths. // int32_t and uint32_t have the same size/alignment, and string // lengths are always non-negative, so the bit patterns are identical. @@ -303,13 +315,11 @@ Status BinaryDictPageDecoder::next_batch(size_t* n, MutableColumnPtr& dst) { dst->insert_offsets_from_lengths(reinterpret_cast(_buffer.data()), max_fetch); } else { - const auto* data_array = reinterpret_cast(_bit_shuffle_ptr->get_data(0)); - size_t start_index = _bit_shuffle_ptr->_cur_index; - dst->insert_many_dict_data(data_array, start_index, _dict_word_info, max_fetch, _num_dict_items); } + *n = max_fetch; _bit_shuffle_ptr->_cur_index += max_fetch; return Status::OK(); @@ -330,33 +340,6 @@ Status BinaryDictPageDecoder::read_by_rowids(const rowid_t* rowids, ordinal_t pa } auto total = *n; - - if (_options.only_read_offsets) { - // OFFSET_ONLY mode: resolve dict codes to get real string lengths - // without copying actual char data. This allows length() to work correctly. - // ColumnDictI32 does not implement insert_offsets_from_lengths, so convert - // it to a predicate column (ColumnString) first. - dst = dst->convert_to_predicate_column_if_dictionary(); - const auto* data_array = reinterpret_cast(_bit_shuffle_ptr->get_data(0)); - size_t read_count = 0; - _buffer.resize(total); - for (size_t i = 0; i < total; ++i) { - ordinal_t ord = rowids[i] - page_first_ordinal; - if (ord >= _bit_shuffle_ptr->_num_elements) [[unlikely]] { - break; - } - int32_t codeword = data_array[ord]; - _buffer[read_count] = static_cast(_dict_word_info[codeword].size); - read_count++; - } - if (read_count > 0) { - dst->insert_offsets_from_lengths(reinterpret_cast(_buffer.data()), - read_count); - } - *n = read_count; - return Status::OK(); - } - const auto* data_array = reinterpret_cast(_bit_shuffle_ptr->get_data(0)); size_t read_count = 0; _buffer.resize(total); @@ -368,8 +351,22 @@ Status BinaryDictPageDecoder::read_by_rowids(const rowid_t* rowids, ordinal_t pa _buffer[read_count++] = data_array[ord]; } + RETURN_IF_ERROR(_validate_dict_codes(_buffer.data(), 0, read_count)); - if (LIKELY(read_count > 0)) { + if (_options.only_read_offsets) { + // OFFSET_ONLY mode: resolve dict codes to get real string lengths + // without copying actual char data. This allows length() to work correctly. + // ColumnDictI32 does not implement insert_offsets_from_lengths, so convert + // it to a predicate column (ColumnString) first. + dst = dst->convert_to_predicate_column_if_dictionary(); + for (size_t i = 0; i < read_count; ++i) { + _buffer[i] = static_cast(_dict_word_info[_buffer[i]].size); + } + if (read_count > 0) { + dst->insert_offsets_from_lengths(reinterpret_cast(_buffer.data()), + read_count); + } + } else if (LIKELY(read_count > 0)) { dst->insert_many_dict_data(_buffer.data(), 0, _dict_word_info, read_count, _num_dict_items); } *n = read_count; diff --git a/be/src/storage/segment/binary_dict_page.h b/be/src/storage/segment/binary_dict_page.h index 60b39cc693d827..761b245a431ad3 100644 --- a/be/src/storage/segment/binary_dict_page.h +++ b/be/src/storage/segment/binary_dict_page.h @@ -149,6 +149,8 @@ class BinaryDictPageDecoder : public PageDecoder { ~BinaryDictPageDecoder() override; private: + Status _validate_dict_codes(const int32_t* codes, size_t start_index, size_t count) const; + Slice _data; PageDecoderOptions _options; std::unique_ptr _data_page_decoder; diff --git a/be/test/storage/segment/binary_dict_page_test.cpp b/be/test/storage/segment/binary_dict_page_test.cpp index f402e98dc258c4..ab8aa3fd5d192f 100644 --- a/be/test/storage/segment/binary_dict_page_test.cpp +++ b/be/test/storage/segment/binary_dict_page_test.cpp @@ -24,8 +24,10 @@ #include #include +#include "common/cast_set.h" #include "common/config.h" #include "common/logging.h" +#include "core/column/column_dictionary.h" #include "core/column/column_string.h" #include "runtime/exec_env.h" #include "storage/olap_common.h" @@ -35,9 +37,11 @@ #include "storage/segment/binary_plain_page_v2_pre_decoder.h" #include "storage/segment/binary_plain_page_v3.h" #include "storage/segment/binary_plain_page_v3_pre_decoder.h" +#include "storage/segment/bitshuffle_page.h" #include "storage/segment/page_builder.h" #include "storage/segment/page_decoder.h" #include "storage/types.h" +#include "util/coding.h" #include "util/debug_util.h" namespace doris { @@ -569,6 +573,110 @@ class BinaryDictPageTest : public testing::Test { std::unique_ptr _resolver; }; +// NOLINTNEXTLINE(readability-function-cognitive-complexity): gtest macros inflate the score. +TEST_F(BinaryDictPageTest, RejectInvalidDictionaryCodes) { + std::vector values = {"first", "last"}; + std::vector slices; + std::vector dictionary; + for (const auto& value : values) { + slices.emplace_back(value); + dictionary.emplace_back(value); + } + + PageBuilderOptions options; + options.data_page_size = 256 * 1024; + options.dict_page_size = 256 * 1024; + auto page_builder = create_and_add_data(slices, options); + ASSERT_NE(nullptr, page_builder); + OwnedSlice encoded_page; + ASSERT_TRUE(page_builder->finish(&encoded_page).ok()); + + Slice decoded_slice = encoded_page.slice(); + std::unique_ptr decoded_page; + ASSERT_TRUE(apply_pre_decode(decoded_slice, decoded_page).ok()); + ASSERT_GE(decoded_slice.size, static_cast(BINARY_DICT_PAGE_HEADER_SIZE) + + static_cast(BITSHUFFLE_PAGE_HEADER_SIZE) + + 2 * sizeof(int32_t)); + + auto check_valid = [&](bool only_read_offsets, bool read_by_rowids) { + PageDecoderOptions decoder_options; + decoder_options.only_read_offsets = only_read_offsets; + BinaryDictPageDecoder decoder(decoded_slice, decoder_options); + ASSERT_TRUE(decoder.init().ok()); + ASSERT_TRUE(decoder.is_dict_encoding()); + decoder.set_dict_decoder(cast_set(dictionary.size()), dictionary.data()); + + MutableColumnPtr output = ColumnString::create(); + size_t count = values.size(); + Status status; + if (read_by_rowids) { + const rowid_t rowids[] = {0, 1}; + status = decoder.read_by_rowids(rowids, 0, &count, output); + } else { + status = decoder.next_batch(&count, output); + } + ASSERT_TRUE(status.ok()) << status; + ASSERT_EQ(values.size(), count); + ASSERT_EQ(values.size(), output->size()); + const auto& string_output = assert_cast(*output); + for (size_t i = 0; i < values.size(); ++i) { + if (only_read_offsets) { + EXPECT_EQ(values[i].size(), string_output.size_at(i)); + } else { + EXPECT_EQ(values[i], string_output.get_data_at(i).to_string()); + } + } + }; + + check_valid(false, false); + check_valid(false, true); + check_valid(true, false); + check_valid(true, true); + + auto check_rejected = [&](bool only_read_offsets, bool read_by_rowids) { + PageDecoderOptions decoder_options; + decoder_options.only_read_offsets = only_read_offsets; + BinaryDictPageDecoder decoder(decoded_slice, decoder_options); + ASSERT_TRUE(decoder.init().ok()); + ASSERT_TRUE(decoder.is_dict_encoding()); + decoder.set_dict_decoder(cast_set(dictionary.size()), dictionary.data()); + + auto dict_output = ColumnDictI32::create(); + dict_output->reserve(1); + const int32_t existing_code = 0; + dict_output->insert_many_dict_data(&existing_code, 0, dictionary.data(), 1, + cast_set(dictionary.size())); + MutableColumnPtr output = std::move(dict_output); + const auto* original_output = output.get(); + size_t count = 2; + Status status; + if (read_by_rowids) { + const rowid_t rowids[] = {0, 1}; + status = decoder.read_by_rowids(rowids, 0, &count, output); + } else { + status = decoder.next_batch(&count, output); + } + EXPECT_TRUE(status.is()) << status; + EXPECT_EQ(2, count); + EXPECT_EQ(original_output, output.get()); + EXPECT_EQ(1, output->size()); + EXPECT_EQ(existing_code, assert_cast(*output).get_data()[0]); + EXPECT_EQ(0, decoder.current_index()); + }; + + constexpr size_t second_code_offset = static_cast(BINARY_DICT_PAGE_HEADER_SIZE) + + static_cast(BITSHUFFLE_PAGE_HEADER_SIZE) + + sizeof(int32_t); + for (int32_t invalid_code : {-1, static_cast(dictionary.size())}) { + encode_fixed32_le(reinterpret_cast(decoded_slice.data + second_code_offset), + static_cast(invalid_code)); + check_rejected(false, false); + check_rejected(false, true); + check_rejected(true, false); + check_rejected(true, true); + } +} + // Local behavior tests - test specific config behavior TEST_F(BinaryDictPageTest, TestConfigUsePlainBinaryV2False) { std::vector slices;