diff --git a/src/overlaybd/lsmt/file.cpp b/src/overlaybd/lsmt/file.cpp index 0bac3181..0669478c 100644 --- a/src/overlaybd/lsmt/file.cpp +++ b/src/overlaybd/lsmt/file.cpp @@ -1029,6 +1029,9 @@ class LSMTFile : public LSMTReadOnlyFile { virtual int flatten(IFile *as) override { unique_ptr pmi((IComboIndex*)(m_index->make_read_only_index())); + if (!pmi) + LOG_ERROR_RETURN(0, -1, "failed to make read only index."); + CommitArgs args(as); atomic_uint64_t _no_use_var(0); CompactOptions opts(&m_files, (SegmentMapping*)(pmi->buffer()), pmi->size(), m_vsize, &args); @@ -1350,6 +1353,10 @@ static HeaderTrailer *verify_ht(IFile *file, char *buf, bool is_trailer, ssize_t LOG_ERROR_RETURN(0, nullptr, "trailer magic, trailer type, " "file type or sealedness doesn't match"); + + if (pht->index_size > MAX_LSMT_RO_INDEX_SIZE) + LOG_ERROR_RETURN(0, nullptr, "LSMT RO index size ` exceeds maximum `", + pht->index_size + 0, MAX_LSMT_RO_INDEX_SIZE); return pht; } @@ -1387,6 +1394,10 @@ static SegmentMapping *do_load_index(IFile *file, HeaderTrailer *pheader_trailer LOG_ERROR_RETURN(0, nullptr, "index offset wrong"); index_bytes = stat.st_size - HeaderTrailer::SPACE; pht->index_size = index_bytes / sizeof(SegmentMapping); + + if (pht->index_size > MAX_LSMT_INDEX_SIZE) + LOG_ERROR_RETURN(0, nullptr, "LSMT RW index size ` exceeds maximum `", + pht->index_size + 0, MAX_LSMT_INDEX_SIZE); } SegmentMapping *ibuf = nullptr; diff --git a/src/overlaybd/lsmt/index.cpp b/src/overlaybd/lsmt/index.cpp index 7c1c59ec..803698c9 100644 --- a/src/overlaybd/lsmt/index.cpp +++ b/src/overlaybd/lsmt/index.cpp @@ -622,9 +622,10 @@ class Index0 : public IComboIndex { UNIMPLEMENTED(int commit_index0() override); }; -static void merge_indexes(uint8_t level, vector &mapping, const Index **pindexes, +static bool merge_indexes(uint8_t level, vector &mapping, const Index **pindexes, std::size_t n, uint64_t begin, uint64_t end, bool change_tag = true, - size_t max_level = 0); + size_t max_level = 0, + size_t max_index_size = MAX_LSMT_INDEX_SIZE); class ComboIndex : public Index0 { public: @@ -740,7 +741,8 @@ class ComboIndex : public Index0 { virtual Index *rebuild_backing_index(Index *highlevel_idx, size_t max_level) { vector mappings; const Index *indexes[2] = {highlevel_idx, const_cast(m_backing_index)}; - merge_indexes(0, mappings, indexes, 2, 0, UINT64_MAX, false, max_level); + if (!merge_indexes(0, mappings, indexes, 2, 0, UINT64_MAX, false, max_level)) + return nullptr; return new Index(std::move(mappings)); } @@ -753,7 +755,10 @@ class ComboIndex : public Index0 { return ro_idx0; } const Index *indexes[2] = {ro_idx0, const_cast(m_backing_index)}; - merge_indexes(0, mappings, indexes, 2, 0, UINT64_MAX, false, 2); + if (!merge_indexes(0, mappings, indexes, 2, 0, UINT64_MAX, false, 2)) { + delete ro_idx0; + return nullptr; + } delete ro_idx0; return new Index(std::move(mappings)); } @@ -839,35 +844,41 @@ IMemoryIndex *create_level_index(const SegmentMapping *pmappings, size_t n, uint return (ok1 && ok2) ? new LevelIndex(pmappings, n, copy_mode) : nullptr; } -static void merge_indexes(uint8_t level, vector &mapping, const Index **pindexes, +static bool merge_indexes(uint8_t level, vector &mapping, const Index **pindexes, size_t n, uint64_t begin, uint64_t end, bool change_tag, - size_t max_level) { + size_t max_level, size_t max_index_size) { if (pindexes == nullptr) - return; + return true; if (change_tag) { if (n == 0) - return; + return true; } else { if (max_level == 0) - return; + return true; } if (begin >= end) - return; + return true; auto begin0 = begin; auto size0 = mapping.size(); auto pi0 = pindexes[0]; for (auto it = pi0->lower_bound(begin); it != pi0->end() && it->offset < end; ++it) { if (it->offset > begin) { - if (change_tag) - merge_indexes(level + 1, mapping, pindexes + 1, n - 1, begin, it->offset); - else { + if (change_tag) { + if (!merge_indexes(level + 1, mapping, pindexes + 1, n - 1, begin, it->offset, + true, max_level, max_index_size)) + return false; + } else { int k = (n <= 1 ? 0 : 1); - merge_indexes(level + 1, mapping, pindexes + k, 0, begin, it->offset, false, - max_level - 1); + if (!merge_indexes(level + 1, mapping, pindexes + k, 0, begin, it->offset, false, + max_level - 1, max_index_size)) + return false; } } + if (mapping.size() >= max_index_size) + LOG_ERROR_RETURN(0, false, "Merged LSMT index size ` exceeds maximum `", + mapping.size() + 1, max_index_size); mapping.push_back(*it); if (change_tag) { @@ -876,11 +887,15 @@ static void merge_indexes(uint8_t level, vector &mapping, const begin = it->end(); } if (begin < end) { - if (change_tag) - merge_indexes(level + 1, mapping, pindexes + 1, n - 1, begin, end); - else { + if (change_tag) { + if (!merge_indexes(level + 1, mapping, pindexes + 1, n - 1, begin, end, + true, max_level, max_index_size)) + return false; + } else { int k = (n <= 1 ? 0 : 1); - merge_indexes(level + 1, mapping, pindexes + k, 0, begin, end, false, max_level - 1); + if (!merge_indexes(level + 1, mapping, pindexes + k, 0, begin, end, false, + max_level - 1, max_index_size)) + return false; } } if (mapping.size() > size0) { @@ -889,6 +904,8 @@ static void merge_indexes(uint8_t level, vector &mapping, const if (mapping.back().end() > end) mapping.back().backward_end_to(end); } + + return true; } IComboIndex *create_combo_index(IMemoryIndex0 *index0, const IMemoryIndex *index, @@ -952,7 +969,8 @@ IMemoryIndex *merge_memory_indexes(const IMemoryIndex **pindexes, size_t n) { vector mapping; auto pi = (const Index **)pindexes; mapping.reserve(pi[0]->size()); - merge_indexes(0, mapping, pi, n, 0, UINT64_MAX); + if (!merge_indexes(0, mapping, pi, n, 0, UINT64_MAX)) + return nullptr; if (pindexes[0]->vsize() < static_cast(UINT32_MAX) * ALIGNMENT && mapping.size() < NODES_PER_LEVEL_32[MAX_LEVEL_32-1]) { diff --git a/src/overlaybd/lsmt/index.h b/src/overlaybd/lsmt/index.h index ce6931c4..22852513 100644 --- a/src/overlaybd/lsmt/index.h +++ b/src/overlaybd/lsmt/index.h @@ -29,6 +29,8 @@ IMemoryIndex -> IMemoryIndex0 -> IComboIndex -> Index0 ( set ) -> Co #include namespace LSMT { +static const uint64_t MAX_LSMT_RO_INDEX_SIZE = 1000000; +static const uint64_t MAX_LSMT_INDEX_SIZE = 128000000; struct Segment { uint64_t offset : 50; uint32_t length : 14; diff --git a/src/overlaybd/lsmt/test/test.cpp b/src/overlaybd/lsmt/test/test.cpp index 0bc85240..d08f273a 100644 --- a/src/overlaybd/lsmt/test/test.cpp +++ b/src/overlaybd/lsmt/test/test.cpp @@ -62,6 +62,53 @@ void lookup_test(const SegmentMapping (&mapping)[N1], Segment s, lookup_test(mapping, N1, s, stdrst, N2); } +TEST_F(FileTest, reject_oversized_index) { + const char *filename = "oversized-index.lsmt"; + auto file = lfs->open(filename, O_RDWR | O_CREAT | O_TRUNC, S_IRWXU); + ASSERT_NE(file, nullptr); + + LayerInfo info; + info.virtual_size = HeaderTrailer::SPACE; + ASSERT_EQ(write_header_trailer(file, true, true, true, 0, 0, info), + (int)HeaderTrailer::SPACE); + + const uint64_t index_size = MAX_LSMT_RO_INDEX_SIZE + 1; + const uint64_t index_offset = HeaderTrailer::SPACE; + const uint64_t trailer_offset = + index_offset + index_size * sizeof(SegmentMapping); + ASSERT_EQ(file->lseek(trailer_offset, SEEK_SET), (off_t)trailer_offset); + ASSERT_EQ(write_header_trailer(file, false, true, true, index_offset, + index_size, info), + (int)HeaderTrailer::SPACE); + + EXPECT_EQ(LSMT::open_file_ro(file), nullptr); + delete file; + lfs->unlink(filename); +} + +TEST_F(FileTest, reject_oversized_rw_index) { + auto rw = create_file_rw(); + ASSERT_NE(rw, nullptr); + delete rw; + + auto fdata = lfs->open(data_name.back().c_str(), O_RDWR, S_IRWXU); + auto findex = lfs->open(idx_name.back().c_str(), O_RDWR, S_IRWXU); + ASSERT_NE(fdata, nullptr); + ASSERT_NE(findex, nullptr); + + const uint64_t index_file_size = + HeaderTrailer::SPACE + + (MAX_LSMT_INDEX_SIZE + 1) * sizeof(SegmentMapping); + ASSERT_EQ(findex->ftruncate(index_file_size), 0); + + auto reopened = LSMT::open_file_rw(fdata, findex, false); + EXPECT_EQ(reopened, nullptr); + + delete reopened; + delete fdata; + delete findex; +} + void lookup_test(IMemoryIndex &idx); TEST(Index, lookup) { @@ -267,6 +314,20 @@ inline void test_merge_combo(const IMemoryIndex *indexes[], size_t ni, // num of test_combo(indexes, ni, stdrst, NR); } +TEST(Index, reject_oversized_merge) { + SegmentMapping mapping0[] = {{0, 1, 0}, {2, 1, 2}}; + SegmentMapping mapping1[] = {{1, 1, 1}, {3, 1, 3}}; + + Index index0(mapping0, LEN(mapping0), false); + Index index1(mapping1, LEN(mapping1), false); + const Index *indexes[] = {&index0, &index1}; + + vector merged; + EXPECT_FALSE(merge_indexes(0, merged, indexes, LEN(indexes), 0, UINT64_MAX, + true, 0, 3)); + EXPECT_EQ(merged.size(), 3); +} + TEST(Index, merge) { const static SegmentMapping mapping0[] = {{5, 5, 0}, {10, 10, 50}, {100, 10, 20}}; const static SegmentMapping mapping1[] = {{0, 1, 7}, {2, 4, 5}, {15, 10, 22}, diff --git a/src/overlaybd/zfile/test/test.cpp b/src/overlaybd/zfile/test/test.cpp index 4e145f3b..b60df5b4 100644 --- a/src/overlaybd/zfile/test/test.cpp +++ b/src/overlaybd/zfile/test/test.cpp @@ -133,6 +133,31 @@ class ZFileTest : public ::testing::Test { } }; +TEST_F(ZFileTest, reject_oversized_index) { + const char *filename = "oversized-index.zfile"; + auto file = lfs->open(filename, O_RDWR | O_CREAT | O_TRUNC, 0644); + ASSERT_NE(file, nullptr); + + char ht_buf[CompressionFile::HeaderTrailer::SPACE]{}; + auto ht = new (ht_buf) CompressionFile::HeaderTrailer; + CompressOptions opt; + ht->set_compress_option(opt); + ht->index_offset = CompressionFile::HeaderTrailer::SPACE; + ht->index_size = MAX_ZFILE_INDEX_SIZE + 1; + ASSERT_EQ(write_header_trailer(file, true, false, true, ht), + (int)CompressionFile::HeaderTrailer::SPACE); + + const uint64_t trailer_offset = + ht->index_offset + ht->index_size * sizeof(uint32_t); + ASSERT_EQ(file->lseek(trailer_offset, SEEK_SET), (off_t)trailer_offset); + ASSERT_EQ(write_header_trailer(file, false, true, true, ht), + (int)CompressionFile::HeaderTrailer::SPACE); + + EXPECT_EQ(zfile_open_ro(file, false), nullptr); + delete file; + lfs->unlink(filename); +} + /* testcases: checksum{disable, enable} x algorithm{lz4, zstd} x bs{4K, 8K, 16K, 32K, 64K} diff --git a/src/overlaybd/zfile/zfile.cpp b/src/overlaybd/zfile/zfile.cpp index a2a36baa..e250d079 100644 --- a/src/overlaybd/zfile/zfile.cpp +++ b/src/overlaybd/zfile/zfile.cpp @@ -47,6 +47,7 @@ uint64_t zfile_blk_cnt = 0; namespace ZFile { const static size_t BUF_SIZE = 512; +const static uint64_t MAX_ZFILE_INDEX_SIZE = 1000000000; const static uint32_t NOI_WELL_KNOWN_PRIME = 100007; const static uint8_t FLAG_VALID_FALSE = 0; const static uint8_t FLAG_VALID_TRUE = 1; @@ -1095,6 +1096,10 @@ bool load_jump_table(IFile *file, CompressionFile::HeaderTrailer *pheader_traile "trailer magic, trailer type, file type or sealedness doesn't match"); } + if (pht->index_size > MAX_ZFILE_INDEX_SIZE) + LOG_ERROR_RETURN(0, false, "ZFile index size ` exceeds maximum `", + pht->index_size + 0, MAX_ZFILE_INDEX_SIZE); + index_bytes = pht->index_size * sizeof(uint32_t); LOG_INFO("trailer_offset: `, idx_offset: `, idx_bytes: `, dict_size: `, use_dict: `", trailer_offset, pht->index_offset, index_bytes, pht->opt.dict_size, @@ -1103,6 +1108,10 @@ bool load_jump_table(IFile *file, CompressionFile::HeaderTrailer *pheader_traile if (index_bytes > trailer_offset - pht->index_offset) LOG_ERROR_RETURN(0, false, "invalid index bytes or size. "); } else { + if (pht->index_size > MAX_ZFILE_INDEX_SIZE) + LOG_ERROR_RETURN(0, false, "ZFile index size ` exceeds maximum `", + pht->index_size + 0, MAX_ZFILE_INDEX_SIZE); + index_bytes = pht->index_size * sizeof(uint32_t); LOG_INFO("read overwrite header. idx_offset: `, idx_bytes: `, dict_size: `, use_dict: `", pht->index_offset, index_bytes, pht->opt.dict_size, pht->opt.use_dict);