1 files changed, 78 insertions, 19 deletions

diff --git a/src/tests/test_rng.cpp b/src/tests/test_rng.cpp
index 84e291c6a..6d10de9a7 100644
--- a/src/tests/test_rng.cpp
+++ b/src/tests/test_rng.cpp
@@ -6,29 +6,51 @@
 #include <iostream>
 #include <fstream>
 
+#if defined(BOTAN_HAS_HMAC_DRBG)
+  #include <botan/hmac_drbg.h>
+#endif
+
 #if defined(BOTAN_HAS_X931_RNG)
   #include <botan/x931_rng.h>
-  #include <botan/aes.h>
-  #include <botan/des.h>
 #endif
 
 using namespace Botan;
 
 namespace {
 
-RandomNumberGenerator* get_rng(const std::string& algo, const std::string& ikm_hex)
+RandomNumberGenerator* get_rng(const std::string& algo_str, const std::string& ikm_hex)
    {
+   class AllOnce_RNG : public Fixed_Output_RNG
+      {
+      public:
+         AllOnce_RNG(const std::vector<byte>& in) : Fixed_Output_RNG(in) {}
+
+         Botan::secure_vector<byte> random_vec(size_t)
+            {
+            Botan::secure_vector<byte> vec(this->remaining());
+            this->randomize(&vec[0], vec.size());
+            return vec;
+            }
+      };
+
    const auto ikm = hex_decode(ikm_hex);
 
+   Algorithm_Factory& af = global_state().algorithm_factory();
+
+   const auto algo_name = parse_algorithm_name(algo_str);
+
+   const std::string rng_name = algo_name[0];
+
+#if defined(BOTAN_HAS_HMAC_DRBG)
+   if(rng_name == "HMAC_DRBG")
+      return new HMAC_DRBG(af.make_mac("HMAC(" + algo_name[1] + ")"),
+                               new AllOnce_RNG(ikm));
+#endif
+
 #if defined(BOTAN_HAS_X931_RNG)
-   if(algo == "X9.31-RNG(TripleDES)")
-      return new ANSI_X931_RNG(new TripleDES, new Fixed_Output_RNG(ikm));
-   else if(algo == "X9.31-RNG(AES-128)")
-      return new ANSI_X931_RNG(new AES_128, new Fixed_Output_RNG(ikm));
-   else if(algo == "X9.31-RNG(AES-192)")
-      return new ANSI_X931_RNG(new AES_192, new Fixed_Output_RNG(ikm));
-   else if(algo == "X9.31-RNG(AES-256)")
-      return new ANSI_X931_RNG(new AES_256, new Fixed_Output_RNG(ikm));
+   if(rng_name == "X9.31-RNG")
+      return new ANSI_X931_RNG(af.make_block_cipher(algo_name[1]),
+                               new Fixed_Output_RNG(ikm));
 #endif
 
    return nullptr;
@@ -44,8 +66,6 @@ size_t x931_test(const std::string& algo,
    if(!rng)
       throw std::runtime_error("Unknown RNG " + algo);
 
-   rng->reseed(0);
-
    const std::string got = hex_encode(rng->random_vec(L));
 
    if(got != out)
@@ -57,15 +77,54 @@ size_t x931_test(const std::string& algo,
    return 0;
    }
 
+size_t hmac_drbg_test(std::map<std::string, std::string> m)
+   {
+   const std::string algo = m["RNG"];
+   const std::string ikm = m["EntropyInput"];
+
+   std::unique_ptr<RandomNumberGenerator> rng(get_rng(algo, ikm));
+   if(!rng)
+      throw std::runtime_error("Unknown RNG " + algo);
+
+   rng->reseed(0); // force initialization
+
+   // now reseed
+   const auto reseed_input = hex_decode(m["EntropyInputReseed"]);
+   rng->add_entropy(&reseed_input[0], reseed_input.size());
+
+   const std::string out = m["Out"];
+
+   const size_t out_len = out.size() / 2;
+
+   rng->random_vec(out_len); // gen 1st block (discarded)
+
+   const std::string got = hex_encode(rng->random_vec(out_len));
+
+   if(got != out)
+      {
+      std::cout << algo << " " << got << " != " << out << "\n";
+      return 1;
+      }
+
+   return 0;
+   }
+
 }
 
 size_t test_rngs()
    {
-   std::ifstream vec(TEST_DATA_DIR "/x931.vec");
+   std::ifstream hmac_drbg_vec(TEST_DATA_DIR "/hmac_drbg.vec");
+   std::ifstream x931_vec(TEST_DATA_DIR "/x931.vec");
+
+   size_t fails = 0;
+
+   fails += run_tests_bb(hmac_drbg_vec, "RNG", "Out", true, hmac_drbg_test);
+
+   fails += run_tests_bb(x931_vec, "RNG", "Out", true,
+                         [](std::map<std::string, std::string> m) -> size_t
+                         {
+                         return x931_test(m["RNG"], m["IKM"], m["Out"], to_u32bit(m["L"]));
+                         });
 
-   return run_tests_bb(vec, "RNG", "Out", true,
-             [](std::map<std::string, std::string> m) -> size_t
-             {
-             return x931_test(m["RNG"], m["IKM"], m["Out"], to_u32bit(m["L"]));
-             });
+   return fails;
    }