1 files changed, 8 insertions, 39 deletions

diff --git a/src/lib.rs b/src/lib.rs
index b00612f..d1e0cab 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -86,11 +86,14 @@
 use constant_time_eq::constant_time_eq;
 
 mod accessor;
+mod aesround;
 mod block;
+
 #[cfg(test)]
 mod testvectors;
 
 use accessor::BlockAccessor;
+use aesround::AesRound;
 use block::Block;
 type Key = [u8; 48];
 type Tweak<'a> = &'a [&'a [u8]];
@@ -247,34 +250,6 @@ impl Aez {
     }
 }
 
-fn aesenc(block: &mut Block, key: &Block){
-    aes::hazmat::cipher_round((&mut block.0).into(), &key.0.into());
-}
-
-fn aes4(key_i: &Block, key_j: &Block, key_l: &Block, block: &Block) -> Block {
-    let mut r = *block;
-    aesenc(&mut r, key_j);
-    aesenc(&mut r, key_i);
-    aesenc(&mut r, key_l);
-    aesenc(&mut r, &Block::NULL);
-    r
-}
-
-fn aes10(key_i: &Block, key_j: &Block, key_l: &Block, block: &Block) -> Block {
-    let mut r = *block;
-    aesenc(&mut r, key_i);
-    aesenc(&mut r, key_j);
-    aesenc(&mut r, key_l);
-    aesenc(&mut r, key_i);
-    aesenc(&mut r, key_j);
-    aesenc(&mut r, key_l);
-    aesenc(&mut r, key_i);
-    aesenc(&mut r, key_j);
-    aesenc(&mut r, key_l);
-    aesenc(&mut r, key_i);
-    r
-}
-
 fn extract(key: &[u8]) -> [u8; 48] {
     if key.len() == 48 {
         key.try_into().unwrap()
@@ -717,12 +692,10 @@ fn aez_prf(key: &Key, tweaks: Tweak, buffer: &mut [u8]) {
 ///
 /// As we usually need multiple values with a fixed j and ascending i, this struct saves the
 /// temporary values and makes it much faster to compute E_K^{j, i+1}, E_K^{j, i+2}, ...
-#[derive(Clone, Debug)]
 struct E {
-    key_i: Block,
-    key_j: Block,
     key_l: Block,
     state: Estate,
+    aes: aesround::AesImpl,
 }
 
 #[derive(Clone, Debug)]
@@ -741,6 +714,7 @@ impl E {
     /// Create a new "suspended" computation of E_K^{j,i}.
     fn new(j: i32, i: u32, key: &Key) -> Self {
         let (key_i, key_j, key_l) = split_key(key);
+        let aes = aesround::AesImpl::new(key_i, key_j, key_l);
         let state = if j == -1 {
             Estate::Neg { i }
         } else {
@@ -752,12 +726,7 @@ impl E {
                 ki_p_i: key_i.exp(exponent),
             }
         };
-        E {
-            key_i,
-            key_j,
-            key_l,
-            state,
-        }
+        E { key_l, state, aes }
     }
 
     /// Complete this computation to evaluate E_K^{j,i}(block).
@@ -765,11 +734,11 @@ impl E {
         match self.state {
             Estate::Neg { i } => {
                 let delta = self.key_l * i;
-                aes10(&self.key_i, &self.key_j, &self.key_l, &(block ^ delta))
+                self.aes.aes10(block ^ delta)
             }
             Estate::Pos { i, kj_t_j, ki_p_i } => {
                 let delta = kj_t_j ^ ki_p_i ^ (self.key_l * (i % 8));
-                aes4(&self.key_i, &self.key_j, &self.key_l, &(block ^ delta))
+                self.aes.aes4(block ^ delta)
             }
         }
     }