Cleaner Rust block encryption code

utils/src/aes.rs

@@ -136,17 +136,15 @@ fn encrypt_block(block: &mut Block, round_keys: &[Block]) -> Result<()> {
     }
 
     add_round_key(block, &round_keys[0]);
-    for round in 1..ROUNDS {
+    for round in 1..=ROUNDS {
         substitute_block(block);
         shift_rows_left(block);
-        mix_words(block);
+        if round != ROUNDS {
+            mix_words(block);
+        }
         add_round_key(block, &round_keys[round]);
     }
 
-    substitute_block(block);
-    shift_rows_left(block);
-    add_round_key(block, &round_keys[round_keys.len() - 1]);
-
     Ok(())
 }
 
@@ -156,13 +154,11 @@ fn decrypt_block(block: &mut Block, round_keys: &[Block]) -> Result<()> {
         return Err(AesError::InvalidRoundKeys);
     }
 
-    add_round_key(block, &round_keys[round_keys.len() - 1]);
-    shift_rows_right(block);
-    inverse_substitute_block(block);
-
-    for round in (1..ROUNDS).rev() {
+    for round in (1..=ROUNDS).rev() {
         add_round_key(block, &round_keys[round]);
-        inverse_mix_words(block);
+        if round != ROUNDS {
+            inverse_mix_words(block);
+        }
         shift_rows_right(block);
         inverse_substitute_block(block);
     }
@@ -562,9 +558,7 @@ mod test {
     #[test]
     fn test_encrypt_block() {
         let round_keys = expand_key(&TEST_KEY);
-
-        let mut input = empty_block();
-        input[2][2] = 10;
+        let input = TEST_BLOCK.clone();
 
         let mut buf = input.clone();
         let _ = encrypt_block(&mut buf, &round_keys);