Skip to content

[fix](be) Validate dictionary page codes #66832

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
YanzhiJin5 wants to merge 1 commit into
base: master
Choose a base branch
from
+141 −34
Open

[fix](be) Validate dictionary page codes #66832

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
65 changes: 31 additions & 34 deletions be/src/storage/segment/binary_dict_page.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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();
Expand All @@ -281,7 +293,9 @@ Status BinaryDictPageDecoder::next_batch(size_t* n, MutableColumnPtr& dst) {

size_t max_fetch = std::min(*n, static_cast<size_t>(_bit_shuffle_ptr->_num_elements -
_bit_shuffle_ptr->_cur_index));
*n = max_fetch;
const auto* data_array = reinterpret_cast<const int32_t*>(_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
Expand All @@ -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<const int32_t*>(_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.
Expand All @@ -303,13 +315,11 @@ Status BinaryDictPageDecoder::next_batch(size_t* n, MutableColumnPtr& dst) {
dst->insert_offsets_from_lengths(reinterpret_cast<const uint32_t*>(_buffer.data()),
max_fetch);
} else {
const auto* data_array = reinterpret_cast<const int32_t*>(_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();
Expand All @@ -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<const int32_t*>(_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<int32_t>(_dict_word_info[codeword].size);
read_count++;
}
if (read_count > 0) {
dst->insert_offsets_from_lengths(reinterpret_cast<const uint32_t*>(_buffer.data()),
read_count);
}
*n = read_count;
return Status::OK();
}

const auto* data_array = reinterpret_cast<const int32_t*>(_bit_shuffle_ptr->get_data(0));
size_t read_count = 0;
_buffer.resize(total);
Expand All @@ -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<int32_t>(_dict_word_info[_buffer[i]].size);
}
if (read_count > 0) {
dst->insert_offsets_from_lengths(reinterpret_cast<const uint32_t*>(_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;
Expand Down
2 changes: 2 additions & 0 deletions be/src/storage/segment/binary_dict_page.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<PageDecoder> _data_page_decoder;
Expand Down
108 changes: 108 additions & 0 deletions be/test/storage/segment/binary_dict_page_test.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -24,8 +24,10 @@
#include <memory>
#include <vector>

#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"
Expand All @@ -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 {
Expand Down Expand Up @@ -569,6 +573,110 @@ class BinaryDictPageTest : public testing::Test {
std::unique_ptr<segment_v2::EncodingInfoResolver> _resolver;
};

// NOLINTNEXTLINE(readability-function-cognitive-complexity): gtest macros inflate the score.
TEST_F(BinaryDictPageTest, RejectInvalidDictionaryCodes) {
std::vector<std::string> values = {"first", "last"};
std::vector<Slice> slices;
std::vector<StringRef> 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<DataPage> decoded_page;
ASSERT_TRUE(apply_pre_decode(decoded_slice, decoded_page).ok());
ASSERT_GE(decoded_slice.size, static_cast<size_t>(BINARY_DICT_PAGE_HEADER_SIZE) +
static_cast<size_t>(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<uint32_t>(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<const ColumnString&>(*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<uint32_t>(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<uint32_t>(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<ErrorCode::CORRUPTION>()) << status;
EXPECT_EQ(2, count);
EXPECT_EQ(original_output, output.get());
EXPECT_EQ(1, output->size());
EXPECT_EQ(existing_code, assert_cast<const ColumnDictI32&>(*output).get_data()[0]);
EXPECT_EQ(0, decoder.current_index());
};

constexpr size_t second_code_offset = static_cast<size_t>(BINARY_DICT_PAGE_HEADER_SIZE) +
static_cast<size_t>(BITSHUFFLE_PAGE_HEADER_SIZE) +
sizeof(int32_t);
for (int32_t invalid_code : {-1, static_cast<int32_t>(dictionary.size())}) {
encode_fixed32_le(reinterpret_cast<uint8_t*>(decoded_slice.data + second_code_offset),
static_cast<uint32_t>(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<Slice> slices;
Expand Down
Loading