repo-sync-2024-09-23T14:50:14+0800

yacl/crypto/rand/rand.h

@@ -17,15 +17,18 @@
 #include <algorithm>
 #include <climits>
 #include <cstdint>
+#include <limits>
 #include <memory>
 #include <random>
 #include <type_traits>
 #include <vector>
 
 #include "yacl/base/dynamic_bitset.h"
 #include "yacl/base/int128.h"
+#include "yacl/crypto/block_cipher/symmetric_crypto.h"
 #include "yacl/crypto/openssl_wrappers.h"
 #include "yacl/crypto/rand/drbg/drbg.h"
+#include "yacl/crypto/tools/prg.h"
 #include "yacl/math/mpint/mp_int.h"
 #include "yacl/secparam.h"
 
@@ -114,6 +117,18 @@ inline std::vector<T> RandVec(uint64_t len, bool fast_mode = false) {
   return out;
 }
 
+// Generate random T-type vectors in fast mode
+template <typename T, std::enable_if_t<std::is_standard_layout_v<T>, int> = 0>
+inline std::vector<T> FastRandVec(uint64_t len) {
+  return RandVec<T>(len, true);
+}
+
+// Generate random T-type vectors in secure mode
+template <typename T, std::enable_if_t<std::is_standard_layout_v<T>, int> = 0>
+inline std::vector<T> SecureRandVec(uint64_t len) {
+  return RandVec<T>(len, false);
+}
+
 // -----------------------------------
 // Random Support for Integral Numbers
 // -----------------------------------
@@ -139,4 +154,75 @@ inline T RandLtN(T n) {
   return r.Get<T>();
 }
 
+// Implementation of standard (cpp) UniformRandomBitGenerator
+//
+// see: https://en.cppreference.com/w/cpp/named_req/UniformRandomBitGenerator
+//
+// This implementation is provided to be used as the random bit source of
+// std::shuffle. std::shuffle should internally call YaclStdUrbd with
+// std::uniform_int_distribution, which again internally uses the method defined
+// in A.5.1 of SP800-90A (reviewed on MacOS with libc++).
+//
+template <typename T, std::enable_if_t<std::is_unsigned_v<T>, bool> = true>
+class YaclStdUrbg {
+ public:
+  using result_type = T;
+  YaclStdUrbg() { drbg_ = DrbgFactory::Instance().Create("ctr-drbg"); };
+
+  static constexpr T min() { return std::numeric_limits<T>::min(); }
+  static constexpr T max() { return std::numeric_limits<T>::max(); }
+
+  T operator()() {
+    T ret;
+    drbg_->Fill((char *)&ret, sizeof(T));
+    return ret;
+  }
+
+ private:
+  std::unique_ptr<Drbg> drbg_;
+};
+
+// Implementation of standard (cpp) and replayable UniformRandomBitGenerator
+//
+// see: https://en.cppreference.com/w/cpp/named_req/UniformRandomBitGenerator
+//
+// This implementation is provided to be used as the random bit source of
+// std::shuffle. std::shuffle should internally call YaclStdUrbd with
+// std::uniform_int_distribution, which again internally uses the method defined
+// in A.5.1 of SP800-90A (reviewed on MacOS with libc++).
+//
+// NOTE this implementation is not compatiable with NIST standards as it allows
+// the manipulation of internal random states.
+//
+template <typename T, std::enable_if_t<std::is_unsigned_v<T>, bool> = true>
+class YaclReplayUrbg {
+ public:
+  using result_type = T;
+  using CType = yacl::crypto::SymmetricCrypto::CryptoType;
+
+  YaclReplayUrbg(uint128_t seed, uint64_t ctr, uint64_t iv = 0,
+                 CType ctype = CType::AES128_CTR)
+      : seed_(seed), ctr_(ctr), iv_(iv), ctype_(ctype) {}
+
+  static constexpr T min() { return std::numeric_limits<T>::min(); }
+  static constexpr T max() { return std::numeric_limits<T>::max(); }
+
+  T operator()() {
+    T ret;
+    ctr_ = FillPRand(ctype_, seed_, iv_, ctr_, (char *)&ret, sizeof(T));
+    return ret;
+  }
+
+  uint64_t GetSeed() const { return seed_; }
+  uint64_t GetCounter() const { return ctr_; }
+  uint64_t GetIV() const { return iv_; }
+  CType GetCType() const { return ctype_; }
+
+ private:
+  const uint128_t seed_;
+  uint64_t ctr_;  // NOTE ctr_ is mutable
+  const uint64_t iv_;
+  const CType ctype_;
+};
+
 }  // namespace yacl::crypto

yacl/crypto/rand/rand_test.cc

@@ -14,11 +14,14 @@
 
 #include "yacl/crypto/rand/rand.h"
 
+#include <algorithm>
+#include <cstring>
 #include <limits>
 
 #include "gtest/gtest.h"
 
 #include "yacl/base/int128.h"
+#include "yacl/crypto/block_cipher/symmetric_crypto.h"
 
 namespace yacl::crypto {
 
@@ -105,11 +108,82 @@ TEST(GenericRandTest, RandLtnTest) {
   }
   {
     uint64_t u64max = std::numeric_limits<uint64_t>::max();
+    // should be 170141183460469231731687303715884105727
     uint128_t mersenne_prime = MakeUint128(u64max >> 1, u64max);
-    SPDLOG_INFO(mersenne_prime);  // 170141183460469231731687303715884105727
     auto rand = RandLtN(mersenne_prime);
     EXPECT_TRUE(rand < mersenne_prime);
   }
 }
 
+TEST(GenericRandTest, RandomShuffleTest) {
+  auto vec = FastRandVec<uint128_t>(129);
+  YaclStdUrbg<uint32_t> g;
+  std::shuffle(vec.begin(), vec.end(), g);
+}
+
+TEST(GenericRandTest, ReplayRandomShuffleTest) {
+  int n = 129;
+  auto vec = FastRandVec<uint128_t>(n);
+
+  // same drbg internal states
+  {
+    auto seed = SecureRandSeed();
+    auto ctr = FastRandU64();
+    auto iv = FastRandU64();
+    auto ctype = yacl::crypto::SymmetricCrypto::CryptoType::AES128_CTR;
+    auto vec_copy = vec;
+
+    YaclReplayUrbg<uint32_t> g1(seed, ctr, iv, ctype);
+    std::shuffle(vec.begin(), vec.end(), g1);
+
+    YaclReplayUrbg<uint32_t> g2(seed, ctr, iv, ctype);
+    std::shuffle(vec_copy.begin(), vec_copy.end(), g2);
+
+    EXPECT_EQ(std::memcmp(vec.data(), vec_copy.data(), sizeof(uint128_t) * n),
+              0);
+  }
+
+  // different drbg internal states (seed)
+  {
+    auto seed = SecureRandSeed();
+    auto ctr = FastRandU64();
+    auto iv = FastRandU64();
+    auto ctype = yacl::crypto::SymmetricCrypto::CryptoType::AES128_ECB;
+    auto vec1 = vec;
+    auto vec2 = vec;
+    auto vec3 = vec;
+    auto vec4 = vec;
+
+    YaclReplayUrbg<uint32_t> g1(seed, ctr, iv, ctype);
+    std::shuffle(vec1.begin(), vec1.end(), g1);
+
+    // different seed will almost always result in different shuffles
+    YaclReplayUrbg<uint32_t> g2(seed + 1, ctr, iv, ctype);
+    std::shuffle(vec2.begin(), vec2.end(), g2);
+
+    // NOTE g.GetCounter() will return the after-shuffle prg counter
+    YaclReplayUrbg<uint32_t> g3(seed, g1.GetCounter() + 1, iv, ctype);
+    std::shuffle(vec3.begin(), vec3.end(), g3);
+
+    // NOTE different iv does not gurantee different shuffle, it's
+    // recommended to use different seed to generate different shuffles
+    YaclReplayUrbg<uint32_t> g4(seed, ctr, iv + 1, ctype);
+    std::shuffle(vec4.begin(), vec4.end(), g4);
+
+    // g1 is a random shuffle, different from the original vector
+    EXPECT_NE(std::memcmp(vec.data(), vec1.data(), sizeof(uint128_t) * n), 0);
+
+    // g2 is a different shuffle as g1
+    EXPECT_NE(std::memcmp(vec.data(), vec2.data(), sizeof(uint128_t) * n), 0);
+    EXPECT_NE(std::memcmp(vec1.data(), vec2.data(), sizeof(uint128_t) * n), 0);
+
+    // g3 is a different shuffle as g1
+    EXPECT_NE(std::memcmp(vec.data(), vec3.data(), sizeof(uint128_t) * n), 0);
+    EXPECT_NE(std::memcmp(vec1.data(), vec3.data(), sizeof(uint128_t) * n), 0);
+
+    // NOTE g4 is a SAME shuffle as g1!!!! even though they differ in iv
+    EXPECT_NE(std::memcmp(vec.data(), vec4.data(), sizeof(uint128_t) * n), 0);
+    EXPECT_EQ(std::memcmp(vec1.data(), vec4.data(), sizeof(uint128_t) * n), 0);
+  }
+}
 }  // namespace yacl::crypto

yacl/crypto/tools/prg.h

@@ -41,10 +41,11 @@ namespace yacl::crypto {
 // ---------------------------
 
 // Core implementation of filling deterministic pseudorandomness, return the
-// increased counter (count++, presumably).
-// Note: FillPRand is different from drbg, NIST800-90A since FillPRand will
-// never perform healthcheck, reseed. FillPRand is only an abstract API for the
-// theoretical tool: PRG.
+// increased counter (count++, presumably). FillPRand-like implementations never
+// perform healthcheck, reseed.
+//
+// NOTE FillPRand is not an instantiation of NIST800-90A.
+//
 uint64_t FillPRand(SymmetricCrypto::CryptoType type, uint128_t seed,
                    uint64_t iv, uint64_t count, char* buf, size_t len);