Split ChaCha20 into aligned/unaligned variants

2 files changed, 97 insertions, 45 deletions

diff --git a/src/crypto/chacha20.cpp b/src/crypto/chacha20.cpp
index 25d7baa8cc..cdeeee192e 100644
--- a/src/crypto/chacha20.cpp
+++ b/src/crypto/chacha20.cpp
@@ -8,6 +8,7 @@
 #include <crypto/common.h>
 #include <crypto/chacha20.h>
 
+#include <algorithm>
 #include <string.h>
 
 constexpr static inline uint32_t rotl32(uint32_t v, int c) { return (v << c) | (v >> (32 - c)); }
@@ -23,7 +24,7 @@ constexpr static inline uint32_t rotl32(uint32_t v, int c) { return (v << c) | (
 static const unsigned char sigma[] = "expand 32-byte k";
 static const unsigned char tau[] = "expand 16-byte k";
 
-void ChaCha20::SetKey(const unsigned char* k, size_t keylen)
+void ChaCha20Aligned::SetKey(const unsigned char* k, size_t keylen)
 {
     const unsigned char *constants;
 
@@ -51,37 +52,34 @@ void ChaCha20::SetKey(const unsigned char* k, size_t keylen)
     input[15] = 0;
 }
 
-ChaCha20::ChaCha20()
+ChaCha20Aligned::ChaCha20Aligned()
 {
     memset(input, 0, sizeof(input));
 }
 
-ChaCha20::ChaCha20(const unsigned char* k, size_t keylen)
+ChaCha20Aligned::ChaCha20Aligned(const unsigned char* k, size_t keylen)
 {
     SetKey(k, keylen);
 }
 
-void ChaCha20::SetIV(uint64_t iv)
+void ChaCha20Aligned::SetIV(uint64_t iv)
 {
     input[14] = iv;
     input[15] = iv >> 32;
 }
 
-void ChaCha20::Seek(uint64_t pos)
+void ChaCha20Aligned::Seek(uint64_t pos)
 {
     input[12] = pos;
     input[13] = pos >> 32;
 }
 
-void ChaCha20::Keystream(unsigned char* c, size_t bytes)
+inline void ChaCha20Aligned::Keystream64(unsigned char* c, size_t blocks)
 {
     uint32_t x0, x1, x2, x3, x4, x5, x6, x7, x8, x9, x10, x11, x12, x13, x14, x15;
     uint32_t j0, j1, j2, j3, j4, j5, j6, j7, j8, j9, j10, j11, j12, j13, j14, j15;
-    unsigned char *ctarget = nullptr;
-    unsigned char tmp[64];
-    unsigned int i;
 
-    if (!bytes) return;
+    if (!blocks) return;
 
     j0 = input[0];
     j1 = input[1];
@@ -101,10 +99,6 @@ void ChaCha20::Keystream(unsigned char* c, size_t bytes)
     j15 = input[15];
 
     for (;;) {
-        if (bytes < 64) {
-            ctarget = c;
-            c = tmp;
-        }
         x0 = j0;
         x1 = j1;
         x2 = j2;
@@ -171,28 +165,22 @@ void ChaCha20::Keystream(unsigned char* c, size_t bytes)
         WriteLE32(c + 56, x14);
         WriteLE32(c + 60, x15);
 
-        if (bytes <= 64) {
-            if (bytes < 64) {
-                for (i = 0;i < bytes;++i) ctarget[i] = c[i];
-            }
+        if (blocks == 1) {
             input[12] = j12;
             input[13] = j13;
             return;
         }
-        bytes -= 64;
+        blocks -= 1;
         c += 64;
     }
 }
 
-void ChaCha20::Crypt(const unsigned char* m, unsigned char* c, size_t bytes)
+inline void ChaCha20Aligned::Crypt64(const unsigned char* m, unsigned char* c, size_t blocks)
 {
     uint32_t x0, x1, x2, x3, x4, x5, x6, x7, x8, x9, x10, x11, x12, x13, x14, x15;
     uint32_t j0, j1, j2, j3, j4, j5, j6, j7, j8, j9, j10, j11, j12, j13, j14, j15;
-    unsigned char *ctarget = nullptr;
-    unsigned char tmp[64];
-    unsigned int i;
 
-    if (!bytes) return;
+    if (!blocks) return;
 
     j0 = input[0];
     j1 = input[1];
@@ -212,14 +200,6 @@ void ChaCha20::Crypt(const unsigned char* m, unsigned char* c, size_t bytes)
     j15 = input[15];
 
     for (;;) {
-        if (bytes < 64) {
-            // if m has fewer than 64 bytes available, copy m to tmp and
-            // read from tmp instead
-            for (i = 0;i < bytes;++i) tmp[i] = m[i];
-            m = tmp;
-            ctarget = c;
-            c = tmp;
-        }
         x0 = j0;
         x1 = j1;
         x2 = j2;
@@ -303,16 +283,48 @@ void ChaCha20::Crypt(const unsigned char* m, unsigned char* c, size_t bytes)
         WriteLE32(c + 56, x14);
         WriteLE32(c + 60, x15);
 
-        if (bytes <= 64) {
-            if (bytes < 64) {
-                for (i = 0;i < bytes;++i) ctarget[i] = c[i];
-            }
+        if (blocks == 1) {
             input[12] = j12;
             input[13] = j13;
             return;
         }
-        bytes -= 64;
+        blocks -= 1;
         c += 64;
         m += 64;
     }
 }
+
+void ChaCha20::Keystream(unsigned char* c, size_t bytes)
+{
+    if (!bytes) return;
+    if (bytes >= 64) {
+        size_t blocks = bytes / 64;
+        m_aligned.Keystream64(c, blocks);
+        c += blocks * 64;
+        bytes -= blocks * 64;
+    }
+    if (bytes) {
+        unsigned char buffer[64];
+        m_aligned.Keystream64(buffer, 1);
+        memcpy(c, buffer, bytes);
+    }
+}
+
+void ChaCha20::Crypt(const unsigned char* m, unsigned char* c, size_t bytes)
+{
+    if (!bytes) return;
+    if (bytes >= 64) {
+        size_t blocks = bytes / 64;
+        m_aligned.Crypt64(m, c, blocks);
+        c += blocks * 64;
+        m += blocks * 64;
+        bytes -= blocks * 64;
+    }
+    if (bytes) {
+        unsigned char buffer[64];
+        m_aligned.Keystream64(buffer, 1);
+        for (unsigned i = 0; i < bytes; i++) {
+            c[i] = m[i] ^ buffer[i];
+        }
+    }
+}
diff --git a/src/crypto/chacha20.h b/src/crypto/chacha20.h
index 624c083191..12ddef9256 100644
--- a/src/crypto/chacha20.h
+++ b/src/crypto/chacha20.h
@@ -8,19 +8,59 @@
 #include <cstdlib>
 #include <stdint.h>
 
-/** A class for ChaCha20 256-bit stream cipher developed by Daniel J. Bernstein
-    https://cr.yp.to/chacha/chacha-20080128.pdf */
-class ChaCha20
+// classes for ChaCha20 256-bit stream cipher developed by Daniel J. Bernstein
+// https://cr.yp.to/chacha/chacha-20080128.pdf */
+
+/** ChaCha20 cipher that only operates on multiples of 64 bytes. */
+class ChaCha20Aligned
 {
 private:
     uint32_t input[16];
 
 public:
-    ChaCha20();
-    ChaCha20(const unsigned char* key, size_t keylen);
-    void SetKey(const unsigned char* key, size_t keylen); //!< set key with flexible keylength; 256bit recommended */
-    void SetIV(uint64_t iv); // set the 64bit nonce
-    void Seek(uint64_t pos); // set the 64bit block counter
+    ChaCha20Aligned();
+
+    /** Initialize a cipher with specified key (see SetKey for arguments). */
+    ChaCha20Aligned(const unsigned char* key, size_t keylen);
+
+    /** set key with flexible keylength (16 or 32 bytes; 32 recommended). */
+    void SetKey(const unsigned char* key, size_t keylen);
+
+    /** set the 64-bit nonce. */
+    void SetIV(uint64_t iv);
+
+    /** set the 64bit block counter (pos seeks to byte position 64*pos). */
+    void Seek(uint64_t pos);
+
+    /** outputs the keystream of size <64*blocks> into <c> */
+    void Keystream64(unsigned char* c, size_t blocks);
+
+    /** enciphers the message <input> of length <64*blocks> and write the enciphered representation into <output>
+     *  Used for encryption and decryption (XOR)
+     */
+    void Crypt64(const unsigned char* input, unsigned char* output, size_t blocks);
+};
+
+/** Unrestricted ChaCha20 cipher. Seeks forward to a multiple of 64 bytes after every operation. */
+class ChaCha20
+{
+private:
+    ChaCha20Aligned m_aligned;
+
+public:
+    ChaCha20() = default;
+
+    /** Initialize a cipher with specified key (see SetKey for arguments). */
+    ChaCha20(const unsigned char* key, size_t keylen) : m_aligned(key, keylen) {}
+
+    /** set key with flexible keylength (16 or 32 bytes; 32 recommended). */
+    void SetKey(const unsigned char* key, size_t keylen) { m_aligned.SetKey(key, keylen); }
+
+    /** set the 64-bit nonce. */
+    void SetIV(uint64_t iv) { m_aligned.SetIV(iv); }
+
+    /** set the 64bit block counter (pos seeks to byte position 64*pos). */
+    void Seek(uint64_t pos) { m_aligned.Seek(pos); }
 
     /** outputs the keystream of size <bytes> into <c> */
     void Keystream(unsigned char* c, size_t bytes);