fix overflow for long messages

2 files changed, 26 insertions, 3 deletions

diff --git a/src/block.rs b/src/block.rs
index e63062e..abd7ff7 100644
--- a/src/block.rs
+++ b/src/block.rs
@@ -53,6 +53,17 @@ impl Block {
         }
         Block(block)
     }
+
+    /// Computes self * 2^exponent
+    ///
+    /// Ensures that there's no overflow in computing 2^exponent.
+    pub fn exp(&self, exponent: u32) -> Block {
+        match exponent {
+            _ if exponent < 32 => *self * (1 << exponent),
+            _ if exponent % 2 == 0 => self.exp(exponent / 2).exp(exponent / 2),
+            _ => (*self * 2).exp(exponent - 1),
+        }
+    }
 }
 
 impl From<[u8; 16]> for Block {
diff --git a/src/lib.rs b/src/lib.rs
index 8b4a5e6..06eb7cd 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -216,7 +216,7 @@ fn encrypt(key: &Key, nonce: &[u8], ad: &[&[u8]], tau: u32, message: &[u8]) -> V
         .chain(iter::repeat_n(0, tau as usize))
         .collect::<Vec<_>>();
     // We treat tau as bytes, but according to the spec, tau is actually in bits.
-    let tau_block = Block::from_int(tau * 8);
+    let tau_block = Block::from_int(tau as u128 * 8);
     let mut tweaks = vec![&tau_block.0, nonce];
     tweaks.extend(ad);
     if message.is_empty() {
@@ -647,10 +647,10 @@ fn e(j: i32, i: i32, key: &Key, block: Block) -> Block {
         aes10(&k, &(block ^ delta))
     } else {
         let k = [&Block::NULL, &key_j, &key_i, &key_l, &Block::NULL];
-        let exponent = if i % 8 == 0 { i / 8 } else { i / 8 + 1 };
         let j: u32 = j.try_into().expect("j was negative");
         let i: u32 = i.try_into().expect("i was negative");
-        let delta = (key_j * j) ^ (key_i * (1 << exponent)) ^ (key_l * (i % 8));
+        let exponent = if i % 8 == 0 { i / 8 } else { i / 8 + 1 };
+        let delta = (key_j * j) ^ key_i.exp(exponent) ^ (key_l * (i % 8));
         aes4(&k, &(block ^ delta))
     }
 }
@@ -793,4 +793,16 @@ mod test {
         assert!(aez.decrypt(&[0], &[b"foobar"], 16, &hash).is_some());
         assert!(aez.decrypt(&[0], &[b"boofar"], 16, &hash).is_none());
     }
+
+    #[test]
+    fn test_fuzzed_1() {
+        let aez = Aez::new(b"");
+        aez.encrypt(b"", &[], 2220241, &[0]);
+    }
+
+    #[test]
+    fn test_fuzzed_2() {
+        let aez = Aez::new(b"");
+        aez.encrypt(b"", &[], 673261693, &[]);
+    }
 }