Fix potential memory leak of scratch buffer (#6879)

Summary:
If `req.scratch` is an internally allocated buffer, but `raw_block_contents` is not constructed to own `req.scratch`, then `req.scratch` will be leaked.
Pull Request resolved: https://github.com/facebook/rocksdb/pull/6879

Test Plan: make asan_check

Reviewed By: anand1976

Differential Revision: D21728498

Pulled By: cheng-chang

fbshipit-source-id: 8fc6a4f2543918c565ddc16ecfad1807eb9a42cf

table/block_based/block_based_table_reader.cc

@@ -1717,7 +1717,8 @@ void BlockBasedTable::RetrieveMultipleBlocks(
     FSReadRequest& req = read_reqs[req_idx];
     Status s = req.status;
     if (s.ok()) {
-      if (req.result.size() != req.len) {
+      if ((req.result.size() != req.len) ||
+          (req_offset + block_size(handle) > req.result.size())) {
         s = Status::Corruption(
             "truncated block read from " + rep_->file->file_name() +
             " offset " + ToString(handle.offset()) + ", expected " +
@@ -1726,49 +1727,39 @@ void BlockBasedTable::RetrieveMultipleBlocks(
     }
 
     BlockContents raw_block_contents;
-    size_t cur_read_end = req_offset + block_size(handle);
-    if (cur_read_end > req.result.size()) {
-      s = Status::Corruption(
-          "truncated block read from " + rep_->file->file_name() + " offset " +
-          ToString(handle.offset()) + ", expected " + ToString(req.len) +
-          " bytes, got " + ToString(req.result.size()));
+    if (!use_shared_buffer) {
+      // We allocated a buffer for this block. Give ownership of it to
+      // BlockContents so it can free the memory
+      assert(req.result.data() == req.scratch);
+      assert(req.result.size() == block_size(handle));
+      assert(req_offset == 0);
+      std::unique_ptr<char[]> raw_block(req.scratch);
+      raw_block_contents = BlockContents(std::move(raw_block), handle.size());
+    } else {
+      // We used the scratch buffer or direct io buffer
+      // which are shared by the blocks.
+      // raw_block_contents does not have the ownership.
+      raw_block_contents =
+          BlockContents(Slice(req.result.data() + req_offset, handle.size()));
     }
-
-    if (s.ok()) {
-      if (!use_shared_buffer) {
-        // We allocated a buffer for this block. Give ownership of it to
-        // BlockContents so it can free the memory
-        assert(req.result.data() == req.scratch);
-        assert(req.result.size() == block_size(handle));
-        assert(req_offset == 0);
-        std::unique_ptr<char[]> raw_block(req.scratch);
-        raw_block_contents = BlockContents(std::move(raw_block), handle.size());
-      } else {
-        // We used the scratch buffer or direct io buffer
-        // which are shared by the blocks.
-        // raw_block_contents does not have the ownership.
-        raw_block_contents =
-            BlockContents(Slice(req.result.data() + req_offset, handle.size()));
-      }
-
 #ifndef NDEBUG
       raw_block_contents.is_raw_block = true;
 #endif
-      if (options.verify_checksums) {
-        PERF_TIMER_GUARD(block_checksum_time);
-        const char* data = req.result.data();
-        uint32_t expected =
-            DecodeFixed32(data + req_offset + handle.size() + 1);
-        // Since the scratch might be shared. the offset of the data block in
-        // the buffer might not be 0. req.result.data() only point to the
-        // begin address of each read request, we need to add the offset
-        // in each read request. Checksum is stored in the block trailer,
-        // which is handle.size() + 1.
-        s = ROCKSDB_NAMESPACE::VerifyChecksum(footer.checksum(),
-                                              data + req_offset,
-                                              handle.size() + 1, expected);
-        TEST_SYNC_POINT_CALLBACK("RetrieveMultipleBlocks:VerifyChecksum", &s);
-      }
+
+    if (s.ok() && options.verify_checksums) {
+      PERF_TIMER_GUARD(block_checksum_time);
+      const char* data = req.result.data();
+      uint32_t expected =
+          DecodeFixed32(data + req_offset + handle.size() + 1);
+      // Since the scratch might be shared. the offset of the data block in
+      // the buffer might not be 0. req.result.data() only point to the
+      // begin address of each read request, we need to add the offset
+      // in each read request. Checksum is stored in the block trailer,
+      // which is handle.size() + 1.
+      s = ROCKSDB_NAMESPACE::VerifyChecksum(footer.checksum(),
+                                            data + req_offset,
+                                            handle.size() + 1, expected);
+      TEST_SYNC_POINT_CALLBACK("RetrieveMultipleBlocks:VerifyChecksum", &s);
     }
 
     if (s.ok()) {