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