Hunting a bug that causes test failure

pyproject.toml

@@ -30,22 +30,10 @@ keywords = [
 [project.urls]
 
 [project.optional-dependencies]
-dev = ["pytest", "ruff"]
+dev = ["pytest", "ruff", "cryptography"]
 
 [tool.hatchling]
 
-[tool.hatch.version]
-source = "hatchling.version:Version"
-
-[tool.hatch.build]
-hooks.vcs.version-file = "randquik/_version.py"
-hooks.vcs.template = """
-# This file is automatically generated by hatch build.
-__version__ = {version!r}
-"""
-
-targets.sdist.include = ["/randquik"]
-
 [tool.ruff]
 extend-select = ["I", "W", "UP", "C4", "ISC", "S"]
 # Worth selecting but still too broken: ASYNC, B, DTZ, FA
@@ -70,3 +58,6 @@ ignore = [
 ]
 show-source = true
 show-fixes = true
+
+[tool.pytest.ini_options]
+pythonpath = ["."]

randquik/cha.py

@@ -11,7 +11,11 @@ if not src.is_dir():
 ffi = cffi.FFI()
 ffi.cdef(
     """
-    typedef struct cha_ctx { uint32_t input[16]; } cha_ctx;
+    typedef struct cha_ctx {
+        uint32_t input[16];
+        uint8_t unconsumed[64];
+        uint8_t uncount;
+    } cha_ctx;
 
     int cha_generate(uint8_t* out, uint64_t outlen, const uint8_t key[32], const uint8_t iv[16]);
 
@@ -20,7 +24,9 @@ ffi.cdef(
     int cha_update(cha_ctx* ctx, uint8_t* out, uint64_t outlen);
     """
 )
-lib = ffi.dlopen("../build/librandquik-chacha20.so")
+lib = ffi.dlopen(
+    (Path(__file__).parent.parent / "build/librandquik-chacha20.so").as_posix()
+)
 
 
 def _processKeys(key, iv):
@@ -58,14 +64,14 @@ class Cha:
 
     def __call__(self, out: bytearray | Any):
         """Fill the parameter with random bytes"""
-        out, outlen = _processBuffer(out)
-        lib.cha_update(self.ctx, out, outlen)
+        outbuf, outlen = _processBuffer(out)
+        lib.cha_update(self.ctx, outbuf, outlen)
         return out
 
 
 def generate(out: bytearray | Any, key: bytes | Any, iv: bytes | Any):
     """Setup a generator, fill the out buffer and dispose the generator"""
-    key, iv =_processKeys(key, iv)
-    out, outlen = _processBuffer(out)
-    lib.cha_generate(out, outlen, key, iv)
+    key, iv = _processKeys(key, iv)
+    outbuf, outlen = _processBuffer(out)
+    lib.cha_generate(outbuf, outlen, key, iv)
     return out

src/cha1block.h

@@ -41,6 +41,9 @@ static inline uint64_t _cha_block(cha_ctx* ctx, uint8_t* begin, uint8_t* end) {
         if (bytes < 64) {
             memcpy(c, x, bytes);
             c = end;
+            // Leftover bytes are stored in ctx for next call
+            ctx->uncount = 64 - bytes;
+            memcpy(ctx->unconsumed, x + bytes, ctx->uncount);
             break;
         }
         memcpy(c, x, 64);

src/chacha20.c

@@ -30,6 +30,8 @@ void cha_init(cha_ctx* ctx, const uint8_t* key, const uint8_t* iv) {
     ctx->state[3] = 0x6b206574;
     memcpy(ctx->state + 4, key, 32);
     memcpy(ctx->state + 12, iv, 16);
+    memset(ctx->unconsumed, 0, sizeof ctx->unconsumed);
+    ctx->uncount = 0;
 }
 
 void cha_wipe(cha_ctx* ctx) { memset(ctx, 0, sizeof(cha_ctx)); }
@@ -38,6 +40,19 @@ int cha_update(cha_ctx* ctx, uint8_t* out, uint64_t outlen) {
     // The included header will mess with these variables
     uint8_t* c = out;
     uint8_t* end = out + outlen;
+    if (ctx->uncount) {
+        // Deliver stored bytes first
+        uint64_t N = ctx->uncount >= outlen ? outlen : ctx->uncount;
+        fprintf(stderr, "%lu, %i, %lu\n", N, ctx->uncount, outlen);
+        memcpy(c, ctx->unconsumed, N);
+        ctx->uncount -= N;
+        c += N;
+        if (ctx->uncount) {
+            memmove(ctx->unconsumed, ctx->unconsumed + N, ctx->uncount);
+        }
+        if (c == out + outlen)
+            return 0;
+    }
     // TODO: Handle resume if we are not at block boundary
     if (__builtin_cpu_supports("ssse3")) {
         if (__builtin_cpu_supports("avx2")) {

src/chacha20.h

@@ -6,6 +6,8 @@ static const uint64_t CHA_BLOCK_SIZE = 64;
 
 typedef struct cha_ctx {
     uint32_t state[16];
+    uint8_t unconsumed[64];
+    uint8_t uncount;
 } cha_ctx;
 
 /// @brief Initialize cha_ctx

tests/test_chacha.py

@@ -0,0 +1,34 @@
+from secrets import randbelow, token_bytes
+
+from cryptography.hazmat.primitives.ciphers import Cipher
+from cryptography.hazmat.primitives.ciphers.algorithms import ChaCha20
+
+from randquik import cha
+
+
+def test_cipherstreams_fullblocks():
+    key = token_bytes(32)
+    nonce = token_bytes(16)
+    c0 = Cipher(ChaCha20(key, nonce), None, None).encryptor()
+    c1 = cha.Cha(key, nonce)
+
+    for i in range(2048):
+        N = 64 * (1 + randbelow(2048))
+        ct0 = c0.update(bytes(N))
+        ct1 = c1(bytearray(N))
+        assert len(ct0) == len(ct1)
+        assert ct0.hex() == ct1.hex()
+
+
+def test_cipherstreams_partial_updates():
+    key = token_bytes(32)
+    nonce = token_bytes(16)
+    c0 = Cipher(ChaCha20(key, nonce), None, None).encryptor()
+    c1 = cha.Cha(key, nonce)
+
+    for i in range(2048):
+        N = 1 + randbelow(2048)
+        ct0 = c0.update(bytes(N))
+        ct1 = c1(bytearray(N))
+        assert len(ct0) == len(ct1)
+        assert ct0.hex() == ct1.hex()