Move poll logic to Entropy_Sources

3 files changed, 27 insertions, 18 deletions

diff --git a/src/lib/entropy/entropy_src.h b/src/lib/entropy/entropy_src.h
index 539df809a..1726d05b8 100644
--- a/src/lib/entropy/entropy_src.h
+++ b/src/lib/entropy/entropy_src.h
@@ -1,6 +1,6 @@
 /*
 * EntropySource
-* (C) 2008,2009,2014,2015 Jack Lloyd
+* (C) 2008,2009,2014,2015,2016 Jack Lloyd
 *
 * Botan is released under the Simplified BSD License (see license.txt)
 */
@@ -10,10 +10,13 @@
 
 #include <botan/secmem.h>
 #include <string>
+#include <chrono>
 #include <functional>
 
 namespace Botan {
 
+class RandomNumberGenerator;
+
 /**
 * Class used to accumulate the poll results of EntropySources
 */
@@ -109,7 +112,10 @@ class BOTAN_DLL Entropy_Sources final
 
       std::vector<std::string> enabled_sources() const;
 
-      void poll(Entropy_Accumulator& accum);
+      size_t poll(RandomNumberGenerator& rng,
+                  size_t bits,
+                  std::chrono::milliseconds timeout);
+
       bool poll_just(Entropy_Accumulator& accum, const std::string& src);
 
       Entropy_Sources() {}
diff --git a/src/lib/entropy/entropy_srcs.cpp b/src/lib/entropy/entropy_srcs.cpp
index a5dc0a819..81076b0e9 100644
--- a/src/lib/entropy/entropy_srcs.cpp
+++ b/src/lib/entropy/entropy_srcs.cpp
@@ -6,6 +6,7 @@
 */
 
 #include <botan/entropy_src.h>
+#include <botan/rng.h>
 
 #if defined(BOTAN_HAS_ENTROPY_SRC_HIGH_RESOLUTION_TIMER)
   #include <botan/internal/hres_timer.h>
@@ -154,14 +155,30 @@ std::vector<std::string> Entropy_Sources::enabled_sources() const
    return sources;
    }
 
-void Entropy_Sources::poll(Entropy_Accumulator& accum)
+size_t Entropy_Sources::poll(RandomNumberGenerator& rng,
+                             size_t poll_bits,
+                             std::chrono::milliseconds timeout)
    {
+   typedef std::chrono::system_clock clock;
+
+   auto deadline = clock::now() + timeout;
+
+   double bits_collected = 0;
+
+   Entropy_Accumulator accum([&](const byte in[], size_t in_len, double entropy_estimate) {
+      rng.add_entropy(in, in_len);
+      bits_collected += entropy_estimate;
+      return (bits_collected >= poll_bits || clock::now() > deadline);
+      });
+
    for(size_t i = 0; i != m_srcs.size(); ++i)
       {
       m_srcs[i]->poll(accum);
       if(accum.polling_goal_achieved())
          break;
       }
+
+   return static_cast<size_t>(bits_collected);
    }
 
 bool Entropy_Sources::poll_just(Entropy_Accumulator& accum, const std::string& the_src)
diff --git a/src/lib/rng/rng.cpp b/src/lib/rng/rng.cpp
index 8144ac293..923b417dc 100644
--- a/src/lib/rng/rng.cpp
+++ b/src/lib/rng/rng.cpp
@@ -39,21 +39,7 @@ size_t RandomNumberGenerator::reseed_with_sources(Entropy_Sources& srcs,
                                                   size_t poll_bits,
                                                   std::chrono::milliseconds poll_timeout)
    {
-   typedef std::chrono::system_clock clock;
-
-   auto deadline = clock::now() + poll_timeout;
-
-   double bits_collected = 0;
-
-   Entropy_Accumulator accum([&](const byte in[], size_t in_len, double entropy_estimate) {
-      add_entropy(in, in_len);
-      bits_collected += entropy_estimate;
-      return (bits_collected >= poll_bits || clock::now() > deadline);
-      });
-
-   srcs.poll(accum);
-
-   return bits_collected;
+   return srcs.poll(*this, poll_bits, poll_timeout);
    }
 
 Stateful_RNG::Stateful_RNG(size_t bytes_before_reseed) : m_bytes_before_reseed(bytes_before_reseed)