Skip to content

Harden LOUDS trie topology validation #14905

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
zaidoon1 wants to merge 1 commit into
base: main
Choose a base branch
from
+711 −28
Open

Harden LOUDS trie topology validation #14905

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
165 changes: 139 additions & 26 deletions utilities/trie_index/louds_trie.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1069,6 +1069,10 @@ Status LoudsTrie::InitFromData(const Slice& data) {
if (max_depth_ > kMaxReasonableDepth) {
return Status::Corruption("Trie index: max_depth exceeds reasonable limit");
}
if (static_cast<uint64_t>(cutoff_level_) >
static_cast<uint64_t>(max_depth_) + 1) {
return Status::Corruption("Trie index: cutoff_level exceeds max_depth");
}

memcpy(&dense_leaf_count_, p, 8);
p += 8;
Expand All @@ -1080,6 +1084,17 @@ Status LoudsTrie::InitFromData(const Slice& data) {
p += 8;
remaining -= 8;

// d_labels_ stores exactly 256 bits per dense node, and Bitvector rejects
// serialized bitvectors whose logical size exceeds UINT32_MAX bits because
// its rank LUT entries are uint32_t. Check the matching dense-node bound
// explicitly before using this header field to bound validation work.
static constexpr uint64_t kMaxReasonableDenseNodes =
std::numeric_limits<uint32_t>::max() / 256;
if (dense_node_count_ > kMaxReasonableDenseNodes) {
return Status::Corruption(
"Trie index: dense_node_count exceeds reasonable limit");
}

// Read flags field.
uint32_t flags;
memcpy(&flags, p, 4);
Expand Down Expand Up @@ -1120,26 +1135,72 @@ Status LoudsTrie::InitFromData(const Slice& data) {
// ordinal computation during traversal. Inconsistent values would cause
// incorrect rank arithmetic leading to wrong leaf indices -> wrong block
// handles.
if (dense_node_count_ > 0) {
// Each dense node uses 256 bits in d_labels_.
if (d_labels_.NumBits() != dense_node_count_ * 256) {
return Status::Corruption(
"Trie index: d_labels size inconsistent with dense_node_count");
}
// d_has_child_ has one bit per set bit in d_labels_ (one per label).
if (d_has_child_.NumBits() != d_labels_.NumOnes()) {
return Status::Corruption(
"Trie index: d_has_child size inconsistent with d_labels");
if (dense_node_count_ > std::numeric_limits<uint64_t>::max() / 256 ||
d_labels_.NumBits() != dense_node_count_ * 256) {
return Status::Corruption(
"Trie index: d_labels size inconsistent with dense_node_count");
}
// d_has_child_ has one bit per set bit in d_labels_ (one per label).
if (d_has_child_.NumBits() != d_labels_.NumOnes()) {
return Status::Corruption(
"Trie index: d_has_child size inconsistent with d_labels");
}
// d_is_prefix_key_ has one bit per dense node.
if (d_is_prefix_key_.NumBits() != dense_node_count_) {
return Status::Corruption(
"Trie index: d_is_prefix_key size inconsistent with "
"dense_node_count");
}
uint64_t computed_dense_leaf_count =
d_is_prefix_key_.NumOnes() +
(d_labels_.NumOnes() - d_has_child_.NumOnes());
if (dense_leaf_count_ != computed_dense_leaf_count) {
return Status::Corruption(
"Trie index: dense_leaf_count inconsistent with dense topology");
}
if (dense_leaf_count_ > num_keys_) {
return Status::Corruption("Trie index: dense_leaf_count exceeds num_keys");
}
if (cutoff_level_ > 0 && dense_node_count_ == 0) {
return Status::Corruption(
"Trie index: cutoff_level inconsistent with dense_node_count");
}

uint64_t expected_dense_child_count = 0;
if (cutoff_level_ == 0) {
expected_dense_child_count = num_keys_ == 0 ? 0 : 1;
} else {
uint64_t level_start_node = 0;
uint64_t level_node_count = dense_node_count_ == 0 ? 0 : 1;
uint64_t dense_nodes_seen = 0;
for (uint32_t level = 0; level < cutoff_level_; level++) {
if (level_start_node + level_node_count > dense_node_count_) {
return Status::Corruption(
"Trie index: dense_node_count inconsistent with dense topology");
}
uint64_t internal_children = 0;
for (uint64_t i = 0; i < level_node_count; i++) {
uint64_t node_num = level_start_node + i;
uint64_t label_begin = d_labels_.Rank1(node_num * 256);
uint64_t label_end = d_labels_.Rank1((node_num + 1) * 256);
internal_children +=
d_has_child_.Rank1(label_end) - d_has_child_.Rank1(label_begin);
}
dense_nodes_seen += level_node_count;
if (level + 1 == cutoff_level_) {
expected_dense_child_count = internal_children;
}
level_start_node += level_node_count;
level_node_count = internal_children;
}
// d_is_prefix_key_ has one bit per dense node.
if (d_is_prefix_key_.NumBits() != dense_node_count_) {
if (dense_nodes_seen != dense_node_count_) {
return Status::Corruption(
"Trie index: d_is_prefix_key size inconsistent with "
"dense_node_count");
"Trie index: dense_node_count inconsistent with dense topology");
}
}
if (dense_leaf_count_ > num_keys_) {
return Status::Corruption("Trie index: dense_leaf_count exceeds num_keys");
if (dense_child_count_ != expected_dense_child_count) {
return Status::Corruption(
"Trie index: dense_child_count inconsistent with dense topology");
}

// Sparse section.
Expand Down Expand Up @@ -1190,15 +1251,30 @@ Status LoudsTrie::InitFromData(const Slice& data) {
// label has exactly one bit in s_has_child_ and one bit in s_louds_.
// A mismatch means the serialized data is inconsistent, which would cause
// GetBit/Rank1 OOB reads during traversal.
if (s_labels_size_ > 0) {
if (s_has_child_.NumBits() != s_labels_size_) {
return Status::Corruption(
"Trie index: s_has_child size inconsistent with s_labels_size");
}
if (s_louds_.NumBits() != s_labels_size_) {
return Status::Corruption(
"Trie index: s_louds size inconsistent with s_labels_size");
}
if (s_has_child_.NumBits() != s_labels_size_) {
return Status::Corruption(
"Trie index: s_has_child size inconsistent with s_labels_size");
}
if (s_louds_.NumBits() != s_labels_size_) {
return Status::Corruption(
"Trie index: s_louds size inconsistent with s_labels_size");
}
if (s_labels_size_ > 0 && !s_louds_.GetBit(0)) {
return Status::Corruption("Trie index: sparse LOUDS root is missing");
}
if (s_is_prefix_key_.NumBits() != s_louds_.NumOnes()) {
return Status::Corruption(
"Trie index: s_is_prefix_key size inconsistent with sparse nodes");
}
if (dense_child_count_ > s_louds_.NumOnes()) {
return Status::Corruption(
"Trie index: dense_child_count exceeds sparse node count");
}
uint64_t computed_sparse_leaf_count =
s_is_prefix_key_.NumOnes() + (s_labels_size_ - s_has_child_.NumOnes());
if (dense_leaf_count_ + computed_sparse_leaf_count != num_keys_) {
return Status::Corruption(
"Trie index: num_keys inconsistent with trie topology");
}

// Child position lookup tables for Select-free sparse traversal.
Expand All @@ -1219,6 +1295,18 @@ Status LoudsTrie::InitFromData(const Slice& data) {
return Status::Corruption(
"Trie index: num_internal exceeds reasonable limit");
}
if (num_internal != s_has_child_.NumOnes()) {
return Status::Corruption(
"Trie index: child position table count inconsistent with sparse "
"topology");
}
if (dense_child_count_ >
std::numeric_limits<uint64_t>::max() - num_internal ||
s_louds_.NumOnes() != dense_child_count_ + num_internal) {
return Status::Corruption(
"Trie index: sparse node count inconsistent with child position "
"table");
}

if (num_internal > 0) {
// Read s_child_start_pos_
Expand Down Expand Up @@ -1251,10 +1339,26 @@ Status LoudsTrie::InitFromData(const Slice& data) {
for (uint64_t k = 0; k < num_internal; k++) {
if (s_child_start_pos_[k] >= s_labels_size_ ||
s_child_end_pos_[k] > s_labels_size_ ||
s_child_end_pos_[k] < s_child_start_pos_[k]) {
s_child_end_pos_[k] <= s_child_start_pos_[k]) {
return Status::Corruption(
"Trie index: child position out of range for sparse labels");
}
if (!s_louds_.GetBit(s_child_start_pos_[k])) {
return Status::Corruption(
"Trie index: child position does not start a sparse node");
}
uint64_t child_node = dense_child_count_ + k;
uint64_t expected_start = s_louds_.FindNthOneBit(child_node);
uint64_t expected_end = s_louds_.NextSetBit(expected_start + 1);
if (expected_end > s_labels_size_) {
expected_end = s_labels_size_;
}
if (s_child_start_pos_[k] != expected_start ||
s_child_end_pos_[k] != expected_end) {
return Status::Corruption(
"Trie index: child position table inconsistent with sparse "
"topology");
}
}
}
}
Expand All @@ -1269,6 +1373,10 @@ Status LoudsTrie::InitFromData(const Slice& data) {
memcpy(&num_chains, p, 8);
p += 8;
remaining -= 8;
if (num_chains > num_internal) {
return Status::Corruption(
"Trie index: chain count exceeds sparse internal node count");
}

if (num_chains > 0) {
// Read chain bitmap.
Expand All @@ -1278,6 +1386,11 @@ Status LoudsTrie::InitFromData(const Slice& data) {
}
p += consumed;
remaining -= consumed;
if (s_chain_bitmap_.NumBits() != num_internal ||
s_chain_bitmap_.NumOnes() != num_chains) {
return Status::Corruption(
"Trie index: chain bitmap inconsistent with chain count");
}

// Read compact chain_offsets (uint32_t per chain).
size_t offsets_bytes = num_chains * sizeof(uint32_t);
Expand Down
Loading
Loading