Refactor stateful RNG tests to cover both HMAC_DRBG and ChaCha_RNG

2 files changed, 21 insertions, 14 deletions

diff --git a/src/lib/rng/chacha_rng/chacha_rng.cpp b/src/lib/rng/chacha_rng/chacha_rng.cpp
index 86c71f9fe..ceccfef2c 100644
--- a/src/lib/rng/chacha_rng/chacha_rng.cpp
+++ b/src/lib/rng/chacha_rng/chacha_rng.cpp
@@ -57,7 +57,6 @@ void ChaCha_RNG::clear()
    Stateful_RNG::clear();
 
    m_hmac->set_key(std::vector<uint8_t>(m_hmac->output_length(), 0x00));
-   m_chacha->set_key(m_hmac->final());
    }
 
 void ChaCha_RNG::randomize(uint8_t output[], size_t output_len)
@@ -68,28 +67,34 @@ void ChaCha_RNG::randomize(uint8_t output[], size_t output_len)
 void ChaCha_RNG::randomize_with_input(uint8_t output[], size_t output_len,
                                       const uint8_t input[], size_t input_len)
    {
-   add_entropy(input, input_len);
    reseed_check();
 
+   if(input_len > 0)
+      {
+      update(input, input_len);
+      }
+
    clear_mem(output, output_len);
    m_chacha->cipher1(output, output_len);
    }
 
-void ChaCha_RNG::add_entropy(const uint8_t input[], size_t input_len)
+void ChaCha_RNG::update(const uint8_t input[], size_t input_len)
    {
-   if(input_len > 0)
-      {
-      m_hmac->update(input, input_len);
-      m_chacha->set_key(m_hmac->final());
+   m_hmac->update(input, input_len);
+   m_chacha->set_key(m_hmac->final());
 
-      secure_vector<uint8_t> mac_key(m_hmac->output_length());
-      m_chacha->cipher1(mac_key.data(), mac_key.size());
-      m_hmac->set_key(mac_key);
+   secure_vector<uint8_t> mac_key(m_hmac->output_length());
+   m_chacha->cipher1(mac_key.data(), mac_key.size());
+   m_hmac->set_key(mac_key);
+   }
 
-      if(8*input_len >= security_level())
-         {
-         m_reseed_counter = 1;
-         }
+void ChaCha_RNG::add_entropy(const uint8_t input[], size_t input_len)
+   {
+   update(input, input_len);
+
+   if(8*input_len >= security_level())
+      {
+      m_reseed_counter = 1;
       }
    }
 
diff --git a/src/lib/rng/chacha_rng/chacha_rng.h b/src/lib/rng/chacha_rng/chacha_rng.h
index b6a763f62..7deaa2d89 100644
--- a/src/lib/rng/chacha_rng/chacha_rng.h
+++ b/src/lib/rng/chacha_rng/chacha_rng.h
@@ -115,6 +115,8 @@ class BOTAN_DLL ChaCha_RNG final : public Stateful_RNG
 
       size_t security_level() const override;
 
+      size_t max_number_of_bytes_per_request() const override { return 0; }
+
    private:
       void update(const uint8_t input[], size_t input_len);