From 1d34cd797ed59d90a9cc0b0f87a008f374aabb3c Mon Sep 17 00:00:00 2001
From: Peter Dillinger <peterd@fb.com>
Date: Thu, 29 Jul 2021 14:58:35 -0700
Subject: [PATCH] Fix insecure internal API for GetImpl (#8590)

Summary:
Calling the GetImpl function could leave reference to a local
callback function in a field of a parameter struct. As this is
performance-critical code, I'm not going to attempt to sanitize this
code too much, but make the existing hack a bit cleaner by reverting
what it overwrites in the input struct.

Added SaveAndRestore utility class to make that easier.

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

Test Plan:
added unit test for SaveAndRestore; existing tests for
GetImpl

Reviewed By: riversand963

Differential Revision: D29947983

Pulled By: pdillinger

fbshipit-source-id: 2f608853f970bc06724e834cc84dcc4b8599ddeb
---
 db/db_impl/db_impl.cc |  3 +++
 util/defer.h          | 21 ++++++++++++++++++++-
 util/defer_test.cc    | 11 +++++++++++
 3 files changed, 34 insertions(+), 1 deletion(-)

diff --git a/db/db_impl/db_impl.cc b/db/db_impl/db_impl.cc
index 95eedbf49..4f0bcf859 100644
--- a/db/db_impl/db_impl.cc
+++ b/db/db_impl/db_impl.cc
@@ -99,6 +99,7 @@
 #include "util/coding.h"
 #include "util/compression.h"
 #include "util/crc32c.h"
+#include "util/defer.h"
 #include "util/mutexlock.h"
 #include "util/stop_watch.h"
 #include "util/string_util.h"
@@ -1794,6 +1795,8 @@ Status DBImpl::GetImpl(const ReadOptions& read_options, const Slice& key,
   }
   // If timestamp is used, we use read callback to ensure <key,t,s> is returned
   // only if t <= read_opts.timestamp and s <= snapshot.
+  // HACK: temporarily overwrite input struct field but restore
+  SaveAndRestore<ReadCallback*> restore_callback(&get_impl_options.callback);
   if (ts_sz > 0) {
     assert(!get_impl_options
                 .callback);  // timestamp with callback is not supported
diff --git a/util/defer.h b/util/defer.h
index cb0b42a36..81b0df255 100644
--- a/util/defer.h
+++ b/util/defer.h
@@ -38,7 +38,7 @@ namespace ROCKSDB_NAMESPACE {
 // lambda passed to Defer, and you can return immediately on failure when necessary.
 class Defer final {
  public:
-  Defer(std::function<void()>&& fn) : fn_(std::move(fn)) {}
+  explicit Defer(std::function<void()>&& fn) : fn_(std::move(fn)) {}
   ~Defer() { fn_(); }
 
   // Disallow copy.
@@ -49,4 +49,23 @@ class Defer final {
   std::function<void()> fn_;
 };
 
+// An RAII utility object that saves the current value of an object so that
+// it can be overwritten, and restores it to the saved value when the
+// SaveAndRestore object goes out of scope.
+template <typename T>
+class SaveAndRestore {
+ public:
+  // obj is non-null pointer to value to be saved and later restored.
+  explicit SaveAndRestore(T* obj) : obj_(obj), saved_(*obj) {}
+  ~SaveAndRestore() { *obj_ = std::move(saved_); }
+
+  // No copies
+  SaveAndRestore(const SaveAndRestore&) = delete;
+  SaveAndRestore& operator=(const SaveAndRestore&) = delete;
+
+ private:
+  T* const obj_;
+  T saved_;
+};
+
 }  // namespace ROCKSDB_NAMESPACE
diff --git a/util/defer_test.cc b/util/defer_test.cc
index e13b25efb..1334e68b2 100644
--- a/util/defer_test.cc
+++ b/util/defer_test.cc
@@ -30,6 +30,17 @@ TEST(DeferTest, FunctionScope) {
   ASSERT_EQ(4, v);
 }
 
+TEST(SaveAndRestoreTest, BlockScope) {
+  int v = 1;
+  {
+    SaveAndRestore<int> sr(&v);
+    ASSERT_EQ(v, 1);
+    v = 2;
+    ASSERT_EQ(v, 2);
+  }
+  ASSERT_EQ(v, 1);
+}
+
 }  // namespace ROCKSDB_NAMESPACE
 
 int main(int argc, char** argv) {