1 files changed, 46 insertions, 38 deletions

diff --git a/src/lib/pbkdf/pbkdf2/pbkdf2.cpp b/src/lib/pbkdf/pbkdf2/pbkdf2.cpp
index 146dd15b0..0ff412bc5 100644
--- a/src/lib/pbkdf/pbkdf2/pbkdf2.cpp
+++ b/src/lib/pbkdf/pbkdf2/pbkdf2.cpp
@@ -27,51 +27,49 @@ PKCS5_PBKDF2* PKCS5_PBKDF2::make(const Spec& spec)
    return nullptr;
    }
 
-/*
-* Return a PKCS #5 PBKDF2 derived key
-*/
-std::pair<size_t, OctetString>
-PKCS5_PBKDF2::key_derivation(size_t key_len,
-                             const std::string& passphrase,
-                             const byte salt[], size_t salt_len,
-                             size_t iterations,
-                             std::chrono::milliseconds msec) const
+size_t
+pbkdf2(MessageAuthenticationCode& prf,
+       byte out[],
+       size_t out_len,
+       const std::string& passphrase,
+       const byte salt[], size_t salt_len,
+       size_t iterations,
+       std::chrono::milliseconds msec)
    {
-   if(key_len == 0)
-      return std::make_pair(iterations, OctetString());
+   clear_mem(out, out_len);
+
+   if(out_len == 0)
+      return 0;
 
    try
       {
-      mac->set_key(reinterpret_cast<const byte*>(passphrase.data()),
-                   passphrase.length());
+      prf.set_key(reinterpret_cast<const byte*>(passphrase.data()), passphrase.size());
       }
    catch(Invalid_Key_Length)
       {
-      throw Exception(name() + " cannot accept passphrases of length " +
-                      std::to_string(passphrase.length()));
+      throw std::runtime_error("PBKDF2 with " + prf.name() +
+                               " cannot accept passphrases of length " +
+                               std::to_string(passphrase.size()));
       }
 
-   secure_vector<byte> key(key_len);
-
-   byte* T = &key[0];
-
-   secure_vector<byte> U(mac->output_length());
+   const size_t prf_sz = prf.output_length();
+   secure_vector<byte> U(prf_sz);
 
-   const size_t blocks_needed = round_up(key_len, mac->output_length()) / mac->output_length();
+   const size_t blocks_needed = round_up(out_len, prf_sz) / prf_sz;
 
    std::chrono::microseconds usec_per_block =
       std::chrono::duration_cast<std::chrono::microseconds>(msec) / blocks_needed;
 
    u32bit counter = 1;
-   while(key_len)
+   while(out_len)
       {
-      size_t T_size = std::min<size_t>(mac->output_length(), key_len);
+      const size_t prf_output = std::min<size_t>(prf_sz, out_len);
 
-      mac->update(salt, salt_len);
-      mac->update_be(counter);
-      mac->final(&U[0]);
+      prf.update(salt, salt_len);
+      prf.update_be(counter++);
+      prf.final(&U[0]);
 
-      xor_buf(T, &U[0], T_size);
+      xor_buf(out, &U[0], prf_output);
 
       if(iterations == 0)
          {
@@ -86,9 +84,9 @@ PKCS5_PBKDF2::key_derivation(size_t key_len,
 
          while(true)
             {
-            mac->update(U);
-            mac->final(&U[0]);
-            xor_buf(T, &U[0], T_size);
+            prf.update(U);
+            prf.final(&U[0]);
+            xor_buf(out, &U[0], prf_output);
             iterations++;
 
             /*
@@ -96,7 +94,7 @@ PKCS5_PBKDF2::key_derivation(size_t key_len,
             avoids confusion, and likely some broken implementations
             break on getting completely randomly distributed values
             */
-            if(iterations % 1000 == 0)
+            if(iterations % 10000 == 0)
                {
                auto time_taken = std::chrono::high_resolution_clock::now() - start;
                auto usec_taken = std::chrono::duration_cast<std::chrono::microseconds>(time_taken);
@@ -109,18 +107,28 @@ PKCS5_PBKDF2::key_derivation(size_t key_len,
          {
          for(size_t i = 1; i != iterations; ++i)
             {
-            mac->update(U);
-            mac->final(&U[0]);
-            xor_buf(T, &U[0], T_size);
+            prf.update(U);
+            prf.final(&U[0]);
+            xor_buf(out, &U[0], prf_output);
             }
          }
 
-      key_len -= T_size;
-      T += T_size;
-      ++counter;
+      out_len -= prf_output;
+      out += prf_output;
       }
 
-   return std::make_pair(iterations, key);
+   return iterations;
    }
 
+size_t
+PKCS5_PBKDF2::pbkdf(byte key[], size_t key_len,
+                    const std::string& passphrase,
+                    const byte salt[], size_t salt_len,
+                    size_t iterations,
+                    std::chrono::milliseconds msec) const
+   {
+   return pbkdf2(*mac.get(), key, key_len, passphrase, salt, salt_len, iterations, msec);
+   }
+
+
 }