From b1c27a52d2a749b7acd39515e332b3e33342333e Mon Sep 17 00:00:00 2001
From: Levi Tamasi <ltamasi@fb.com>
Date: Mon, 1 Nov 2021 12:31:13 -0700
Subject: [PATCH] Add a consistency check that prevents the overflow of garbage
 in blob files (#9100)

Summary:
The number or total size of garbage blobs in any given blob file can
never exceed the number or total size of all blobs in the file. (This
would be a similar error to e.g. attempting to remove from the LSM tree
an SST file that has already been removed.) The patch builds on
https://github.com/facebook/rocksdb/issues/9085 and adds a
consistency check to `VersionBuilder` that prevents the above from
happening.

Pull Request resolved: https://github.com/facebook/rocksdb/pull/9100

Test Plan: `make check`

Reviewed By: riversand963

Differential Revision: D32048982

Pulled By: ltamasi

fbshipit-source-id: 6f7e0793bf534ad04c3359cc0f696b8e4e5ef81c
---
 db/version_builder.cc      | 19 ++++++++++--
 db/version_builder_test.cc | 63 ++++++++++++++++++++++++++++++++++++++
 2 files changed, 79 insertions(+), 3 deletions(-)

diff --git a/db/version_builder.cc b/db/version_builder.cc
index bd1eae935..69dc34fe0 100644
--- a/db/version_builder.cc
+++ b/db/version_builder.cc
@@ -190,11 +190,20 @@ class VersionBuilder::Rep {
 
     uint64_t GetGarbageBlobBytes() const { return garbage_blob_bytes_; }
 
-    void AddGarbage(uint64_t count, uint64_t bytes) {
+    bool AddGarbage(uint64_t count, uint64_t bytes) {
+      assert(shared_meta_);
+
+      if (garbage_blob_count_ + count > shared_meta_->GetTotalBlobCount() ||
+          garbage_blob_bytes_ + bytes > shared_meta_->GetTotalBlobBytes()) {
+        return false;
+      }
+
       delta_.AddGarbage(count, bytes);
 
       garbage_blob_count_ += count;
       garbage_blob_bytes_ += bytes;
+
+      return true;
     }
 
     void LinkSst(uint64_t sst_file_number) {
@@ -555,8 +564,12 @@ class VersionBuilder::Rep {
       return Status::Corruption("VersionBuilder", oss.str());
     }
 
-    mutable_meta->AddGarbage(blob_file_garbage.GetGarbageBlobCount(),
-                             blob_file_garbage.GetGarbageBlobBytes());
+    if (!mutable_meta->AddGarbage(blob_file_garbage.GetGarbageBlobCount(),
+                                  blob_file_garbage.GetGarbageBlobBytes())) {
+      std::ostringstream oss;
+      oss << "Garbage overflow for blob file #" << blob_file_number;
+      return Status::Corruption("VersionBuilder", oss.str());
+    }
 
     return Status::OK();
   }
diff --git a/db/version_builder_test.cc b/db/version_builder_test.cc
index 502b70c7b..67a9a135a 100644
--- a/db/version_builder_test.cc
+++ b/db/version_builder_test.cc
@@ -923,6 +923,69 @@ TEST_F(VersionBuilderTest, ApplyBlobFileGarbageFileNotFound) {
   ASSERT_TRUE(std::strstr(s.getState(), "Blob file #1234 not found"));
 }
 
+TEST_F(VersionBuilderTest, BlobFileGarbageOverflow) {
+  // Test that VersionEdits that would result in the count/total size of garbage
+  // exceeding the count/total size of all blobs are rejected.
+
+  EnvOptions env_options;
+  constexpr TableCache* table_cache = nullptr;
+  constexpr VersionSet* version_set = nullptr;
+
+  VersionBuilder builder(env_options, &ioptions_, table_cache, &vstorage_,
+                         version_set);
+
+  VersionEdit addition;
+
+  constexpr uint64_t blob_file_number = 1234;
+  constexpr uint64_t total_blob_count = 5678;
+  constexpr uint64_t total_blob_bytes = 999999;
+  constexpr char checksum_method[] = "SHA1";
+  constexpr char checksum_value[] =
+      "\xbd\xb7\xf3\x4a\x59\xdf\xa1\x59\x2c\xe7\xf5\x2e\x99\xf9\x8c\x57\x0c\x52"
+      "\x5c\xbd";
+
+  addition.AddBlobFile(blob_file_number, total_blob_count, total_blob_bytes,
+                       checksum_method, checksum_value);
+
+  // Add dummy table file to ensure the blob file is referenced.
+  constexpr uint64_t table_file_number = 1;
+  AddDummyFileToEdit(&addition, table_file_number, blob_file_number);
+
+  ASSERT_OK(builder.Apply(&addition));
+
+  {
+    // Garbage blob count overflow
+    constexpr uint64_t garbage_blob_count = 5679;
+    constexpr uint64_t garbage_blob_bytes = 999999;
+
+    VersionEdit garbage;
+
+    garbage.AddBlobFileGarbage(blob_file_number, garbage_blob_count,
+                               garbage_blob_bytes);
+
+    const Status s = builder.Apply(&garbage);
+    ASSERT_TRUE(s.IsCorruption());
+    ASSERT_TRUE(
+        std::strstr(s.getState(), "Garbage overflow for blob file #1234"));
+  }
+
+  {
+    // Garbage blob bytes overflow
+    constexpr uint64_t garbage_blob_count = 5678;
+    constexpr uint64_t garbage_blob_bytes = 1000000;
+
+    VersionEdit garbage;
+
+    garbage.AddBlobFileGarbage(blob_file_number, garbage_blob_count,
+                               garbage_blob_bytes);
+
+    const Status s = builder.Apply(&garbage);
+    ASSERT_TRUE(s.IsCorruption());
+    ASSERT_TRUE(
+        std::strstr(s.getState(), "Garbage overflow for blob file #1234"));
+  }
+}
+
 TEST_F(VersionBuilderTest, SaveBlobFilesTo) {
   // Add three blob files to base version.
   for (uint64_t i = 3; i >= 1; --i) {