From 486f815aed00c7078daa62b83003f9fcdc0aaf24 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Milan=20=C5=A0pinka?= <mspi21@protonmail.com>
Date: Mon, 27 Jan 2025 03:41:11 +0100
Subject: [PATCH] Test ChaCha20 counter edge cases.

---
 src/primitive/streamcipher/chacha20.zig | 40 +++++++++++++++++++++++--
 1 file changed, 38 insertions(+), 2 deletions(-)

diff --git a/src/primitive/streamcipher/chacha20.zig b/src/primitive/streamcipher/chacha20.zig
index 4903637..ea9cc65 100644
--- a/src/primitive/streamcipher/chacha20.zig
+++ b/src/primitive/streamcipher/chacha20.zig
@@ -345,6 +345,42 @@ test "ChaCha20 Cipher" {
     try testing.expectEqualSlices(u8, reference[0..], buffer[0..]);
 }
 
-test "ChaCha20 counter increment edge cases" {
-    // TODO: test
+test "ChaCha20 32-bit counter increment edge cases" {
+    const key = [CHACHA20_KEY_SIZE]u8{
+        0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0a, 0x0b, 0x0c, 0x0d, 0x0e, 0x0f,
+        0x10, 0x11, 0x12, 0x13, 0x14, 0x15, 0x16, 0x17, 0x18, 0x19, 0x1a, 0x1b, 0x1c, 0x1d, 0x1e, 0x1f,
+    };
+    const nonce = [ChaCha20_RFC7539_Parameters.NONCE_SIZE]u8{
+        0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x4a, 0x00, 0x00, 0x00, 0x00,
+    };
+    const counter = word_to_bytes_le(std.math.maxInt(u32) - 1);
+
+    var chacha = chacha20_rfc7539_new(&key, &nonce, &counter);
+    defer chacha20_destroy(&chacha);
+
+    // Counter: 2^32 - 2 -> 2^32 - 1, OK
+    try chacha20_increment_counter(&chacha);
+
+    // Counter: 2^32 - 1 -> 2^32, overflow
+    try testing.expectError(KeyStreamDepleted, chacha20_increment_counter(&chacha));
+}
+
+test "ChaCha20 64-bit counter increment edge cases" {
+    const key = [CHACHA20_KEY_SIZE]u8{
+        0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0a, 0x0b, 0x0c, 0x0d, 0x0e, 0x0f,
+        0x10, 0x11, 0x12, 0x13, 0x14, 0x15, 0x16, 0x17, 0x18, 0x19, 0x1a, 0x1b, 0x1c, 0x1d, 0x1e, 0x1f,
+    };
+    const nonce = [ChaCha20_Bernstein_Parameters.NONCE_SIZE]u8{ 0x00, 0x00, 0x00, 0x4a, 0x00, 0x00, 0x00, 0x00 };
+    const counter = [ChaCha20_Bernstein_Parameters.COUNTER_SIZE]u8{ 0xff, 0xff, 0xff, 0xff, 0xef, 0xbe, 0xad, 0xde };
+
+    var chacha = chacha20_bernstein_new(&key, &nonce, &counter);
+    defer chacha20_destroy(&chacha);
+
+    // Counter: 0xdeadbeefffffffff -> 0xdeadbef000000000, OK
+    try chacha20_increment_counter(&chacha);
+    try testing.expectEqualSlices(u32, &.{ 0x00000000, 0xdeadbef0 }, chacha.nonce[0..2]);
+
+    // Counter: 0xffffffffffffffff -> 0x10000000000000000, overflow
+    @memset(chacha.nonce[0..2], std.math.maxInt(u32));
+    try testing.expectError(KeyStreamDepleted, chacha20_increment_counter(&chacha));
 }