1 files changed, 22 insertions, 15 deletions

diff --git a/src/lib.rs b/src/lib.rs
index fe7a985..210c7f8 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -315,7 +315,9 @@ fn encrypt(key: &Key, nonce: &[u8], ad: &[&[u8]], tau: u32, buffer: &mut [u8]) {
     tweaks.extend(ad);
     assert!(buffer.len() >= tau as usize);
     if buffer.len() == tau as usize {
-        buffer.copy_from_slice(&aez_prf(key, &tweaks, tau));
+        // As aez_prf only xor's the input in, we have to clear the buffer first
+        buffer.fill(0);
+        aez_prf(key, &tweaks, buffer);
     } else {
         encipher(key, &tweaks, buffer);
     }
@@ -337,7 +339,8 @@ fn decrypt<'a>(
     tweaks.extend(ad);
 
     if ciphertext.len() == tau as usize {
-        if constant_time_eq(&ciphertext, &aez_prf(key, &tweaks, tau)) {
+        aez_prf(key, &tweaks, ciphertext);
+        if is_zeroes(&ciphertext) {
             return Some(&[]);
         } else {
             return None;
@@ -347,16 +350,19 @@ fn decrypt<'a>(
     decipher(key, &tweaks, ciphertext);
     let (m, auth) = ciphertext.split_at(ciphertext.len() - tau as usize);
     assert!(auth.len() == tau as usize);
-    let comparator = if tau as usize <= ZEROES.len() {
-        &ZEROES[..tau as usize]
+
+    if is_zeroes(&auth) { Some(m) } else { None }
+}
+
+fn is_zeroes(data: &[u8]) -> bool {
+    let comparator = if data.len() <= ZEROES.len() {
+        &ZEROES[..data.len()]
     } else {
-        &vec![0; tau as usize]
+        // We should find a way to do this without allocating a separate buffer full of zeroes, but
+        // I don't want to hand-roll my constant-time-is-zeroes yet.
+        &vec![0; data.len()]
     };
-    if constant_time_eq(&auth, comparator) {
-        Some(m)
-    } else {
-        None
-    }
+    constant_time_eq(data, comparator)
 }
 
 fn encipher(key: &Key, tweaks: Tweak, message: &mut [u8]) {
@@ -702,16 +708,17 @@ fn aez_hash(key: &Key, tweaks: Tweak) -> Block {
     hash
 }
 
-fn aez_prf(key: &Key, tweaks: Tweak, tau: u32) -> Vec<u8> {
-    let mut result = Vec::new();
+/// XOR's the result of aez_prf into the given buffer
+fn aez_prf(key: &Key, tweaks: Tweak, buffer: &mut [u8]) {
     let mut index = 0u128;
     let delta = aez_hash(key, tweaks);
-    while result.len() < tau as usize {
+    for chunk in buffer.chunks_mut(16) {
         let block = e(-1, 3, key, delta ^ Block::from_int(index));
-        result.extend_from_slice(&block.0[..16.min(tau as usize - result.len())]);
+        for (a, b) in chunk.iter_mut().zip(block.0.iter()) {
+            *a ^= b;
+        }
         index += 1;
     }
-    result
 }
 
 /// Represents a computation of E_K^{j,i}.