From abc9540b56337cdcd9dce9399b66c4807ff2e8bf Mon Sep 17 00:00:00 2001
From: myl7 <myl@myl.moe>
Date: Thu, 11 Apr 2024 01:09:11 +0800
Subject: [PATCH] Fix dcf prg

Fix a misuse of zip.
Add a test for the mistake.
---
 src/dcf/prg.rs | 20 ++++++++++++--------
 1 file changed, 12 insertions(+), 8 deletions(-)

diff --git a/src/dcf/prg.rs b/src/dcf/prg.rs
index 2fd118c..d005f7e 100644
--- a/src/dcf/prg.rs
+++ b/src/dcf/prg.rs
@@ -44,14 +44,16 @@ impl<const LAMBDA: usize, const N: usize> Prg<LAMBDA> for Aes256HirosePrg<LAMBDA
         let mut result_buf0 = [[0; LAMBDA]; 2];
         let mut result_buf1 = [[0; LAMBDA]; 2];
         let mut out_blocks = [GenericArray::default(); 2];
-        (0..2usize).zip(0..LAMBDA / 16).for_each(|(i, j)| {
-            let in_block0 = GenericArray::from_slice(&seed[j * 16..(j + 1) * 16]);
-            let in_block1 = GenericArray::from_slice(&seed_p[j * 16..(j + 1) * 16]);
-            self.ciphers[i * 16 + j]
-                .encrypt_blocks_b2b(&[*in_block0, *in_block1], &mut out_blocks)
-                .unwrap();
-            result_buf0[i][j * 16..(j + 1) * 16].copy_from_slice(out_blocks[0].as_ref());
-            result_buf1[i][j * 16..(j + 1) * 16].copy_from_slice(out_blocks[1].as_ref());
+        (0..2usize).for_each(|i| {
+            (0..LAMBDA / 16).for_each(|j| {
+                let in_block0 = GenericArray::from_slice(&seed[j * 16..(j + 1) * 16]);
+                let in_block1 = GenericArray::from_slice(&seed_p[j * 16..(j + 1) * 16]);
+                self.ciphers[i * (LAMBDA / 16) + j]
+                    .encrypt_blocks_b2b(&[*in_block0, *in_block1], &mut out_blocks)
+                    .unwrap();
+                result_buf0[i][j * 16..(j + 1) * 16].copy_from_slice(out_blocks[0].as_ref());
+                result_buf1[i][j * 16..(j + 1) * 16].copy_from_slice(out_blocks[1].as_ref());
+            });
         });
         result_buf0
             .iter_mut()
@@ -144,6 +146,8 @@ mod tests {
             assert_ne!(out[i].1, [0; 16]);
             assert_ne!(xor(&[&out[i].0, SEED]), [0; 16]);
             assert_ne!(xor(&[&out[i].1, SEED]), [0; 16]);
+            assert_ne!(xor(&[&out[i].0, SEED]), [0xff; 16]);
+            assert_ne!(xor(&[&out[i].1, SEED]), [0xff; 16]);
         });
     }
 }