Fix heap use after free error on retry after checksum mismatch (#12464)

Summary:
Fix the heap use after free bug caused by freeing the file system IO buffer in `BlockFetcher::ReadBlock()` instead of the caller.

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

Test Plan: Update the `DBIOCorruptionTest` tests

Reviewed By: akankshamahajan15

Differential Revision: D55206920

Pulled By: anand1976

fbshipit-source-id: fd6b608a61cd229b20c1e5f348ff3cc92328de0f
This commit is contained in:
anand76 2024-03-21 16:19:09 -07:00 committed by Facebook GitHub Bot
parent 088dc7283b
commit 3b736c4aa3
4 changed files with 48 additions and 21 deletions

View File

@ -20,13 +20,15 @@ class CorruptionFS : public FileSystemWrapper {
bool writable_file_error_; bool writable_file_error_;
int num_writable_file_errors_; int num_writable_file_errors_;
explicit CorruptionFS(const std::shared_ptr<FileSystem>& _target) explicit CorruptionFS(const std::shared_ptr<FileSystem>& _target,
bool fs_buffer)
: FileSystemWrapper(_target), : FileSystemWrapper(_target),
writable_file_error_(false), writable_file_error_(false),
num_writable_file_errors_(0), num_writable_file_errors_(0),
corruption_trigger_(INT_MAX), corruption_trigger_(INT_MAX),
read_count_(0), read_count_(0),
rnd_(300) {} rnd_(300),
fs_buffer_(fs_buffer) {}
~CorruptionFS() override { ~CorruptionFS() override {
// Assert that the corruption was reset, which means it got triggered // Assert that the corruption was reset, which means it got triggered
assert(corruption_trigger_ == INT_MAX); assert(corruption_trigger_ == INT_MAX);
@ -81,7 +83,21 @@ class CorruptionFS : public FileSystemWrapper {
IOStatus MultiRead(FSReadRequest* reqs, size_t num_reqs, IOStatus MultiRead(FSReadRequest* reqs, size_t num_reqs,
const IOOptions& options, const IOOptions& options,
IODebugContext* dbg) override { IODebugContext* dbg) override {
return FSRandomAccessFile::MultiRead(reqs, num_reqs, options, dbg); for (size_t i = 0; i < num_reqs; ++i) {
FSReadRequest& req = reqs[i];
if (fs_.fs_buffer_) {
FSAllocationPtr buffer(new char[req.len], [](void* ptr) {
delete[] static_cast<char*>(ptr);
});
req.fs_scratch = std::move(buffer);
req.status = Read(req.offset, req.len, options, &req.result,
static_cast<char*>(req.fs_scratch.get()), dbg);
} else {
req.status = Read(req.offset, req.len, options, &req.result,
req.scratch, dbg);
}
}
return IOStatus::OK();
} }
private: private:
@ -99,12 +115,16 @@ class CorruptionFS : public FileSystemWrapper {
void SupportedOps(int64_t& supported_ops) override { void SupportedOps(int64_t& supported_ops) override {
supported_ops = 1 << FSSupportedOps::kVerifyAndReconstructRead | supported_ops = 1 << FSSupportedOps::kVerifyAndReconstructRead |
1 << FSSupportedOps::kAsyncIO; 1 << FSSupportedOps::kAsyncIO;
if (fs_buffer_) {
supported_ops |= 1 << FSSupportedOps::kFSBuffer;
}
} }
private: private:
int corruption_trigger_; int corruption_trigger_;
int read_count_; int read_count_;
Random rnd_; Random rnd_;
bool fs_buffer_;
}; };
} // anonymous namespace } // anonymous namespace
@ -674,8 +694,9 @@ TEST_F(DBIOFailureTest, CompactionSstSyncError) {
} }
#endif // !(defined NDEBUG) || !defined(OS_WIN) #endif // !(defined NDEBUG) || !defined(OS_WIN)
class DBIOCorruptionTest : public DBIOFailureTest, class DBIOCorruptionTest
public testing::WithParamInterface<bool> { : public DBIOFailureTest,
public testing::WithParamInterface<std::tuple<bool, bool>> {
public: public:
DBIOCorruptionTest() : DBIOFailureTest() { DBIOCorruptionTest() : DBIOFailureTest() {
BlockBasedTableOptions bbto; BlockBasedTableOptions bbto;
@ -683,7 +704,8 @@ class DBIOCorruptionTest : public DBIOFailureTest,
base_env_ = env_; base_env_ = env_;
EXPECT_NE(base_env_, nullptr); EXPECT_NE(base_env_, nullptr);
fs_.reset(new CorruptionFS(base_env_->GetFileSystem())); fs_.reset(
new CorruptionFS(base_env_->GetFileSystem(), std::get<0>(GetParam())));
env_guard_ = NewCompositeEnv(fs_); env_guard_ = NewCompositeEnv(fs_);
options.env = env_guard_.get(); options.env = env_guard_.get();
bbto.num_file_reads_for_auto_readahead = 0; bbto.num_file_reads_for_auto_readahead = 0;
@ -714,7 +736,7 @@ TEST_P(DBIOCorruptionTest, GetReadCorruptionRetry) {
std::string val; std::string val;
ReadOptions ro; ReadOptions ro;
ro.async_io = GetParam(); ro.async_io = std::get<1>(GetParam());
ASSERT_OK(dbfull()->Get(ReadOptions(), "key1", &val)); ASSERT_OK(dbfull()->Get(ReadOptions(), "key1", &val));
ASSERT_EQ(val, "val1"); ASSERT_EQ(val, "val1");
} }
@ -729,7 +751,7 @@ TEST_P(DBIOCorruptionTest, IterReadCorruptionRetry) {
ReadOptions ro; ReadOptions ro;
ro.readahead_size = 65536; ro.readahead_size = 65536;
ro.async_io = GetParam(); ro.async_io = std::get<1>(GetParam());
Iterator* iter = dbfull()->NewIterator(ro); Iterator* iter = dbfull()->NewIterator(ro);
iter->SeekToFirst(); iter->SeekToFirst();
@ -754,7 +776,7 @@ TEST_P(DBIOCorruptionTest, MultiGetReadCorruptionRetry) {
std::vector<PinnableSlice> values(keys.size()); std::vector<PinnableSlice> values(keys.size());
std::vector<Status> statuses(keys.size()); std::vector<Status> statuses(keys.size());
ReadOptions ro; ReadOptions ro;
ro.async_io = GetParam(); ro.async_io = std::get<1>(GetParam());
dbfull()->MultiGet(ro, dbfull()->DefaultColumnFamily(), keys.size(), dbfull()->MultiGet(ro, dbfull()->DefaultColumnFamily(), keys.size(),
keys.data(), values.data(), statuses.data()); keys.data(), values.data(), statuses.data());
ASSERT_EQ(values[0].ToString(), "val1"); ASSERT_EQ(values[0].ToString(), "val1");
@ -775,7 +797,7 @@ TEST_P(DBIOCorruptionTest, CompactionReadCorruptionRetry) {
std::string val; std::string val;
ReadOptions ro; ReadOptions ro;
ro.async_io = GetParam(); ro.async_io = std::get<1>(GetParam());
ASSERT_OK(dbfull()->Get(ro, "key1", &val)); ASSERT_OK(dbfull()->Get(ro, "key1", &val));
ASSERT_EQ(val, "val1"); ASSERT_EQ(val, "val1");
} }
@ -790,13 +812,13 @@ TEST_P(DBIOCorruptionTest, FlushReadCorruptionRetry) {
std::string val; std::string val;
ReadOptions ro; ReadOptions ro;
ro.async_io = GetParam(); ro.async_io = std::get<1>(GetParam());
ASSERT_OK(dbfull()->Get(ro, "key1", &val)); ASSERT_OK(dbfull()->Get(ro, "key1", &val));
ASSERT_EQ(val, "val1"); ASSERT_EQ(val, "val1");
} }
INSTANTIATE_TEST_CASE_P(DBIOCorruptionTest, DBIOCorruptionTest, INSTANTIATE_TEST_CASE_P(DBIOCorruptionTest, DBIOCorruptionTest,
testing::Bool()); testing::Combine(testing::Bool(), testing::Bool()));
} // namespace ROCKSDB_NAMESPACE } // namespace ROCKSDB_NAMESPACE
int main(int argc, char** argv) { int main(int argc, char** argv) {

View File

@ -789,6 +789,8 @@ class FSSequentialFile {
// SequentialFileWrapper too. // SequentialFileWrapper too.
}; };
using FSAllocationPtr = std::unique_ptr<void, std::function<void(void*)>>;
// A read IO request structure for use in MultiRead and asynchronous Read APIs. // A read IO request structure for use in MultiRead and asynchronous Read APIs.
struct FSReadRequest { struct FSReadRequest {
// Input parameter that represents the file offset in bytes. // Input parameter that represents the file offset in bytes.
@ -855,7 +857,7 @@ struct FSReadRequest {
// - FSReadRequest::result should point to fs_scratch. // - FSReadRequest::result should point to fs_scratch.
// - This is needed only if FSSupportedOps::kFSBuffer support is provided by // - This is needed only if FSSupportedOps::kFSBuffer support is provided by
// underlying FS. // underlying FS.
std::unique_ptr<void, std::function<void(void*)>> fs_scratch; FSAllocationPtr fs_scratch;
}; };
// A file abstraction for randomly reading the contents of a file. // A file abstraction for randomly reading the contents of a file.

View File

@ -241,7 +241,7 @@ inline void BlockFetcher::GetBlockContents() {
// Read a block from the file and verify its checksum. Upon return, io_status_ // Read a block from the file and verify its checksum. Upon return, io_status_
// will be updated with the status of the read, and slice_ will be updated // will be updated with the status of the read, and slice_ will be updated
// with a pointer to the data. // with a pointer to the data.
void BlockFetcher::ReadBlock(bool retry) { void BlockFetcher::ReadBlock(bool retry, FSAllocationPtr& fs_buf) {
FSReadRequest read_req; FSReadRequest read_req;
IOOptions opts; IOOptions opts;
io_status_ = file_->PrepareIOOptions(read_options_, opts); io_status_ = file_->PrepareIOOptions(read_options_, opts);
@ -336,6 +336,7 @@ void BlockFetcher::ReadBlock(bool retry) {
if (io_status_.ok()) { if (io_status_.ok()) {
InsertCompressedBlockToPersistentCacheIfNeeded(); InsertCompressedBlockToPersistentCacheIfNeeded();
fs_buf = std::move(read_req.fs_scratch);
} else { } else {
ReleaseFileSystemProvidedBuffer(&read_req); ReleaseFileSystemProvidedBuffer(&read_req);
direct_io_buf_.reset(); direct_io_buf_.reset();
@ -346,8 +347,7 @@ void BlockFetcher::ReadBlock(bool retry) {
} }
IOStatus BlockFetcher::ReadBlockContents() { IOStatus BlockFetcher::ReadBlockContents() {
FSReadRequest read_req; FSAllocationPtr fs_buf;
read_req.status.PermitUncheckedError();
if (TryGetUncompressBlockFromPersistentCache()) { if (TryGetUncompressBlockFromPersistentCache()) {
compression_type_ = kNoCompression; compression_type_ = kNoCompression;
#ifndef NDEBUG #ifndef NDEBUG
@ -360,13 +360,15 @@ IOStatus BlockFetcher::ReadBlockContents() {
return io_status_; return io_status_;
} }
} else if (!TryGetSerializedBlockFromPersistentCache()) { } else if (!TryGetSerializedBlockFromPersistentCache()) {
ReadBlock(/*retry =*/false); ReadBlock(/*retry =*/false, fs_buf);
// If the file system supports retry after corruption, then try to // If the file system supports retry after corruption, then try to
// re-read the block and see if it succeeds. // re-read the block and see if it succeeds.
if (io_status_.IsCorruption() && retry_corrupt_read_) { if (io_status_.IsCorruption() && retry_corrupt_read_) {
ReadBlock(/*retry=*/true); assert(!fs_buf);
ReadBlock(/*retry=*/true, fs_buf);
} }
if (!io_status_.ok()) { if (!io_status_.ok()) {
assert(!fs_buf);
return io_status_; return io_status_;
} }
} }
@ -390,7 +392,6 @@ IOStatus BlockFetcher::ReadBlockContents() {
} }
InsertUncompressedBlockToPersistentCacheIfNeeded(); InsertUncompressedBlockToPersistentCacheIfNeeded();
ReleaseFileSystemProvidedBuffer(&read_req);
return io_status_; return io_status_;
} }
@ -416,14 +417,16 @@ IOStatus BlockFetcher::ReadAsyncBlockContents() {
return io_s; return io_s;
} }
if (io_s.ok()) { if (io_s.ok()) {
FSAllocationPtr fs_buf;
// Data Block is already in prefetch. // Data Block is already in prefetch.
got_from_prefetch_buffer_ = true; got_from_prefetch_buffer_ = true;
ProcessTrailerIfPresent(); ProcessTrailerIfPresent();
if (io_status_.IsCorruption() && retry_corrupt_read_) { if (io_status_.IsCorruption() && retry_corrupt_read_) {
got_from_prefetch_buffer_ = false; got_from_prefetch_buffer_ = false;
ReadBlock(/*retry = */ true); ReadBlock(/*retry = */ true, fs_buf);
} }
if (!io_status_.ok()) { if (!io_status_.ok()) {
assert(!fs_buf);
return io_status_; return io_status_;
} }
used_buf_ = const_cast<char*>(slice_.data()); used_buf_ = const_cast<char*>(slice_.data());

View File

@ -152,7 +152,7 @@ class BlockFetcher {
void InsertCompressedBlockToPersistentCacheIfNeeded(); void InsertCompressedBlockToPersistentCacheIfNeeded();
void InsertUncompressedBlockToPersistentCacheIfNeeded(); void InsertUncompressedBlockToPersistentCacheIfNeeded();
void ProcessTrailerIfPresent(); void ProcessTrailerIfPresent();
void ReadBlock(bool retry); void ReadBlock(bool retry, FSAllocationPtr& fs_buf);
void ReleaseFileSystemProvidedBuffer(FSReadRequest* read_req) { void ReleaseFileSystemProvidedBuffer(FSReadRequest* read_req) {
if (use_fs_scratch_) { if (use_fs_scratch_) {