From 98c6d7fd8026e6d0a7d1d48da1eba7cffb38774f Mon Sep 17 00:00:00 2001 From: Peter Dillinger Date: Wed, 21 Jun 2023 11:32:49 -0700 Subject: [PATCH] Internal API for generating semi-random salt (#11331) Summary: ... so that a non-cryptographic whole file checksum would be highly resistant to manipulation by a user able to manipulate key-value data (e.g. a user whose data is stored in RocksDB) and able to predict SST metadata such as DB session id and file number based on read access to logs or DB files. The adversary would also need to predict the salt in order to influence the checksum result toward collision with another file's checksum. This change is just internal code to support such a future feature. I think this should be a passive feature, not option-controlled, because you probably won't think about needing it until you discover you do need it, and it should be low cost, in space (16 bytes per SST file) and CPU. Pull Request resolved: https://github.com/facebook/rocksdb/pull/11331 Test Plan: Unit tests added to verify at least pseudorandom behavior. (Actually caught a bug in first draft!) The new "stress" style tests run in ~3ms each on my system. Reviewed By: ajkr Differential Revision: D46129415 Pulled By: pdillinger fbshipit-source-id: 7972dc74487e062b29b1fd9c227425e922c98796 --- env/env_test.cc | 84 ++++++++++++++++++++++++++++++++++++++++++++ env/unique_id_gen.cc | 78 ++++++++++++++++++++++++++++++++++++++++ env/unique_id_gen.h | 34 ++++++++++++++++++ 3 files changed, 196 insertions(+) diff --git a/env/env_test.cc b/env/env_test.cc index eedec9bea..1f2077a45 100644 --- a/env/env_test.cc +++ b/env/env_test.cc @@ -3148,6 +3148,90 @@ TEST_F(EnvTest, SemiStructuredUniqueIdGenTestSmaller) { } } +TEST_F(EnvTest, UnpredictableUniqueIdGenTest1) { + // Must be thread safe and usable as a static. + static UnpredictableUniqueIdGen gen; + + struct MyStressTest + : public NoDuplicateMiniStressTest { + uint64_pair_t Generate() override { + uint64_pair_t p; + gen.GenerateNext(&p.first, &p.second); + return p; + } + }; + + MyStressTest t; + t.Run(); +} + +TEST_F(EnvTest, UnpredictableUniqueIdGenTest2) { + // Even if we completely strip the seeding and entropy of the structure + // down to a bare minimum, we still get quality pseudorandom results. + static UnpredictableUniqueIdGen gen{ + UnpredictableUniqueIdGen::TEST_ZeroInitialized{}}; + + struct MyStressTest + : public NoDuplicateMiniStressTest { + uint64_pair_t Generate() override { + uint64_pair_t p; + // No extra entropy is required to get quality pseudorandom results + gen.GenerateNextWithEntropy(&p.first, &p.second, /*no extra entropy*/ 0); + return p; + } + }; + + MyStressTest t; + t.Run(); +} + +TEST_F(EnvTest, UnpredictableUniqueIdGenTest3) { + struct MyStressTest + : public NoDuplicateMiniStressTest { + uint64_pair_t Generate() override { + uint64_pair_t p; + thread_local UnpredictableUniqueIdGen gen{ + UnpredictableUniqueIdGen::TEST_ZeroInitialized{}}; + // Even without the counter (reset it to thread id), we get quality + // single-threaded results (because part of each result is fed back + // into pool). + gen.TEST_counter().store(Env::Default()->GetThreadID()); + gen.GenerateNext(&p.first, &p.second); + return p; + } + }; + + MyStressTest t; + t.Run(); +} + +TEST_F(EnvTest, UnpredictableUniqueIdGenTest4) { + struct MyStressTest + : public NoDuplicateMiniStressTest { + uint64_pair_t Generate() override { + uint64_pair_t p; + // Even if we reset the state to thread ID each time, RDTSC instruction + // suffices for quality single-threaded results. + UnpredictableUniqueIdGen gen{ + UnpredictableUniqueIdGen::TEST_ZeroInitialized{}}; + gen.TEST_counter().store(Env::Default()->GetThreadID()); + gen.GenerateNext(&p.first, &p.second); + return p; + } + }; + + MyStressTest t; +#ifdef __SSE4_2__ // Our rough check for RDTSC + t.Run(); +#else + ROCKSDB_GTEST_BYPASS("Requires IA32 with RDTSC"); + // because nanosecond time might not be high enough fidelity to have + // incremented after a few hundred instructions, especially in cases where + // we really only have microsecond fidelity. Also, wall clock might not be + // monotonic. +#endif +} + TEST_F(EnvTest, FailureToCreateLockFile) { auto env = Env::Default(); auto fs = env->GetFileSystem(); diff --git a/env/unique_id_gen.cc b/env/unique_id_gen.cc index a1986fa15..7d221d374 100644 --- a/env/unique_id_gen.cc +++ b/env/unique_id_gen.cc @@ -7,14 +7,27 @@ #include #include +#include +#include #include #include +#include "port/lang.h" #include "port/port.h" #include "rocksdb/env.h" #include "rocksdb/version.h" #include "util/hash.h" +#ifdef __SSE4_2__ +#ifdef _WIN32 +#include +#else +#include +#endif +#else +#include "rocksdb/system_clock.h" +#endif + namespace ROCKSDB_NAMESPACE { namespace { @@ -161,4 +174,69 @@ void SemiStructuredUniqueIdGen::GenerateNext(uint64_t* upper, uint64_t* lower) { } } +void UnpredictableUniqueIdGen::Reset() { + for (size_t i = 0; i < pool_.size(); i += 2) { + assert(i + 1 < pool_.size()); + uint64_t a, b; + GenerateRawUniqueId(&a, &b); + pool_[i] = a; + pool_[i + 1] = b; + } +} + +void UnpredictableUniqueIdGen::GenerateNext(uint64_t* upper, uint64_t* lower) { + uint64_t extra_entropy; + // Use timing information (if available) to add to entropy. (Not a disaster + // if unavailable on some platforms. High performance is important.) +#ifdef __SSE4_2__ // More than enough to guarantee rdtsc instruction + extra_entropy = static_cast(_rdtsc()); +#else + extra_entropy = SystemClock::Default()->NowNanos(); +#endif + + GenerateNextWithEntropy(upper, lower, extra_entropy); +} + +void UnpredictableUniqueIdGen::GenerateNextWithEntropy(uint64_t* upper, + uint64_t* lower, + uint64_t extra_entropy) { + // To efficiently ensure unique inputs to the hash function in the presence + // of multithreading, we do not require atomicity on the whole entropy pool, + // but instead only a piece of it (a 64-bit counter) that is sufficient to + // guarantee uniqueness. + uint64_t count = counter_.fetch_add(1, std::memory_order_relaxed); + uint64_t a = count; + uint64_t b = extra_entropy; + // Invoking the hash function several times avoids copying all the inputs + // to a contiguous, non-atomic buffer. + BijectiveHash2x64(a, b, &a, &b); // Based on XXH128 + + // In hashing the rest of the pool with that, we don't need to worry about + // races, but use atomic operations for sanitizer-friendliness. + for (size_t i = 0; i < pool_.size(); i += 2) { + assert(i + 1 < pool_.size()); + a ^= pool_[i].load(std::memory_order_relaxed); + b ^= pool_[i + 1].load(std::memory_order_relaxed); + BijectiveHash2x64(a, b, &a, &b); // Based on XXH128 + } + + // Return result + *lower = a; + *upper = b; + + // Add some back into pool. We don't really care that there's a race in + // storing the result back and another thread computing the next value. + // It's just an entropy pool. + pool_[count & (pool_.size() - 1)].fetch_add(a, std::memory_order_relaxed); +} + +#ifndef NDEBUG +UnpredictableUniqueIdGen::UnpredictableUniqueIdGen(TEST_ZeroInitialized) { + for (auto& p : pool_) { + p.store(0); + } + counter_.store(0); +} +#endif + } // namespace ROCKSDB_NAMESPACE diff --git a/env/unique_id_gen.h b/env/unique_id_gen.h index 69033f8d3..f654c7b11 100644 --- a/env/unique_id_gen.h +++ b/env/unique_id_gen.h @@ -12,10 +12,12 @@ #pragma once +#include #include #include #include +#include "port/port.h" #include "rocksdb/rocksdb_namespace.h" namespace ROCKSDB_NAMESPACE { @@ -82,4 +84,36 @@ class SemiStructuredUniqueIdGen { int64_t saved_process_id_; }; +// A unique id generator that should provide reasonable security against +// predicting the output from previous outputs, but is NOT known to be +// cryptographically secure. Unlike std::random_device, this is guaranteed +// not to block once initialized. +class ALIGN_AS(CACHE_LINE_SIZE) UnpredictableUniqueIdGen { + public: + // Initializes with random starting state (from several GenerateRawUniqueId) + UnpredictableUniqueIdGen() { Reset(); } + // Re-initializes, but not thread safe + void Reset(); + + // Generate next probabilistically unique value. Thread safe. Uses timing + // information to add to the entropy pool. + void GenerateNext(uint64_t* upper, uint64_t* lower); + + // Explicitly include given value for entropy pool instead of timing + // information. + void GenerateNextWithEntropy(uint64_t* upper, uint64_t* lower, + uint64_t extra_entropy); + +#ifndef NDEBUG + struct TEST_ZeroInitialized {}; + explicit UnpredictableUniqueIdGen(TEST_ZeroInitialized); + std::atomic& TEST_counter() { return counter_; } +#endif + private: + // 256 bit entropy pool + std::array, 4> pool_; + // Counter to ensure unique hash inputs + std::atomic counter_; +}; + } // namespace ROCKSDB_NAMESPACE