3 files changed, 33 insertions, 9 deletions
diff --git a/src/lib/stream/chacha/chacha.cpp b/src/lib/stream/chacha/chacha.cpp
index 0a32c720b..ac81fd70d 100644
--- a/src/lib/stream/chacha/chacha.cpp
+++ b/src/lib/stream/chacha/chacha.cpp
@@ -10,8 +10,18 @@
 
 namespace Botan {
 
-void ChaCha::chacha(byte output[64], const u32bit input[16])
+ChaCha::ChaCha(size_t rounds) : m_rounds(rounds)
    {
+   if(m_rounds != 12 && m_rounds != 20)
+      throw Invalid_Argument("ChaCha only supports 12 or 20 rounds");
+   }
+
+namespace {
+
+void chacha(byte output[64], const u32bit input[16], size_t rounds)
+   {
+   BOTAN_ASSERT(rounds % 2 == 0, "Valid rounds");
+
    u32bit x00 = input[ 0], x01 = input[ 1], x02 = input[ 2], x03 = input[ 3],
           x04 = input[ 4], x05 = input[ 5], x06 = input[ 6], x07 = input[ 7],
           x08 = input[ 8], x09 = input[ 9], x10 = input[10], x11 = input[11],
@@ -25,7 +35,7 @@ void ChaCha::chacha(byte output[64], const u32bit input[16])
    c += d; b ^= c; b = rotate_left(b, 7);  \
    } while(0)
 
-   for(size_t i = 0; i != 10; ++i)
+   for(size_t i = 0; i != rounds / 2; ++i)
       {
       CHACHA_QUARTER_ROUND(x00, x04, x08, x12);
       CHACHA_QUARTER_ROUND(x01, x05, x09, x13);
@@ -58,6 +68,8 @@ void ChaCha::chacha(byte output[64], const u32bit input[16])
    store_le(x15 + input[15], output + 4 * 15);
    }
 
+}
+
 /*
 * Combine cipher stream with message
 */
@@ -69,7 +81,7 @@ void ChaCha::cipher(const byte in[], byte out[], size_t length)
       length -= (m_buffer.size() - m_position);
       in += (m_buffer.size() - m_position);
       out += (m_buffer.size() - m_position);
-      chacha(m_buffer.data(), m_state.data());
+      chacha(m_buffer.data(), m_state.data(), m_rounds);
 
       ++m_state[12];
       m_state[13] += (m_state[12] == 0);
@@ -142,7 +154,7 @@ void ChaCha::set_iv(const byte iv[], size_t length)
       m_state[15] = load_le<u32bit>(iv, 2);
       }
 
-   chacha(m_buffer.data(), m_state.data());
+   chacha(m_buffer.data(), m_state.data(), m_rounds);
    ++m_state[12];
    m_state[13] += (m_state[12] == 0);
 
@@ -156,4 +168,9 @@ void ChaCha::clear()
    m_position = 0;
    }
 
+std::string ChaCha::name() const
+   {
+   return "ChaCha(" + std::to_string(m_rounds) + ")";
+   }
+
 }
diff --git a/src/lib/stream/chacha/chacha.h b/src/lib/stream/chacha/chacha.h
index 92f8ef035..ba93d6260 100644
--- a/src/lib/stream/chacha/chacha.h
+++ b/src/lib/stream/chacha/chacha.h
@@ -18,6 +18,14 @@ namespace Botan {
 class BOTAN_DLL ChaCha final : public StreamCipher
    {
    public:
+      StreamCipher* clone() const override { return new ChaCha(m_rounds); }
+
+      /**
+      * Currently only 12 or 20 rounds are supported, all others
+      * will throw an exception
+      */
+      ChaCha(size_t rounds);
+
       void cipher(const byte in[], byte out[], size_t length) override;
 
       void set_iv(const byte iv[], size_t iv_len) override;
@@ -31,14 +39,13 @@ class BOTAN_DLL ChaCha final : public StreamCipher
          }
 
       void clear() override;
-      std::string name() const override { return "ChaCha"; }
 
-      StreamCipher* clone() const override { return new ChaCha; }
-   protected:
-      virtual void chacha(byte output[64], const u32bit input[16]);
+      std::string name() const override;
+
    private:
       void key_schedule(const byte key[], size_t key_len) override;
 
+      size_t m_rounds;
       secure_vector<u32bit> m_state;
       secure_vector<byte> m_buffer;
       size_t m_position = 0;
diff --git a/src/lib/stream/stream_cipher.cpp b/src/lib/stream/stream_cipher.cpp
index 03ef5e329..6f98df1fb 100644
--- a/src/lib/stream/stream_cipher.cpp
+++ b/src/lib/stream/stream_cipher.cpp
@@ -51,7 +51,7 @@ void StreamCipher::set_iv(const byte[], size_t iv_len)
    }
 
 #if defined(BOTAN_HAS_CHACHA)
-BOTAN_REGISTER_T_NOARGS(StreamCipher, ChaCha);
+BOTAN_REGISTER_T_1LEN(StreamCipher, ChaCha, 20);
 #endif
 
 #if defined(BOTAN_HAS_SALSA20)