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

@@ -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, HashUint64Pair> {
+    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, HashUint64Pair> {
+    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, HashUint64Pair> {
+    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, HashUint64Pair> {
+    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();

env/unique_id_gen.cc

@@ -7,14 +7,27 @@
 
 #include <algorithm>
 #include <array>
+#include <atomic>
+#include <cstdint>
 #include <cstring>
 #include <random>
 
+#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 <intrin.h>
+#else
+#include <x86intrin.h>
+#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<uint64_t>(_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

env/unique_id_gen.h

@@ -12,10 +12,12 @@
 
 #pragma once
 
+#include <array>
 #include <atomic>
 #include <cstdint>
 #include <type_traits>
 
+#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<uint64_t>& TEST_counter() { return counter_; }
+#endif
+ private:
+  // 256 bit entropy pool
+  std::array<std::atomic<uint64_t>, 4> pool_;
+  // Counter to ensure unique hash inputs
+  std::atomic<uint64_t> counter_;
+};
+
 }  // namespace ROCKSDB_NAMESPACE