Validate LSMT and ZFile index sizes

src/overlaybd/lsmt/file.cpp

@@ -1367,18 +1367,16 @@ static SegmentMapping *do_load_index(IFile *file, HeaderTrailer *pheader_trailer
         LOG_ERRNO_RETURN(0, nullptr, "failed to stat file.");
     assert(trailer || pht->is_sparse_rw() == false);
     uint64_t index_bytes;
+    uint64_t trailer_offset = 0;
     if (trailer) {
         if (!pht->is_data_file())
             LOG_ERROR_RETURN(0, nullptr, "uncognized file type");
         pht = verify_ht(file, buf, true, stat.st_size);
         if (pht == nullptr) {
             return nullptr;
         }
-        auto trailer_offset = stat.st_size - HeaderTrailer::SPACE;
+        trailer_offset = stat.st_size - HeaderTrailer::SPACE;
         LOG_DEBUG("index_size: `, trailer offset: `", pht->index_size + 0, trailer_offset);
-        index_bytes = pht->index_size * sizeof(SegmentMapping);
-        if (index_bytes > trailer_offset - pht->index_offset)
-            LOG_ERROR_RETURN(0, nullptr, "invalid index bytes or size");
 
     } else {
         if (!pht->is_index_file() || pht->is_sealed())
@@ -1389,6 +1387,16 @@ static SegmentMapping *do_load_index(IFile *file, HeaderTrailer *pheader_trailer
         pht->index_size = index_bytes / sizeof(SegmentMapping);
     }
 
+    if (pht->index_size > MAX_LSMT_INDEX_SIZE)
+        LOG_ERROR_RETURN(0, nullptr, "LSMT index size ` exceeds maximum `",
+                         pht->index_size + 0, MAX_LSMT_INDEX_SIZE);
+
+    if (trailer) {
+        index_bytes = pht->index_size * sizeof(SegmentMapping);
+        if (index_bytes > trailer_offset - pht->index_offset)
+            LOG_ERROR_RETURN(0, nullptr, "invalid index bytes or size");
+    }
+
     SegmentMapping *ibuf = nullptr;
     posix_memalign((void **)&ibuf, ALIGNMENT4K, pht->index_size * sizeof(*ibuf));
     ret = file->pread(ibuf, index_bytes, pht->index_offset);

src/overlaybd/lsmt/index.cpp

@@ -954,6 +954,11 @@ IMemoryIndex *merge_memory_indexes(const IMemoryIndex **pindexes, size_t n) {
     mapping.reserve(pi[0]->size());
     merge_indexes(0, mapping, pi, n, 0, UINT64_MAX);
 
+    if (mapping.size() > MAX_LSMT_INDEX_SIZE)
+        LOG_ERROR_RETURN(0, nullptr,
+                         "Merged LSMT index size ` exceeds maximum `",
+                         mapping.size(), MAX_LSMT_INDEX_SIZE);
+
     if (pindexes[0]->vsize() < static_cast<uint64_t>(UINT32_MAX) * ALIGNMENT
         && mapping.size() < NODES_PER_LEVEL_32[MAX_LEVEL_32-1]) {
         return new_index_with_lineriazed_bptree<uint32_t>(std::move(mapping), pindexes[0]->vsize());

src/overlaybd/lsmt/index.h

@@ -29,6 +29,7 @@ IMemoryIndex -> IMemoryIndex0 -> IComboIndex -> Index0 ( set<SegmentMap> ) -> Co
 #include <sys/types.h>
 
 namespace LSMT {
+static const uint64_t MAX_LSMT_INDEX_SIZE = 1000000;
 struct Segment {          
     uint64_t offset : 50;
     uint32_t length : 14;

src/overlaybd/lsmt/test/test.cpp

@@ -62,6 +62,30 @@ void lookup_test(const SegmentMapping (&mapping)[N1], Segment s,
     lookup_test<IDX>(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_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);
+}
+
 void lookup_test(IMemoryIndex &idx);
 
 TEST(Index, lookup) {
@@ -267,6 +291,27 @@ inline void test_merge_combo(const IMemoryIndex *indexes[], size_t ni, // num of
     test_combo(indexes, ni, stdrst, NR);
 }
 
+TEST(Index, reject_oversized_merge) {
+    const size_t per_index_size = MAX_LSMT_INDEX_SIZE / 2 + 1;
+    vector<SegmentMapping> mapping0;
+    vector<SegmentMapping> mapping1;
+    mapping0.reserve(per_index_size);
+    mapping1.reserve(per_index_size);
+
+    for (size_t i = 0; i < per_index_size; ++i) {
+        mapping0.emplace_back(i * 2, 1, i * 2);
+        mapping1.emplace_back(i * 2 + 1, 1, i * 2 + 1);
+    }
+
+    Index index0(mapping0.data(), mapping0.size(), false);
+    Index index1(mapping1.data(), mapping1.size(), false);
+    const IMemoryIndex *indexes[] = {&index0, &index1};
+
+    auto merged = merge_memory_indexes(indexes, LEN(indexes));
+    EXPECT_EQ(merged, nullptr);
+    delete merged;
+}
+
 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},

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}

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;
@@ -1074,7 +1075,10 @@ bool load_jump_table(IFile *file, CompressionFile::HeaderTrailer *pheader_traile
         LOG_ERRNO_RETURN(0, false, "failed to stat file.");
     }
     uint64_t index_bytes = 0;
-    if (!pht->is_header_overwrite()) {
+    uint64_t trailer_offset = 0;
+    bool header_overwrite = pht->is_header_overwrite();
+
+    if (!header_overwrite) {
         struct stat stat;
         ret = file->fstat(&stat);
         if (ret < 0) {
@@ -1084,7 +1088,7 @@ bool load_jump_table(IFile *file, CompressionFile::HeaderTrailer *pheader_traile
             LOG_ERROR_RETURN(0, false, "uncognized file type");
         }
 
-        auto trailer_offset = stat.st_size - CompressionFile::HeaderTrailer::SPACE;
+        trailer_offset = stat.st_size - CompressionFile::HeaderTrailer::SPACE;
         ret = file->pread(buf, CompressionFile::HeaderTrailer::SPACE, trailer_offset);
         if (ret < (ssize_t)CompressionFile::HeaderTrailer::SPACE)
             LOG_ERRNO_RETURN(0, false, "failed to read file trailer.");
@@ -1094,19 +1098,26 @@ bool load_jump_table(IFile *file, CompressionFile::HeaderTrailer *pheader_traile
             LOG_ERROR_RETURN(0, false,
                              "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);
+    index_bytes = pht->index_size * sizeof(uint32_t);
+
+    if (!header_overwrite) {
         LOG_INFO("trailer_offset: `, idx_offset: `, idx_bytes: `, dict_size: `, use_dict: `",
                  trailer_offset, pht->index_offset, index_bytes, pht->opt.dict_size,
                  pht->opt.use_dict);
 
         if (index_bytes > trailer_offset - pht->index_offset)
             LOG_ERROR_RETURN(0, false, "invalid index bytes or size. ");
     } else {
-        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);
     }
+
     auto ibuf = std::unique_ptr<uint32_t[]>(new uint32_t[pht->index_size]);
     LOG_DEBUG("index_offset: `", pht->index_offset);