diff --git a/table/block_based/block_based_table_builder.cc b/table/block_based/block_based_table_builder.cc index 938da76495..24ec4bf574 100644 --- a/table/block_based/block_based_table_builder.cc +++ b/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 = 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 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 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)); diff --git a/table/block_based/block_based_table_builder.h b/table/block_based/block_based_table_builder.h index f883f5c21c..358324d04b 100644 --- a/table/block_based/block_based_table_builder.h +++ b/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(); diff --git a/util/compression_test.cc b/util/compression_test.cc index 5743f3ffa5..1b0bdb7f5c 100644 --- a/util/compression_test.cc +++ b/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 tamper_count{0}; + SyncPoint::GetInstance()->SetCallBack( + "BlockBasedTableBuilder::CompressAndVerifyBlock:" + "TamperWithCompressedDataBeforeVerify", + [&](void* arg) { + auto* output = static_cast(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;