Skip to content

Commit f35b09d

Browse files
zaidoon1zaidoon1
committed
Harden LOUDS trie topology validation
Validate LOUDS trie topology-derived leaf counts while loading serialized trie data. Reject malformed blocks whose dense/sparse topology does not match the declared key count, child table count, or chain bitmap count before handle and seqno arrays are trusted. This keeps validation on the cold InitFromData path and preserves the on-disk format, while making malformed custom index blocks fail closed with Corruption instead of allowing release builds to index handle/seqno arrays with unchecked topology-derived leaf ordinals.
1 parent 7e2223f commit f35b09d

2 files changed

Lines changed: 287 additions & 26 deletions

File tree

‎utilities/trie_index/louds_trie.cc‎

Lines changed: 106 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -1060,6 +1060,11 @@ Status LoudsTrie::InitFromData(const Slice& data) {
10601060
p += 4;
10611061
remaining -= 4;
10621062

1063+
if (static_cast<uint64_t>(cutoff_level_) >
1064+
static_cast<uint64_t>(max_depth_) + 1) {
1065+
return Status::Corruption("Trie index: cutoff_level exceeds max_depth");
1066+
}
1067+
10631068
// Validate max_depth_ from untrusted data. The iterator allocates a
10641069
// key_buf_ of size MaxDepth()+1; if max_depth_ == UINT32_MAX, the +1
10651070
// overflows uint32_t to 0, causing a zero-length allocation and subsequent
@@ -1120,26 +1125,68 @@ Status LoudsTrie::InitFromData(const Slice& data) {
11201125
// ordinal computation during traversal. Inconsistent values would cause
11211126
// incorrect rank arithmetic leading to wrong leaf indices -> wrong block
11221127
// handles.
1123-
if (dense_node_count_ > 0) {
1124-
// Each dense node uses 256 bits in d_labels_.
1125-
if (d_labels_.NumBits() != dense_node_count_ * 256) {
1126-
return Status::Corruption(
1127-
"Trie index: d_labels size inconsistent with dense_node_count");
1128-
}
1129-
// d_has_child_ has one bit per set bit in d_labels_ (one per label).
1130-
if (d_has_child_.NumBits() != d_labels_.NumOnes()) {
1131-
return Status::Corruption(
1132-
"Trie index: d_has_child size inconsistent with d_labels");
1128+
if (dense_node_count_ > std::numeric_limits<uint64_t>::max() / 256 ||
1129+
d_labels_.NumBits() != dense_node_count_ * 256) {
1130+
return Status::Corruption(
1131+
"Trie index: d_labels size inconsistent with dense_node_count");
1132+
}
1133+
// d_has_child_ has one bit per set bit in d_labels_ (one per label).
1134+
if (d_has_child_.NumBits() != d_labels_.NumOnes()) {
1135+
return Status::Corruption(
1136+
"Trie index: d_has_child size inconsistent with d_labels");
1137+
}
1138+
// d_is_prefix_key_ has one bit per dense node.
1139+
if (d_is_prefix_key_.NumBits() != dense_node_count_) {
1140+
return Status::Corruption(
1141+
"Trie index: d_is_prefix_key size inconsistent with "
1142+
"dense_node_count");
1143+
}
1144+
uint64_t computed_dense_leaf_count =
1145+
d_is_prefix_key_.NumOnes() +
1146+
(d_labels_.NumOnes() - d_has_child_.NumOnes());
1147+
if (dense_leaf_count_ != computed_dense_leaf_count) {
1148+
return Status::Corruption(
1149+
"Trie index: dense_leaf_count inconsistent with dense topology");
1150+
}
1151+
if (dense_leaf_count_ > num_keys_) {
1152+
return Status::Corruption("Trie index: dense_leaf_count exceeds num_keys");
1153+
}
1154+
1155+
uint64_t expected_dense_child_count = 0;
1156+
if (cutoff_level_ == 0) {
1157+
expected_dense_child_count = num_keys_ == 0 ? 0 : 1;
1158+
} else {
1159+
uint64_t level_start_node = 0;
1160+
uint64_t level_node_count = dense_node_count_ == 0 ? 0 : 1;
1161+
uint64_t dense_nodes_seen = 0;
1162+
for (uint32_t level = 0; level < cutoff_level_; level++) {
1163+
if (level_start_node + level_node_count > dense_node_count_) {
1164+
return Status::Corruption(
1165+
"Trie index: dense_node_count inconsistent with dense topology");
1166+
}
1167+
uint64_t internal_children = 0;
1168+
for (uint64_t i = 0; i < level_node_count; i++) {
1169+
uint64_t node_num = level_start_node + i;
1170+
uint64_t label_begin = d_labels_.Rank1(node_num * 256);
1171+
uint64_t label_end = d_labels_.Rank1((node_num + 1) * 256);
1172+
internal_children +=
1173+
d_has_child_.Rank1(label_end) - d_has_child_.Rank1(label_begin);
1174+
}
1175+
dense_nodes_seen += level_node_count;
1176+
if (level + 1 == cutoff_level_) {
1177+
expected_dense_child_count = internal_children;
1178+
}
1179+
level_start_node += level_node_count;
1180+
level_node_count = internal_children;
11331181
}
1134-
// d_is_prefix_key_ has one bit per dense node.
1135-
if (d_is_prefix_key_.NumBits() != dense_node_count_) {
1182+
if (dense_nodes_seen != dense_node_count_) {
11361183
return Status::Corruption(
1137-
"Trie index: d_is_prefix_key size inconsistent with "
1138-
"dense_node_count");
1184+
"Trie index: dense_node_count inconsistent with dense topology");
11391185
}
11401186
}
1141-
if (dense_leaf_count_ > num_keys_) {
1142-
return Status::Corruption("Trie index: dense_leaf_count exceeds num_keys");
1187+
if (dense_child_count_ != expected_dense_child_count) {
1188+
return Status::Corruption(
1189+
"Trie index: dense_child_count inconsistent with dense topology");
11431190
}
11441191

11451192
// Sparse section.
@@ -1190,15 +1237,30 @@ Status LoudsTrie::InitFromData(const Slice& data) {
11901237
// label has exactly one bit in s_has_child_ and one bit in s_louds_.
11911238
// A mismatch means the serialized data is inconsistent, which would cause
11921239
// GetBit/Rank1 OOB reads during traversal.
1193-
if (s_labels_size_ > 0) {
1194-
if (s_has_child_.NumBits() != s_labels_size_) {
1195-
return Status::Corruption(
1196-
"Trie index: s_has_child size inconsistent with s_labels_size");
1197-
}
1198-
if (s_louds_.NumBits() != s_labels_size_) {
1199-
return Status::Corruption(
1200-
"Trie index: s_louds size inconsistent with s_labels_size");
1201-
}
1240+
if (s_has_child_.NumBits() != s_labels_size_) {
1241+
return Status::Corruption(
1242+
"Trie index: s_has_child size inconsistent with s_labels_size");
1243+
}
1244+
if (s_louds_.NumBits() != s_labels_size_) {
1245+
return Status::Corruption(
1246+
"Trie index: s_louds size inconsistent with s_labels_size");
1247+
}
1248+
if (s_labels_size_ > 0 && !s_louds_.GetBit(0)) {
1249+
return Status::Corruption("Trie index: sparse LOUDS root is missing");
1250+
}
1251+
if (s_is_prefix_key_.NumBits() != s_louds_.NumOnes()) {
1252+
return Status::Corruption(
1253+
"Trie index: s_is_prefix_key size inconsistent with sparse nodes");
1254+
}
1255+
if (dense_child_count_ > s_louds_.NumOnes()) {
1256+
return Status::Corruption(
1257+
"Trie index: dense_child_count exceeds sparse node count");
1258+
}
1259+
uint64_t computed_sparse_leaf_count =
1260+
s_is_prefix_key_.NumOnes() + (s_labels_size_ - s_has_child_.NumOnes());
1261+
if (dense_leaf_count_ + computed_sparse_leaf_count != num_keys_) {
1262+
return Status::Corruption(
1263+
"Trie index: num_keys inconsistent with trie topology");
12021264
}
12031265

12041266
// Child position lookup tables for Select-free sparse traversal.
@@ -1219,6 +1281,11 @@ Status LoudsTrie::InitFromData(const Slice& data) {
12191281
return Status::Corruption(
12201282
"Trie index: num_internal exceeds reasonable limit");
12211283
}
1284+
if (num_internal != s_has_child_.NumOnes()) {
1285+
return Status::Corruption(
1286+
"Trie index: child position table count inconsistent with sparse "
1287+
"topology");
1288+
}
12221289

12231290
if (num_internal > 0) {
12241291
// Read s_child_start_pos_
@@ -1251,10 +1318,14 @@ Status LoudsTrie::InitFromData(const Slice& data) {
12511318
for (uint64_t k = 0; k < num_internal; k++) {
12521319
if (s_child_start_pos_[k] >= s_labels_size_ ||
12531320
s_child_end_pos_[k] > s_labels_size_ ||
1254-
s_child_end_pos_[k] < s_child_start_pos_[k]) {
1321+
s_child_end_pos_[k] <= s_child_start_pos_[k]) {
12551322
return Status::Corruption(
12561323
"Trie index: child position out of range for sparse labels");
12571324
}
1325+
if (!s_louds_.GetBit(s_child_start_pos_[k])) {
1326+
return Status::Corruption(
1327+
"Trie index: child position does not start a sparse node");
1328+
}
12581329
}
12591330
}
12601331
}
@@ -1269,6 +1340,10 @@ Status LoudsTrie::InitFromData(const Slice& data) {
12691340
memcpy(&num_chains, p, 8);
12701341
p += 8;
12711342
remaining -= 8;
1343+
if (num_chains > num_internal) {
1344+
return Status::Corruption(
1345+
"Trie index: chain count exceeds sparse internal node count");
1346+
}
12721347

12731348
if (num_chains > 0) {
12741349
// Read chain bitmap.
@@ -1278,6 +1353,11 @@ Status LoudsTrie::InitFromData(const Slice& data) {
12781353
}
12791354
p += consumed;
12801355
remaining -= consumed;
1356+
if (s_chain_bitmap_.NumBits() != num_internal ||
1357+
s_chain_bitmap_.NumOnes() != num_chains) {
1358+
return Status::Corruption(
1359+
"Trie index: chain bitmap inconsistent with chain count");
1360+
}
12811361

12821362
// Read compact chain_offsets (uint32_t per chain).
12831363
size_t offsets_bytes = num_chains * sizeof(uint32_t);

‎utilities/trie_index/trie_index_test.cc‎

Lines changed: 181 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@
88
#include <cinttypes>
99
#include <cmath>
1010
#include <cstdint>
11+
#include <cstring>
1112
#include <memory>
1213
#include <random>
1314
#include <string>
@@ -1103,6 +1104,88 @@ class LoudsTrieTest : public testing::Test {
11031104
}
11041105
};
11051106

1107+
static void PatchFixed64ForTest(std::string* data, size_t offset,
1108+
uint64_t value) {
1109+
ASSERT_LE(offset + sizeof(value), data->size());
1110+
memcpy(&(*data)[offset], &value, sizeof(value));
1111+
}
1112+
1113+
static uint64_t ReadFixed64ForTest(const std::string& data, size_t offset) {
1114+
assert(offset + sizeof(uint64_t) <= data.size());
1115+
uint64_t value;
1116+
memcpy(&value, data.data() + offset, sizeof(value));
1117+
return value;
1118+
}
1119+
1120+
static Status SkipBitvectorForTest(const std::string& data, size_t* offset) {
1121+
if (*offset > data.size()) {
1122+
return Status::Corruption("offset past end");
1123+
}
1124+
Bitvector bv;
1125+
size_t consumed = 0;
1126+
Status s =
1127+
bv.InitFromData(data.data() + *offset, data.size() - *offset, &consumed);
1128+
if (!s.ok()) {
1129+
return s;
1130+
}
1131+
*offset += consumed;
1132+
return Status::OK();
1133+
}
1134+
1135+
static Status FindChildCountOffsetForTest(const std::string& data,
1136+
size_t* child_count_offset) {
1137+
size_t offset = 56;
1138+
for (int i = 0; i < 3; i++) {
1139+
Status s = SkipBitvectorForTest(data, &offset);
1140+
if (!s.ok()) {
1141+
return s;
1142+
}
1143+
}
1144+
if (offset + sizeof(uint64_t) > data.size()) {
1145+
return Status::Corruption("truncated sparse labels size");
1146+
}
1147+
uint64_t sparse_labels_size = ReadFixed64ForTest(data, offset);
1148+
offset += sizeof(uint64_t);
1149+
size_t sparse_labels_padded = (sparse_labels_size + 7) & ~size_t(7);
1150+
if (offset + sparse_labels_padded > data.size()) {
1151+
return Status::Corruption("truncated sparse labels");
1152+
}
1153+
offset += sparse_labels_padded;
1154+
for (int i = 0; i < 3; i++) {
1155+
Status s = SkipBitvectorForTest(data, &offset);
1156+
if (!s.ok()) {
1157+
return s;
1158+
}
1159+
}
1160+
if (offset + sizeof(uint64_t) > data.size()) {
1161+
return Status::Corruption("truncated child count");
1162+
}
1163+
*child_count_offset = offset;
1164+
return Status::OK();
1165+
}
1166+
1167+
static Status FindChainCountOffsetForTest(const std::string& data,
1168+
size_t* chain_count_offset) {
1169+
size_t offset = 0;
1170+
Status s = FindChildCountOffsetForTest(data, &offset);
1171+
if (!s.ok()) {
1172+
return s;
1173+
}
1174+
uint64_t num_internal = ReadFixed64ForTest(data, offset);
1175+
offset += sizeof(uint64_t);
1176+
size_t child_table_bytes = num_internal * sizeof(uint32_t);
1177+
size_t child_table_padded = (child_table_bytes + 7) & ~size_t(7);
1178+
if (offset + child_table_padded + child_table_padded > data.size()) {
1179+
return Status::Corruption("truncated child tables");
1180+
}
1181+
offset += child_table_padded + child_table_padded;
1182+
if (offset + sizeof(uint64_t) > data.size()) {
1183+
return Status::Corruption("truncated chain count");
1184+
}
1185+
*chain_count_offset = offset;
1186+
return Status::OK();
1187+
}
1188+
11061189
// Helper: build a trie from sorted keys and verify that Seek and Next
11071190
// iterate in exactly the expected order, checking both handles and
11081191
// reconstructed keys.
@@ -1830,6 +1913,104 @@ TEST_F(LoudsTrieTest, InitFromDataMaxDepthCorruption) {
18301913
ASSERT_TRUE(trie3.InitFromData(Slice(header2)).IsCorruption());
18311914
}
18321915

1916+
TEST_F(LoudsTrieTest, InitFromDataRejectsLeafCountMismatch) {
1917+
std::vector<std::string> keys = {"alpha", "beta", "delta", "gamma"};
1918+
auto bt = BuildTrieFromKeys(keys);
1919+
std::string data(bt.builder.GetSerializedData().data(),
1920+
bt.builder.GetSerializedData().size());
1921+
1922+
PatchFixed64ForTest(&data, /*num_keys offset=*/8, keys.size() - 1);
1923+
1924+
LoudsTrie trie;
1925+
Status s = trie.InitFromData(Slice(data));
1926+
ASSERT_TRUE(s.IsCorruption()) << s.ToString();
1927+
}
1928+
1929+
TEST_F(LoudsTrieTest, InitFromDataRejectsDenseLeafCountMismatch) {
1930+
std::vector<std::string> keys = {"alpha", "beta", "delta", "gamma"};
1931+
auto bt = BuildTrieFromKeys(keys);
1932+
std::string data(bt.builder.GetSerializedData().data(),
1933+
bt.builder.GetSerializedData().size());
1934+
uint64_t dense_leaf_count = ReadFixed64ForTest(data, 24);
1935+
1936+
PatchFixed64ForTest(&data, /*dense_leaf_count offset=*/24,
1937+
dense_leaf_count + 1);
1938+
1939+
LoudsTrie trie;
1940+
Status s = trie.InitFromData(Slice(data));
1941+
ASSERT_TRUE(s.IsCorruption()) << s.ToString();
1942+
}
1943+
1944+
TEST_F(LoudsTrieTest, InitFromDataRejectsSeqnoLeafCountMismatch) {
1945+
LoudsTrieBuilder builder;
1946+
builder.SetHasSeqnoEncoding(true);
1947+
std::vector<std::string> keys = {"alpha", "beta", "delta", "gamma"};
1948+
for (size_t i = 0; i < keys.size(); i++) {
1949+
builder.AddKeyWithSeqno(Slice(keys[i]), TrieBlockHandle{i * 100, 50},
1950+
/*seqno=*/i + 1, /*block_count=*/1);
1951+
}
1952+
builder.Finish();
1953+
std::string data(builder.GetSerializedData().data(),
1954+
builder.GetSerializedData().size());
1955+
1956+
PatchFixed64ForTest(&data, /*num_keys offset=*/8, keys.size() - 1);
1957+
1958+
LoudsTrie trie;
1959+
Status s = trie.InitFromData(Slice(data));
1960+
ASSERT_TRUE(s.IsCorruption()) << s.ToString();
1961+
}
1962+
1963+
TEST_F(LoudsTrieTest, InitFromDataRejectsSparseChildCountMismatch) {
1964+
std::vector<std::string> keys;
1965+
for (int i = 0; i < 20; i++) {
1966+
std::string key(1, static_cast<char>('A' + i));
1967+
key += 'x';
1968+
keys.push_back(key);
1969+
}
1970+
auto bt = BuildTrieFromKeys(keys);
1971+
std::string data(bt.builder.GetSerializedData().data(),
1972+
bt.builder.GetSerializedData().size());
1973+
size_t child_count_offset = 0;
1974+
ASSERT_OK(FindChildCountOffsetForTest(data, &child_count_offset));
1975+
uint64_t num_internal = ReadFixed64ForTest(data, child_count_offset);
1976+
ASSERT_GT(num_internal, 0u);
1977+
1978+
PatchFixed64ForTest(&data, child_count_offset, num_internal - 1);
1979+
1980+
LoudsTrie trie;
1981+
Status s = trie.InitFromData(Slice(data));
1982+
ASSERT_TRUE(s.IsCorruption()) << s.ToString();
1983+
}
1984+
1985+
TEST_F(LoudsTrieTest, InitFromDataRejectsChainBitmapCountMismatch) {
1986+
std::vector<std::string> keys;
1987+
std::string prefix(10, 'A');
1988+
for (int i = 0; i < 20; i++) {
1989+
char suffix[8];
1990+
snprintf(suffix, sizeof(suffix), "%02d", i);
1991+
keys.push_back(prefix + suffix);
1992+
}
1993+
std::sort(keys.begin(), keys.end());
1994+
auto bt = BuildTrieFromKeys(keys);
1995+
ASSERT_TRUE(bt.trie.HasChains());
1996+
std::string data(bt.builder.GetSerializedData().data(),
1997+
bt.builder.GetSerializedData().size());
1998+
size_t child_count_offset = 0;
1999+
ASSERT_OK(FindChildCountOffsetForTest(data, &child_count_offset));
2000+
uint64_t num_internal = ReadFixed64ForTest(data, child_count_offset);
2001+
size_t chain_count_offset = 0;
2002+
ASSERT_OK(FindChainCountOffsetForTest(data, &chain_count_offset));
2003+
uint64_t num_chains = ReadFixed64ForTest(data, chain_count_offset);
2004+
ASSERT_GT(num_chains, 0u);
2005+
ASSERT_LT(num_chains, num_internal);
2006+
2007+
PatchFixed64ForTest(&data, chain_count_offset, num_chains + 1);
2008+
2009+
LoudsTrie trie;
2010+
Status s = trie.InitFromData(Slice(data));
2011+
ASSERT_TRUE(s.IsCorruption()) << s.ToString();
2012+
}
2013+
18332014
TEST_F(LoudsTrieTest, MoveConstructor) {
18342015
// Verify that move constructor works correctly for LoudsTrie, which
18352016
// contains Bitvector members with RecomputePointers() logic.

0 commit comments

Comments
 (0)