Make unrestricted ChaCha20 cipher not waste keystream bytes

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

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;
     }
 }

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);

src/test/crypto_tests.cpp

@@ -500,6 +500,24 @@ BOOST_AUTO_TEST_CASE(chacha20_testvector)
                  "fab78c9");
 }
 
+BOOST_AUTO_TEST_CASE(chacha20_midblock)
+{
+    auto key = ParseHex("0000000000000000000000000000000000000000000000000000000000000000");
+    ChaCha20 c20{key.data(), 32};
+    // get one block of keystream
+    unsigned char block[64];
+    c20.Keystream(block, CHACHA20_ROUND_OUTPUT);
+    unsigned char b1[5], b2[7], b3[52];
+    c20 = ChaCha20{key.data(), 32};
+    c20.Keystream(b1, 5);
+    c20.Keystream(b2, 7);
+    c20.Keystream(b3, 52);
+
+    BOOST_CHECK_EQUAL(0, memcmp(b1, block, 5));
+    BOOST_CHECK_EQUAL(0, memcmp(b2, block + 5, 7));
+    BOOST_CHECK_EQUAL(0, memcmp(b3, block + 12, 52));
+}
+
 BOOST_AUTO_TEST_CASE(poly1305_testvector)
 {
     // RFC 7539, section 2.5.2.

src/test/fuzz/crypto_diff_fuzz_chacha20.cpp

@@ -310,20 +310,26 @@ FUZZ_TARGET(crypto_diff_fuzz_chacha20)
             },
             [&] {
                 uint32_t integralInRange = fuzzed_data_provider.ConsumeIntegralInRange<size_t>(0, 4096);
+                // DJB's version seeks forward to a multiple of 64 bytes after every operation. Correct for that.
+                uint64_t pos = ctx.input[12] + (((uint64_t)ctx.input[13]) << 32) + ((integralInRange + 63) >> 6);
                 std::vector<uint8_t> output(integralInRange);
                 chacha20.Keystream(output.data(), output.size());
                 std::vector<uint8_t> djb_output(integralInRange);
                 ECRYPT_keystream_bytes(&ctx, djb_output.data(), djb_output.size());
                 assert(output == djb_output);
+                chacha20.Seek64(pos);
             },
             [&] {
                 uint32_t integralInRange = fuzzed_data_provider.ConsumeIntegralInRange<size_t>(0, 4096);
+                // DJB's version seeks forward to a multiple of 64 bytes after every operation. Correct for that.
+                uint64_t pos = ctx.input[12] + (((uint64_t)ctx.input[13]) << 32) + ((integralInRange + 63) >> 6);
                 std::vector<uint8_t> output(integralInRange);
                 const std::vector<uint8_t> input = ConsumeFixedLengthByteVector(fuzzed_data_provider, output.size());
                 chacha20.Crypt(input.data(), output.data(), input.size());
                 std::vector<uint8_t> djb_output(integralInRange);
                 ECRYPT_encrypt_bytes(&ctx, input.data(), djb_output.data(), input.size());
                 assert(output == djb_output);
+                chacha20.Seek64(pos);
             });
     }
 }