From dedf194eeabbc352f4f673ea840fe7fed7f16133 Mon Sep 17 00:00:00 2001
From: Noah Stiltner <nstilt1@lsu.edu>
Date: Sat, 8 Jun 2024 08:15:46 -0500
Subject: [PATCH] adjusted legacy counter test to try different amounts of
 blocks

---
 chacha20/Cargo.toml   |  1 -
 chacha20/tests/mod.rs | 46 +++++++++++++++++++++++++------------------
 2 files changed, 27 insertions(+), 20 deletions(-)

diff --git a/chacha20/Cargo.toml b/chacha20/Cargo.toml
index ed8bc31c..1420ea5b 100644
--- a/chacha20/Cargo.toml
+++ b/chacha20/Cargo.toml
@@ -29,7 +29,6 @@ zeroize = { version = "1.8.1", optional = true }
 cpufeatures = "0.2"
 
 [dev-dependencies]
-chacha20_0_7 = { package = "chacha20", version = "0.7.0", features = ["legacy"] }
 cipher = { version = "=0.5.0-pre.4", features = ["dev"] }
 hex-literal = "0.4"
 rand_chacha = "0.3.1"
diff --git a/chacha20/tests/mod.rs b/chacha20/tests/mod.rs
index ae12b9fe..21bf8e6f 100644
--- a/chacha20/tests/mod.rs
+++ b/chacha20/tests/mod.rs
@@ -234,43 +234,36 @@ mod legacy {
         }
     }
 
-    /// Tests the 64-bit counter
-    #[test]
-    fn legacy_64_bit_counter() {
+    /// Tests the 64-bit counter with a given amount of test blocks
+    fn legacy_counter_over_u32_max<const N: usize>(test: &[u8; N]) {
+        assert!(N % 64 == 0, "N should be a multiple of 64");
         use cipher::StreamCipherSeekCore;
-        use chacha20_0_7::{ChaCha20Legacy as OgLegacy, LegacyNonce as OgLegacyNonce, cipher::{NewCipher, StreamCipher, StreamCipherSeek}};
+        // using rand_chacha v0.3 because it is already a dev-dependency, and
+        // it uses a 64-bit counter
         use rand_chacha::{ChaCha20Rng as OgRng, rand_core::{RngCore, SeedableRng}};
         let mut cipher = ChaCha20Legacy::new(&[0u8; 32].into(), &LegacyNonce::from([0u8; 8]));
-        let mut og_cipher = OgLegacy::new(&[0u8; 32].into(), &OgLegacyNonce::from([0u8; 8]));
-        //let mut rng = ChaCha20Rng
         let mut rng = OgRng::from_seed([0u8; 32]);
         
-        const TEST_BLOCKS: usize = 4;
-        const TEST: [u8; 64 * TEST_BLOCKS] = [0u8; 64 * TEST_BLOCKS];
-        let mut expected = TEST.clone();
+        let mut expected = test.clone();
         rng.fill_bytes(&mut expected);
-        //og_cipher.apply_keystream(&mut expected);
-        let mut result = TEST.clone();
+        let mut result = test.clone();
         cipher.apply_keystream(&mut result);
         assert_eq!(expected, result);
 
         const SEEK_POS: u64 = (u32::MAX - 10) as u64 * 64;
         cipher.seek(SEEK_POS);
         rng.set_word_pos(SEEK_POS as u128 / 4);
-        og_cipher.seek(SEEK_POS);
 
         let pos: u64 = cipher.current_pos();
-        //assert_eq!(pos, og_cipher.current_pos());
         assert_eq!(pos, rng.get_word_pos() as u64 * 4);
         let block_pos = cipher.get_core().get_block_pos();
         assert!(block_pos < u32::MAX as u64);
         // Apply keystream blocks until some point after the u32 boundary
-        for i in 1..80 {
+        for i in 1..16 {
             let starting_block_pos = cipher.get_core().get_block_pos() as i64 - u32::MAX as i64;
-            let mut expected = TEST.clone();
+            let mut expected = test.clone();
             rng.fill_bytes(&mut expected);
-            //og_cipher.apply_keystream(&mut expected);
-            let mut result = TEST.clone();
+            let mut result = test.clone();
             cipher.apply_keystream(&mut result);
             if expected != result {
                 let mut index: usize = 0;
@@ -286,7 +279,7 @@ mod legacy {
                 };
                 panic!("Index {} did not match;\n iteration: {}\n expected: {} != {}\nstart block pos - u32::MAX: {}", index, i, expected_u8, found_u8, starting_block_pos);
             }
-            let expected_block_pos = block_pos + i * TEST_BLOCKS as u64;
+            let expected_block_pos = block_pos + i * (test.len() / 64) as u64;
             assert!(expected_block_pos == cipher.get_core().get_block_pos(), 
                 "Block pos did not increment as expected; Expected block pos: {}\n actual block_pos: {}\n iteration: {}",
                 expected_block_pos,
@@ -296,6 +289,21 @@ mod legacy {
         }
         // this test assures us that the counter is in fact over u32::MAX, in
         // case we change some of the parameters
-        assert!(cipher.get_core().get_block_pos() > u32::MAX as u64);
+        assert!(cipher.get_core().get_block_pos() > u32::MAX as u64, "The 64-bit counter test did not surpass u32::MAX");
+    }
+
+    /// Runs the legacy_64_bit_counter test with different-sized arrays so that
+    /// both `gen_ks_block` and `gen_par_ks_blocks` are called with varying 
+    /// starting positions.
+    #[test]
+    fn legacy_64_bit_counter() {
+        legacy_counter_over_u32_max(&[0u8; 64 * 1]);
+        legacy_counter_over_u32_max(&[0u8; 64 * 2]);
+        legacy_counter_over_u32_max(&[0u8; 64 * 3]);
+        legacy_counter_over_u32_max(&[0u8; 64 * 4]);
+        legacy_counter_over_u32_max(&[0u8; 64 * 5]);
+        legacy_counter_over_u32_max(&[0u8; 64 * 6]);
+        legacy_counter_over_u32_max(&[0u8; 64 * 7]);
+        legacy_counter_over_u32_max(&[0u8; 64 * 8]);
     }
 }