1 files changed, 9 insertions, 4 deletions
diff --git a/src/lib.rs b/src/lib.rs
index 8985c2e..49a83d6 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -626,13 +626,18 @@ fn aez_hash(aez: &Aez, tweaks: Tweak) -> Block {
 fn aez_prf(aez: &Aez, tweaks: Tweak, buffer: &mut [u8]) {
     let mut index = 0u128;
     let delta = aez_hash(aez, tweaks);
-    for chunk in buffer.chunks_mut(16) {
+    for chunk in buffer.chunks_exact_mut(16) {
+        let chunk: &mut [u8; 16] = chunk.try_into().unwrap();
         let block = e(-1, 3, aez, delta ^ Block::from_int(index));
-        for (a, b) in chunk.iter_mut().zip(block.bytes().iter()) {
-            *a ^= b;
-        }
+        (block ^ Block::from(*chunk)).write_to(chunk);
         index += 1;
     }
+    let suffix_start = buffer.len() - buffer.len() % 16;
+    let chunk = &mut buffer[suffix_start..];
+    let block = e(-1, 3, aez, delta ^ Block::from_int(index));
+    for (a, b) in chunk.iter_mut().zip(block.bytes().iter()) {
+        *a ^= *b;
+    }
 }
 
 /// Represents a computation of E_K^{j,i}.