1 files changed, 21 insertions, 4 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) + ")";
+   }
+
 }