Verify compressed block bytes after checksum computation (#15098)

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

Compute the base block checksum in `CompressAndVerifyBlock()` before `verify_compression` decompresses the compressed output. Parallel compression stores that base checksum on the block rep so worker threads can keep the CPU-heavy checksum and verification work off the ordered write path; the writer only applies the offset context modifier before writing the trailer.

Update the regression test to cover both serial and parallel compression paths.

Reviewed By: pdillinger

Differential Revision: D115950785

fbshipit-source-id: b5e47f7d7c41680e7e7907443fc805bc0d043d6e

table/block_based/block_based_table_builder.cc

@@ -365,6 +365,7 @@ struct BlockBasedTableBuilder::ParallelCompressionRep {
     std::string uncompressed;
     GrowableBuffer compressed;
     CompressionType compression_type = kNoCompression;
+    uint32_t contents_checksum = 0;
     // Set by the emit thread: this data block's compression was auto-skipped,
     // so no worker should compress it -- StateTransition auto-completes it
     // straight to the write stage. Set on every emit (true=skip,
@@ -2134,6 +2135,7 @@ void BlockBasedTableBuilder::EmitBlockForParallel(
                                       block_rep->prepared_index_entry.get());
   block_rep->compressed.Reset();
   block_rep->compression_type = kNoCompression;
+  block_rep->contents_checksum = 0;
   // Emit thread makes the auto-skip decision here so skipped blocks are never
   // dispatched to a worker for (no-op) compression; StateTransition
   // auto-completes them straight to the write stage.
@@ -2158,7 +2160,7 @@ void BlockBasedTableBuilder::EmitBlockForParallel(
       Status s = CompressAndVerifyBlock(
           block_rep->uncompressed, /*is_data_block=*/true,
           r->data_block_working_area, &block_rep->compressed,
-          &block_rep->compression_type);
+          &block_rep->compression_type, &block_rep->contents_checksum);
       if (UNLIKELY(!s.ok())) {
         r->SetStatus(s);
         pc_rep.SetAbort(pc_rep.emit_thread_state);
@@ -2218,6 +2220,7 @@ void BlockBasedTableBuilder::WriteBlock(const Slice& uncompressed_block_data,
   // Single-threaded context only
   assert(!r->IsParallelCompressionActive());
   CompressionType type = kNoCompression;
+  uint32_t contents_checksum = 0;
   bool is_data_block = block_type == BlockType::kData;
   // NOTE: only index and data blocks are currently compressed
   assert(is_data_block || block_type == BlockType::kIndex);
@@ -2229,11 +2232,11 @@ void BlockBasedTableBuilder::WriteBlock(const Slice& uncompressed_block_data,
       r->AutoSkipDecideBypassEligible()) {
     r->AutoSkipRecordBypass(uncompressed_block_data.size());
   } else {
-    Status compress_status =
-        CompressAndVerifyBlock(uncompressed_block_data, is_data_block,
-                               is_data_block ? r->data_block_working_area
-                                             : r->index_block_working_area,
-                               &r->single_threaded_compressed_output, &type);
+    Status compress_status = CompressAndVerifyBlock(
+        uncompressed_block_data, is_data_block,
+        is_data_block ? r->data_block_working_area
+                      : r->index_block_working_area,
+        &r->single_threaded_compressed_output, &type, &contents_checksum);
     r->SetStatus(compress_status);
     if (UNLIKELY(!ok())) {
       return;
@@ -2245,7 +2248,8 @@ void BlockBasedTableBuilder::WriteBlock(const Slice& uncompressed_block_data,
   WriteMaybeCompressedBlock(
       type == kNoCompression ? uncompressed_block_data
                              : Slice(r->single_threaded_compressed_output),
-      type, handle, block_type, &uncompressed_block_data, skip_delta_encoding);
+      type, handle, block_type, &uncompressed_block_data, skip_delta_encoding,
+      type == kNoCompression ? nullptr : &contents_checksum);
   r->single_threaded_compressed_output.Reset();
   if (is_data_block) {
     r->props.data_size = r->get_offset();
@@ -2284,7 +2288,8 @@ void BlockBasedTableBuilder::BGWorker(WorkingAreaPair& working_area) {
     auto compress_fn = [this, block_rep, &ios, &working_area]() {
       ios = status_to_io_status(CompressAndVerifyBlock(
           block_rep->uncompressed, /*is_data_block=*/true, working_area,
-          &block_rep->compressed, &block_rep->compression_type));
+          &block_rep->compressed, &block_rep->compression_type,
+          &block_rep->contents_checksum));
     };
     auto write_fn = [this, block_rep, &ios]() {
       Slice compressed = block_rep->compressed;
@@ -2294,7 +2299,10 @@ void BlockBasedTableBuilder::BGWorker(WorkingAreaPair& working_area) {
           block_rep->compression_type == kNoCompression ? uncompressed
                                                         : compressed,
           block_rep->compression_type, &rep_->pending_handle, BlockType::kData,
-          &uncompressed, &skip_delta_encoding);
+          &uncompressed, &skip_delta_encoding,
+          block_rep->compression_type == kNoCompression
+              ? nullptr
+              : &block_rep->contents_checksum);
       if (LIKELY(ios.ok())) {
         rep_->props.data_size = rep_->get_offset();
         rep_->props.uncompressed_data_size += block_rep->uncompressed.size();
@@ -2342,9 +2350,10 @@ void BlockBasedTableBuilder::BGWorker(WorkingAreaPair& working_area) {
 Status BlockBasedTableBuilder::CompressAndVerifyBlock(
     const Slice& uncompressed_block_data, bool is_data_block,
     WorkingAreaPair& working_area, GrowableBuffer* compressed_output,
-    CompressionType* result_compression_type) {
+    CompressionType* result_compression_type, uint32_t* result_checksum) {
   Rep* r = rep_.get();
   Status status;
+  *result_checksum = 0;
 
   UnownedPtr<Compressor> compressor = nullptr;
   Decompressor* verify_decomp = nullptr;
@@ -2390,6 +2399,16 @@ Status BlockBasedTableBuilder::CompressAndVerifyBlock(
           "BlockBasedTableBuilder::CompressAndVerifyBlock:TamperWithResultType",
           &type);
 
+      if (type != kNoCompression) {
+        *result_checksum = ComputeBuiltinChecksumWithLastByte(
+            r->table_options.checksum, compressed_output->data(),
+            compressed_output->size(), /*last_byte*/ type);
+        TEST_SYNC_POINT_CALLBACK(
+            "BlockBasedTableBuilder::CompressAndVerifyBlock:"
+            "TamperWithCompressedDataBeforeVerify",
+            compressed_output);
+      }
+
       // Some of the compression algorithms are known to be unreliable. If
       // the verify_compression flag is set then try to de-compress the
       // compressed data and compare to the input.
@@ -2495,19 +2514,19 @@ Status BlockBasedTableBuilder::CompressAndVerifyBlock(
 void BlockBasedTableBuilder::WriteMaybeCompressedBlock(
     const Slice& block_contents, CompressionType comp_type, BlockHandle* handle,
     BlockType block_type, const Slice* uncompressed_block_data,
-    bool* skip_delta_encoding) {
+    bool* skip_delta_encoding, const uint32_t* precomputed_checksum) {
   // Must have pre-checked status in single-threaded context
   assert(status().ok());
   assert(io_status().ok());
   rep_->SetIOStatus(WriteMaybeCompressedBlockImpl(
       block_contents, comp_type, handle, block_type, uncompressed_block_data,
-      skip_delta_encoding));
+      skip_delta_encoding, precomputed_checksum));
 }
 
 IOStatus BlockBasedTableBuilder::WriteMaybeCompressedBlockImpl(
     const Slice& block_contents, CompressionType comp_type, BlockHandle* handle,
     BlockType block_type, const Slice* uncompressed_block_data,
-    bool* skip_delta_encoding) {
+    bool* skip_delta_encoding, const uint32_t* precomputed_checksum) {
   // File format contains a sequence of blocks where each block has:
   //    block_data: uint8[n]
   //    compression_type: uint8
@@ -2577,6 +2596,18 @@ IOStatus BlockBasedTableBuilder::WriteMaybeCompressedBlockImpl(
     uncompressed_block_data = &block_contents;
     assert(comp_type == kNoCompression);
   }
+  assert(precomputed_checksum == nullptr || comp_type != kNoCompression);
+
+  r->compression_types_used.Add(comp_type);
+  std::array<char, kBlockTrailerSize> trailer;
+  trailer[0] = comp_type;
+  uint32_t checksum =
+      precomputed_checksum != nullptr
+          ? *precomputed_checksum
+          : ComputeBuiltinChecksumWithLastByte(
+                r->table_options.checksum, block_contents.data(),
+                block_contents.size(), /*last_byte*/ comp_type);
+  checksum += ChecksumModifierForContext(r->base_context_checksum, offset);
 
   // TODO: consider a variant of this function that puts the trailer after
   // block_contents (if it comes from a std::string) so we only need one
@@ -2588,14 +2619,6 @@ IOStatus BlockBasedTableBuilder::WriteMaybeCompressedBlockImpl(
     }
   }
 
-  r->compression_types_used.Add(comp_type);
-  std::array<char, kBlockTrailerSize> trailer;
-  trailer[0] = comp_type;
-  uint32_t checksum = ComputeBuiltinChecksumWithLastByte(
-      r->table_options.checksum, block_contents.data(), block_contents.size(),
-      /*last_byte*/ comp_type);
-  checksum += ChecksumModifierForContext(r->base_context_checksum, offset);
-
   if (block_type == BlockType::kFilter) {
     io_s = status_to_io_status(
         r->filter_builder->MaybePostVerifyFilter(block_contents));

table/block_based/block_based_table_builder.h

@@ -154,14 +154,16 @@ class BlockBasedTableBuilder : public TableBuilder {
   void WriteBlock(const Slice& block_contents, BlockHandle* handle,
                   BlockType block_type, bool* skip_delta_encoding = nullptr);
   // Directly write data to the file.
-  void WriteMaybeCompressedBlock(const Slice& block_contents, CompressionType,
-                                 BlockHandle* handle, BlockType block_type,
-                                 const Slice* uncompressed_block_data = nullptr,
-                                 bool* skip_delta_encoding = nullptr);
+  void WriteMaybeCompressedBlock(
+      const Slice& block_contents, CompressionType, BlockHandle* handle,
+      BlockType block_type, const Slice* uncompressed_block_data = nullptr,
+      bool* skip_delta_encoding = nullptr,
+      const uint32_t* precomputed_checksum = nullptr);
   IOStatus WriteMaybeCompressedBlockImpl(
       const Slice& block_contents, CompressionType, BlockHandle* handle,
       BlockType block_type, const Slice* uncompressed_block_data = nullptr,
-      bool* skip_delta_encoding = nullptr);
+      bool* skip_delta_encoding = nullptr,
+      const uint32_t* precomputed_checksum = nullptr);
 
   void SetupCacheKeyPrefix(const TableBuilderOptions& tbo);
 
@@ -240,7 +242,8 @@ class BlockBasedTableBuilder : public TableBuilder {
                                 bool is_data_block,
                                 WorkingAreaPair& working_area,
                                 GrowableBuffer* compressed_output,
-                                CompressionType* result_compression_type);
+                                CompressionType* result_compression_type,
+                                uint32_t* result_checksum);
 
   // If configured, start worker threads for parallel compression
   void MaybeStartParallelCompression();

util/compression_test.cc

@@ -16,6 +16,7 @@
 #include "table/block_based/block_builder.h"
 #include "table/block_based/data_block_footer.h"
 #include "test_util/testutil.h"
+#include "util/aligned_buffer.h"
 #include "util/auto_tune_compressor.h"
 #include "util/coding.h"
 #include "util/random.h"
@@ -931,6 +932,74 @@ TEST_P(CompressionFailuresTest, CompressionFailures) {
   }
 }
 
+TEST_F(DBCompressionTest, VerifyCompressionChecksFinalCompressedBlockContents) {
+  CompressionType compression_type = kNoCompression;
+  for (CompressionType supported : GetSupportedCompressions()) {
+    if (supported != kNoCompression) {
+      compression_type = supported;
+      break;
+    }
+  }
+  if (compression_type == kNoCompression) {
+    return;
+  }
+
+  struct SyncPointCleanup {
+    ~SyncPointCleanup() {
+      SyncPoint::GetInstance()->DisableProcessing();
+      SyncPoint::GetInstance()->ClearAllCallBacks();
+    }
+  } sync_point_cleanup;
+
+  for (uint32_t parallel_threads : {1, 4}) {
+    SCOPED_TRACE("parallel_threads=" + std::to_string(parallel_threads));
+
+    Options options = CurrentOptions();
+    options.compression = compression_type;
+    options.compression_opts.max_compressed_bytes_per_kb = 1024;
+    options.compression_opts.parallel_threads = parallel_threads;
+
+    BlockBasedTableOptions table_options;
+    table_options.block_size = 512;
+    table_options.verify_compression = true;
+    options.table_factory.reset(NewBlockBasedTableFactory(table_options));
+
+    DestroyAndReopen(options);
+
+    SyncPoint::GetInstance()->DisableProcessing();
+    SyncPoint::GetInstance()->ClearAllCallBacks();
+
+    std::atomic<int> tamper_count{0};
+    SyncPoint::GetInstance()->SetCallBack(
+        "BlockBasedTableBuilder::CompressAndVerifyBlock:"
+        "TamperWithCompressedDataBeforeVerify",
+        [&](void* arg) {
+          auto* output = static_cast<GrowableBuffer*>(arg);
+          ASSERT_FALSE(output->empty());
+          output->data()[output->size() - 1]++;
+          tamper_count++;
+        });
+    SyncPoint::GetInstance()->EnableProcessing();
+
+    Random rnd(405);
+    constexpr int kValUnitSize = 16;
+    constexpr int kValSize = 256;
+    for (int i = 0; i < 5; i++) {
+      std::string value_unit = rnd.RandomString(kValUnitSize);
+      std::string value;
+      for (int j = 0; j < kValSize; j += kValUnitSize) {
+        value += value_unit;
+      }
+      ASSERT_OK(Put(Key(i), value));
+    }
+
+    Status s = Flush();
+
+    ASSERT_GT(tamper_count.load(), 0);
+    ASSERT_TRUE(s.IsCorruption()) << s.ToString();
+  }
+}
+
 TEST_F(DBCompressionTest, CompressionOptions) {
   if (!Zlib_Supported() || !Snappy_Supported()) {
     return;