Abstract out seeding/extracting entropy into RNGState::MixExtract

2 files changed, 37 insertions, 25 deletions

diff --git a/src/crypto/sha512.h b/src/crypto/sha512.h
index cd1023bc85..4118ac1b18 100644
--- a/src/crypto/sha512.h
+++ b/src/crypto/sha512.h
@@ -17,7 +17,7 @@ private:
     uint64_t bytes;
 
 public:
-    static const size_t OUTPUT_SIZE = 64;
+    static constexpr size_t OUTPUT_SIZE = 64;
 
     CSHA512();
     CSHA512& Write(const unsigned char* data, size_t len);
diff --git a/src/random.cpp b/src/random.cpp
index 6b7962aa13..c78848d56e 100644
--- a/src/random.cpp
+++ b/src/random.cpp
@@ -304,6 +304,34 @@ struct RNGState {
     {
         InitHardwareRand();
     }
+
+    /** Extract up to 32 bytes of entropy from the RNG state, mixing in new entropy from hasher. */
+    void MixExtract(unsigned char* out, size_t num, CSHA512&& hasher)
+    {
+        assert(num <= 32);
+        unsigned char buf[64];
+        static_assert(sizeof(buf) == CSHA512::OUTPUT_SIZE, "Buffer needs to have hasher's output size");
+        {
+            LOCK(m_mutex);
+            // Write the current state of the RNG into the hasher
+            hasher.Write(m_state, 32);
+            // Write a new counter number into the state
+            hasher.Write((const unsigned char*)&m_counter, sizeof(m_counter));
+            ++m_counter;
+            // Finalize the hasher
+            hasher.Finalize(buf);
+            // Store the last 32 bytes of the hash output as new RNG state.
+            memcpy(m_state, buf + 32, 32);
+        }
+        // If desired, copy (up to) the first 32 bytes of the hash output as output.
+        if (num) {
+            assert(out != nullptr);
+            memcpy(out, buf, num);
+        }
+        // Best effort cleanup of internal state
+        hasher.Reset();
+        memory_cleanse(buf, 64);
+    }
 };
 
 RNGState& GetRNGState()
@@ -315,38 +343,29 @@ RNGState& GetRNGState()
 }
 }
 
-static void AddDataToRng(void* data, size_t len);
+static void AddDataToRng(void* data, size_t len, RNGState& rng);
 
 void RandAddSeedSleep()
 {
+    RNGState& rng = GetRNGState();
+
     int64_t nPerfCounter1 = GetPerformanceCounter();
     std::this_thread::sleep_for(std::chrono::milliseconds(1));
     int64_t nPerfCounter2 = GetPerformanceCounter();
 
     // Combine with and update state
-    AddDataToRng(&nPerfCounter1, sizeof(nPerfCounter1));
-    AddDataToRng(&nPerfCounter2, sizeof(nPerfCounter2));
+    AddDataToRng(&nPerfCounter1, sizeof(nPerfCounter1), rng);
+    AddDataToRng(&nPerfCounter2, sizeof(nPerfCounter2), rng);
 
     memory_cleanse(&nPerfCounter1, sizeof(nPerfCounter1));
     memory_cleanse(&nPerfCounter2, sizeof(nPerfCounter2));
 }
 
-static void AddDataToRng(void* data, size_t len) {
-    RNGState& rng = GetRNGState();
-
+static void AddDataToRng(void* data, size_t len, RNGState& rng) {
     CSHA512 hasher;
     hasher.Write((const unsigned char*)&len, sizeof(len));
     hasher.Write((const unsigned char*)data, len);
-    unsigned char buf[64];
-    {
-        WAIT_LOCK(rng.m_mutex, lock);
-        hasher.Write(rng.m_state, sizeof(rng.m_state));
-        hasher.Write((const unsigned char*)&rng.m_counter, sizeof(rng.m_counter));
-        ++rng.m_counter;
-        hasher.Finalize(buf);
-        memcpy(rng.m_state, buf + 32, 32);
-    }
-    memory_cleanse(buf, 64);
+    rng.MixExtract(nullptr, 0, std::move(hasher));
 }
 
 void GetStrongRandBytes(unsigned char* out, int num)
@@ -372,14 +391,7 @@ void GetStrongRandBytes(unsigned char* out, int num)
     }
 
     // Combine with and update state
-    {
-        WAIT_LOCK(rng.m_mutex, lock);
-        hasher.Write(rng.m_state, sizeof(rng.m_state));
-        hasher.Write((const unsigned char*)&rng.m_counter, sizeof(rng.m_counter));
-        ++rng.m_counter;
-        hasher.Finalize(buf);
-        memcpy(rng.m_state, buf + 32, 32);
-    }
+    rng.MixExtract(out, num, std::move(hasher));
 
     // Produce output
     memcpy(out, buf, num);