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
main
Peter Dillinger 3 years ago committed by Facebook GitHub Bot
parent 9501279d5f
commit 1d34cd797e
  1. 3
      db/db_impl/db_impl.cc
  2. 21
      util/defer.h
  3. 11
      util/defer_test.cc

@ -99,6 +99,7 @@
#include "util/coding.h" #include "util/coding.h"
#include "util/compression.h" #include "util/compression.h"
#include "util/crc32c.h" #include "util/crc32c.h"
#include "util/defer.h"
#include "util/mutexlock.h" #include "util/mutexlock.h"
#include "util/stop_watch.h" #include "util/stop_watch.h"
#include "util/string_util.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 // If timestamp is used, we use read callback to ensure <key,t,s> is returned
// only if t <= read_opts.timestamp and s <= snapshot. // 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) { if (ts_sz > 0) {
assert(!get_impl_options assert(!get_impl_options
.callback); // timestamp with callback is not supported .callback); // timestamp with callback is not supported

@ -38,7 +38,7 @@ namespace ROCKSDB_NAMESPACE {
// lambda passed to Defer, and you can return immediately on failure when necessary. // lambda passed to Defer, and you can return immediately on failure when necessary.
class Defer final { class Defer final {
public: public:
Defer(std::function<void()>&& fn) : fn_(std::move(fn)) {} explicit Defer(std::function<void()>&& fn) : fn_(std::move(fn)) {}
~Defer() { fn_(); } ~Defer() { fn_(); }
// Disallow copy. // Disallow copy.
@ -49,4 +49,23 @@ class Defer final {
std::function<void()> fn_; 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 } // namespace ROCKSDB_NAMESPACE

@ -30,6 +30,17 @@ TEST(DeferTest, FunctionScope) {
ASSERT_EQ(4, v); 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 } // namespace ROCKSDB_NAMESPACE
int main(int argc, char** argv) { int main(int argc, char** argv) {

Loading…
Cancel
Save