5 files changed, 37 insertions, 7 deletions
diff --git a/src/lib/modes/aead/chacha20poly1305/chacha20poly1305.cpp b/src/lib/modes/aead/chacha20poly1305/chacha20poly1305.cpp
index 786e21def..007e2fbe4 100644
--- a/src/lib/modes/aead/chacha20poly1305/chacha20poly1305.cpp
+++ b/src/lib/modes/aead/chacha20poly1305/chacha20poly1305.cpp
@@ -66,11 +66,11 @@ void ChaCha20Poly1305_Mode::start_msg(const uint8_t nonce[], size_t nonce_len)
 
    m_chacha->set_iv(nonce, nonce_len);
 
-   secure_vector<uint8_t> init(64); // zeros
-   m_chacha->encrypt(init);
+   secure_vector<uint8_t> first_block(64);
+   m_chacha->write_keystream(first_block.data(), first_block.size());
 
-   m_poly1305->set_key(init.data(), 32);
-   // Remainder of output is discard
+   m_poly1305->set_key(first_block.data(), 32);
+   // Remainder of first block is discarded
 
    m_poly1305->update(m_ad);
 
diff --git a/src/lib/rng/chacha_rng/chacha_rng.cpp b/src/lib/rng/chacha_rng/chacha_rng.cpp
index ad8ee9ba8..065a19060 100644
--- a/src/lib/rng/chacha_rng/chacha_rng.cpp
+++ b/src/lib/rng/chacha_rng/chacha_rng.cpp
@@ -75,8 +75,7 @@ void ChaCha_RNG::randomize_with_input(uint8_t output[], size_t output_len,
       update(input, input_len);
       }
 
-   clear_mem(output, output_len);
-   m_chacha->cipher1(output, output_len);
+   m_chacha->write_keystream(output, output_len);
    }
 
 void ChaCha_RNG::update(const uint8_t input[], size_t input_len)
@@ -85,7 +84,7 @@ void ChaCha_RNG::update(const uint8_t input[], size_t 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_chacha->write_keystream(mac_key.data(), mac_key.size());
    m_hmac->set_key(mac_key);
    }
 
diff --git a/src/lib/stream/chacha/chacha.cpp b/src/lib/stream/chacha/chacha.cpp
index d2d31a12e..8edb685da 100644
--- a/src/lib/stream/chacha/chacha.cpp
+++ b/src/lib/stream/chacha/chacha.cpp
@@ -173,6 +173,24 @@ void ChaCha::cipher(const uint8_t in[], uint8_t out[], size_t length)
    m_position += length;
    }
 
+void ChaCha::write_keystream(uint8_t out[], size_t length)
+   {
+   verify_key_set(m_state.empty() == false);
+
+   while(length >= m_buffer.size() - m_position)
+      {
+      copy_mem(out, &m_buffer[m_position], m_buffer.size() - m_position);
+      length -= (m_buffer.size() - m_position);
+      out += (m_buffer.size() - m_position);
+      chacha_x4(m_buffer.data(), m_state.data(), m_rounds);
+      m_position = 0;
+      }
+
+   copy_mem(out, &m_buffer[m_position], length);
+
+   m_position += length;
+   }
+
 void ChaCha::initialize_state()
    {
    static const uint32_t TAU[] =
diff --git a/src/lib/stream/chacha/chacha.h b/src/lib/stream/chacha/chacha.h
index e41fd927f..346e25c28 100644
--- a/src/lib/stream/chacha/chacha.h
+++ b/src/lib/stream/chacha/chacha.h
@@ -29,6 +29,8 @@ class BOTAN_PUBLIC_API(2,0) ChaCha final : public StreamCipher
 
       void cipher(const uint8_t in[], uint8_t out[], size_t length) override;
 
+      void write_keystream(uint8_t out[], size_t len) override;
+
       void set_iv(const uint8_t iv[], size_t iv_len) override;
 
       /*
diff --git a/src/lib/stream/stream_cipher.h b/src/lib/stream/stream_cipher.h
index fee37f4e4..25bbc3287 100644
--- a/src/lib/stream/stream_cipher.h
+++ b/src/lib/stream/stream_cipher.h
@@ -58,6 +58,17 @@ class BOTAN_PUBLIC_API(2,0) StreamCipher : public SymmetricAlgorithm
       virtual void cipher(const uint8_t in[], uint8_t out[], size_t len) = 0;
 
       /**
+      * Write keystream bytes to a buffer
+      * @param out the byte array to hold the keystream
+      * @param len the length of out in bytes
+      */
+      virtual void write_keystream(uint8_t out[], size_t len)
+         {
+         clear_mem(out, len);
+         cipher1(out, len);
+         }
+
+      /**
       * Encrypt or decrypt a message
       * The message is encrypted/decrypted in place.
       * @param buf the plaintext / ciphertext