speed up zero appendage

1 files changed, 17 insertions, 7 deletions

diff --git a/src/lib.rs b/src/lib.rs
index 06eb7cd..23ca0ff 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -82,7 +82,6 @@
 //! let plaintext = aez.decrypt(b"nonce", &[b"bar"], 16, &cipher);
 //! assert!(plaintext.is_none());
 //! ```
-use std::iter;
 
 use constant_time_eq::constant_time_eq;
 
@@ -94,6 +93,8 @@ use block::Block;
 type Key = [u8; 48];
 type Tweak<'a> = &'a [&'a [u8]];
 
+static ZEROES: [u8; 1024] = [0; 1024];
+
 /// AEZ encryption scheme.
 pub struct Aez {
     key: Key,
@@ -210,11 +211,15 @@ fn extract(key: &[u8]) -> [u8; 48] {
 }
 
 fn encrypt(key: &Key, nonce: &[u8], ad: &[&[u8]], tau: u32, message: &[u8]) -> Vec<u8> {
-    let auth_message = message
-        .iter()
-        .copied()
-        .chain(iter::repeat_n(0, tau as usize))
-        .collect::<Vec<_>>();
+    let mut auth_message = Vec::with_capacity(message.len() + tau as usize);
+    auth_message.extend_from_slice(&message);
+    while auth_message.len() < message.len() + tau as usize {
+        auth_message.extend_from_slice(
+            &ZEROES[..ZEROES
+                .len()
+                .min(tau as usize - (auth_message.len() - message.len()))],
+        );
+    }
     // We treat tau as bytes, but according to the spec, tau is actually in bits.
     let tau_block = Block::from_int(tau as u128 * 8);
     let mut tweaks = vec![&tau_block.0, nonce];
@@ -246,7 +251,12 @@ fn decrypt(key: &Key, nonce: &[u8], ad: &[&[u8]], tau: u32, ciphertext: &[u8]) -
     let x = decipher(key, &tweaks, ciphertext);
     let (m, auth) = x.split_at(ciphertext.len() - tau as usize);
     assert!(auth.len() == tau as usize);
-    if constant_time_eq(&auth, &vec![0; tau as usize]) {
+    let comparator = if tau as usize <= ZEROES.len() {
+        &ZEROES[..tau as usize]
+    } else {
+        &vec![0; tau as usize]
+    };
+    if constant_time_eq(&auth, comparator) {
         Some(Vec::from(m))
     } else {
         None