Make SimpleRandom faster by using Chacha20 as underlying primitive.

This does alter the output of the RNG, but that's a one-time cost, and one SHA-1 call per byte is highly suboptimal for anything but very short unit tests.

gfe-relnote: n/a (test-only change)
PiperOrigin-RevId: 245501012
Change-Id: I16d9f5b174ef65cfb4ece3ef0699bf1b29b81de5
diff --git a/quic/test_tools/quic_test_utils.cc b/quic/test_tools/quic_test_utils.cc
index cacb28a..52634b7 100644
--- a/quic/test_tools/quic_test_utils.cc
+++ b/quic/test_tools/quic_test_utils.cc
@@ -5,8 +5,10 @@
 #include "net/third_party/quiche/src/quic/test_tools/quic_test_utils.h"
 
 #include <algorithm>
+#include <cstdint>
 #include <memory>
 
+#include "third_party/boringssl/src/include/openssl/chacha.h"
 #include "third_party/boringssl/src/include/openssl/sha.h"
 #include "net/third_party/quiche/src/quic/core/crypto/crypto_framer.h"
 #include "net/third_party/quiche/src/quic/core/crypto/crypto_handshake.h"
@@ -152,20 +154,42 @@
 }
 
 uint64_t SimpleRandom::RandUint64() {
-  std::string hash =
-      Sha1Hash(QuicStringPiece(reinterpret_cast<char*>(&seed_), sizeof(seed_)));
-  DCHECK_EQ(static_cast<size_t>(SHA_DIGEST_LENGTH), hash.length());
-  memcpy(&seed_, hash.data(), sizeof(seed_));
-  return seed_;
+  uint64_t result;
+  RandBytes(&result, sizeof(result));
+  return result;
 }
 
 void SimpleRandom::RandBytes(void* data, size_t len) {
-  uint8_t* real_data = static_cast<uint8_t*>(data);
-  for (size_t offset = 0; offset < len; offset++) {
-    real_data[offset] = RandUint64() & 0xff;
+  uint8_t* data_bytes = reinterpret_cast<uint8_t*>(data);
+  while (len > 0) {
+    const size_t buffer_left = sizeof(buffer_) - buffer_offset_;
+    const size_t to_copy = std::min(buffer_left, len);
+    memcpy(data_bytes, buffer_ + buffer_offset_, to_copy);
+    data_bytes += to_copy;
+    buffer_offset_ += to_copy;
+    len -= to_copy;
+
+    if (buffer_offset_ == sizeof(buffer_)) {
+      FillBuffer();
+    }
   }
 }
 
+void SimpleRandom::FillBuffer() {
+  uint8_t nonce[12];
+  memcpy(nonce, buffer_, sizeof(nonce));
+  CRYPTO_chacha_20(buffer_, buffer_, sizeof(buffer_), key_, nonce, 0);
+  buffer_offset_ = 0;
+}
+
+void SimpleRandom::set_seed(uint64_t seed) {
+  static_assert(sizeof(key_) == SHA256_DIGEST_LENGTH, "Key has to be 256 bits");
+  SHA256(reinterpret_cast<const uint8_t*>(&seed), sizeof(seed), key_);
+
+  memset(buffer_, 0, sizeof(buffer_));
+  FillBuffer();
+}
+
 MockFramerVisitor::MockFramerVisitor() {
   // By default, we want to accept packets.
   ON_CALL(*this, OnProtocolVersionMismatch(_, _))
diff --git a/quic/test_tools/quic_test_utils.h b/quic/test_tools/quic_test_utils.h
index 57831ab..c9db3e5 100644
--- a/quic/test_tools/quic_test_utils.h
+++ b/quic/test_tools/quic_test_utils.h
@@ -213,12 +213,10 @@
 std::string Sha1Hash(QuicStringPiece data);
 
 // Simple random number generator used to compute random numbers suitable
-// for pseudo-randomly dropping packets in tests.  It works by computing
-// the sha1 hash of the current seed, and using the first 64 bits as
-// the next random number, and the next seed.
+// for pseudo-randomly dropping packets in tests.
 class SimpleRandom : public QuicRandom {
  public:
-  SimpleRandom() : seed_(0) {}
+  SimpleRandom() { set_seed(0); }
   SimpleRandom(const SimpleRandom&) = delete;
   SimpleRandom& operator=(const SimpleRandom&) = delete;
   ~SimpleRandom() override {}
@@ -228,10 +226,14 @@
 
   void RandBytes(void* data, size_t len) override;
 
-  void set_seed(uint64_t seed) { seed_ = seed; }
+  void set_seed(uint64_t seed);
 
  private:
-  uint64_t seed_;
+  uint8_t buffer_[4096];
+  size_t buffer_offset_;
+  uint8_t key_[32];
+
+  void FillBuffer();
 };
 
 class MockFramerVisitor : public QuicFramerVisitorInterface {
diff --git a/quic/test_tools/quic_test_utils_test.cc b/quic/test_tools/quic_test_utils_test.cc
index aeccc70..8c30335 100644
--- a/quic/test_tools/quic_test_utils_test.cc
+++ b/quic/test_tools/quic_test_utils_test.cc
@@ -52,8 +52,27 @@
 TEST_F(QuicTestUtilsTest, SimpleRandomStability) {
   SimpleRandom rng;
   rng.set_seed(UINT64_C(0x1234567800010001));
-  EXPECT_EQ(UINT64_C(14865409841904857791), rng.RandUint64());
-  EXPECT_EQ(UINT64_C(12139094019410129741), rng.RandUint64());
+  EXPECT_EQ(UINT64_C(12589383305231984671), rng.RandUint64());
+  EXPECT_EQ(UINT64_C(17775425089941798664), rng.RandUint64());
+}
+
+// Ensure that the output of SimpleRandom does not depend on the size of the
+// read calls.
+TEST_F(QuicTestUtilsTest, SimpleRandomChunks) {
+  SimpleRandom rng;
+  std::string reference(16 * 1024, '\0');
+  rng.RandBytes(&reference[0], reference.size());
+
+  for (size_t chunk_size : {3, 4, 7, 4096}) {
+    rng.set_seed(0);
+    size_t chunks = reference.size() / chunk_size;
+    std::string buffer(chunks * chunk_size, '\0');
+    for (size_t i = 0; i < chunks; i++) {
+      rng.RandBytes(&buffer[i * chunk_size], chunk_size);
+    }
+    EXPECT_EQ(reference.substr(0, buffer.size()), buffer)
+        << "Failed for chunk_size = " << chunk_size;
+  }
 }
 
 }  // namespace test