From 379b79ed6e2cdcc77e54fdaa39be95497dd782a9 Mon Sep 17 00:00:00 2001
From: Loup Vaillant <loup@loup-vaillant.fr>
Date: Mon, 17 Jul 2017 00:52:25 +0200
Subject: [PATCH] added crypto_chacha20_set_ctr()

---
 src/monocypher.c   | 14 +++++++++-----
 src/monocypher.h   |  2 ++
 tests/properties.c | 31 +++++++++++++++++++++++++++++++
 3 files changed, 42 insertions(+), 5 deletions(-)

diff --git a/src/monocypher.c b/src/monocypher.c
index 1c9da4a..68691a0 100644
--- a/src/monocypher.c
+++ b/src/monocypher.c
@@ -130,8 +130,6 @@ sv chacha20_init_key(crypto_chacha_ctx *ctx, const u8 key[32])
     FOR (i, 0, 8) {
         ctx->input[i+4] = load32_le(key + i*4);
     }
-    // pool index (the random pool starts empty)
-    ctx->pool_index = 64;
 }
 
 void crypto_chacha20_H(u8 out[32], const u8 key[32], const u8 in[16])
@@ -155,21 +153,27 @@ void crypto_chacha20_init(crypto_chacha_ctx *ctx,
                           const u8           nonce[8])
 {
     chacha20_init_key(ctx, key  );         // key
-    ctx->input[12] = 0;                    // counter
-    ctx->input[13] = 0;                    // counter
+    crypto_chacha20_set_ctr(ctx, 0);       // counter
     ctx->input[14] = load32_le(nonce + 0); // nonce
     ctx->input[15] = load32_le(nonce + 4); // nonce
 }
 
 void crypto_chacha20_x_init(crypto_chacha_ctx *ctx,
                             const u8           key[32],
-                             const u8           nonce[24])
+                            const u8           nonce[24])
 {
     u8 derived_key[32];
     crypto_chacha20_H(derived_key, key, nonce);
     crypto_chacha20_init(ctx, derived_key, nonce + 16);
 }
 
+void crypto_chacha20_set_ctr(crypto_chacha_ctx *ctx, u64 ctr)
+{
+    ctx->input[12]  = ctr & 0xffffffff;
+    ctx->input[13]  = ctr >> 32;
+    ctx->pool_index = 64;  // The random pool (re)starts empty
+}
+
 void crypto_chacha20_encrypt(crypto_chacha_ctx *ctx,
                              u8                *cipher_text,
                              const u8          *plain_text,
diff --git a/src/monocypher.h b/src/monocypher.h
index 1aa513c..e58665a 100644
--- a/src/monocypher.h
+++ b/src/monocypher.h
@@ -33,6 +33,8 @@ void crypto_chacha20_x_init(crypto_chacha_ctx *ctx,
                             const uint8_t      key[32],
                             const uint8_t      nonce[24]);
 
+void crypto_chacha20_set_ctr(crypto_chacha_ctx *ctx, uint64_t set_ctr);
+
 void crypto_chacha20_encrypt(crypto_chacha_ctx *ctx,
                              uint8_t           *cipher_text,
                              const uint8_t     *plain_text,
diff --git a/tests/properties.c b/tests/properties.c
index 5db6fbe..d5ffdf5 100644
--- a/tests/properties.c
+++ b/tests/properties.c
@@ -83,6 +83,36 @@ static int chacha20()
     return status;
 }
 
+static int chacha20_set_ctr()
+{
+    static const size_t nb_blocks   = 10;
+    static const size_t block_size  = 64;
+    static const size_t stream_size = block_size * nb_blocks;
+    int status = 0;
+    FOR (i, 0, 1000) {
+        u8 output_part[stream_size];
+        u8 output_all [stream_size];
+        u8 key        [32];          p_random(key  , 32);
+        u8 nonce      [8];           p_random(nonce, 8 );
+        size_t ctr   = rand_mod(nb_blocks);
+        size_t limit = ctr * block_size;
+        // Encrypt all at once
+        crypto_chacha_ctx ctx;
+        crypto_chacha20_init(&ctx, key, nonce);
+        crypto_chacha20_stream(&ctx, output_all, stream_size);
+        // Encrypt second part
+        crypto_chacha20_set_ctr(&ctx, ctr);
+        crypto_chacha20_stream(&ctx, output_part + limit, stream_size - limit);
+        // Encrypt first part
+        crypto_chacha20_set_ctr(&ctx, 0);
+        crypto_chacha20_stream(&ctx, output_part, limit);
+        // Compare the results (must be the same)
+        status |= crypto_memcmp(output_part, output_all, stream_size);
+    }
+    printf("%s: Chacha20 (set counter)\n", status != 0 ? "FAILED" : "OK");
+    return status;
+}
+
 // Tests that authenticating bit by bit yields the same mac than
 // authenticating all at once
 static int poly1305()
@@ -200,6 +230,7 @@ int main(void)
 {
     int status = 0;
     status |= chacha20();
+    status |= chacha20_set_ctr();
     status |= poly1305();
     status |= blake2b();
     status |= aead();
-- 
2.47.3