Make unrestricted ChaCha20 cipher not waste keystream bytes

Co-authored-by: dhruv <856960+dhruv@users.noreply.github.com>

2 files changed, 36 insertions, 9 deletions

diff --git a/src/crypto/chacha20.cpp b/src/crypto/chacha20.cpp
index c72ccccc65..80db0a3a83 100644
--- a/src/crypto/chacha20.cpp
+++ b/src/crypto/chacha20.cpp
@@ -297,6 +297,13 @@ inline void ChaCha20Aligned::Crypt64(const unsigned char* m, unsigned char* c, s
 void ChaCha20::Keystream(unsigned char* c, size_t bytes)
 {
     if (!bytes) return;
+    if (m_bufleft) {
+        unsigned reuse = std::min<size_t>(m_bufleft, bytes);
+        memcpy(c, m_buffer + 64 - m_bufleft, reuse);
+        m_bufleft -= reuse;
+        bytes -= reuse;
+        c += reuse;
+    }
     if (bytes >= 64) {
         size_t blocks = bytes / 64;
         m_aligned.Keystream64(c, blocks);
@@ -304,15 +311,25 @@ void ChaCha20::Keystream(unsigned char* c, size_t bytes)
         bytes -= blocks * 64;
     }
     if (bytes) {
-        unsigned char buffer[64];
-        m_aligned.Keystream64(buffer, 1);
-        memcpy(c, buffer, bytes);
+        m_aligned.Keystream64(m_buffer, 1);
+        memcpy(c, m_buffer, bytes);
+        m_bufleft = 64 - bytes;
     }
 }
 
 void ChaCha20::Crypt(const unsigned char* m, unsigned char* c, size_t bytes)
 {
     if (!bytes) return;
+    if (m_bufleft) {
+        unsigned reuse = std::min<size_t>(m_bufleft, bytes);
+        for (unsigned i = 0; i < reuse; i++) {
+            c[i] = m[i] ^ m_buffer[64 - m_bufleft + i];
+        }
+        m_bufleft -= reuse;
+        bytes -= reuse;
+        c += reuse;
+        m += reuse;
+    }
     if (bytes >= 64) {
         size_t blocks = bytes / 64;
         m_aligned.Crypt64(m, c, blocks);
@@ -321,10 +338,10 @@ void ChaCha20::Crypt(const unsigned char* m, unsigned char* c, size_t bytes)
         bytes -= blocks * 64;
     }
     if (bytes) {
-        unsigned char buffer[64];
-        m_aligned.Keystream64(buffer, 1);
+        m_aligned.Keystream64(m_buffer, 1);
         for (unsigned i = 0; i < bytes; i++) {
-            c[i] = m[i] ^ buffer[i];
+            c[i] = m[i] ^ m_buffer[i];
         }
+        m_bufleft = 64 - bytes;
     }
 }
diff --git a/src/crypto/chacha20.h b/src/crypto/chacha20.h
index fdef257ed3..715bf4e8e9 100644
--- a/src/crypto/chacha20.h
+++ b/src/crypto/chacha20.h
@@ -41,11 +41,13 @@ public:
     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. */
+/** Unrestricted ChaCha20 cipher. */
 class ChaCha20
 {
 private:
     ChaCha20Aligned m_aligned;
+    unsigned char m_buffer[64] = {0};
+    unsigned m_bufleft{0};
 
 public:
     ChaCha20() = default;
@@ -54,13 +56,21 @@ public:
     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); }
+    void SetKey(const unsigned char* key, size_t keylen)
+    {
+        m_aligned.SetKey(key, keylen);
+        m_bufleft = 0;
+    }
 
     /** 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 Seek64(uint64_t pos) { m_aligned.Seek64(pos); }
+    void Seek64(uint64_t pos)
+    {
+        m_aligned.Seek64(pos);
+        m_bufleft = 0;
+    }
 
     /** outputs the keystream of size <bytes> into <c> */
     void Keystream(unsigned char* c, size_t bytes);