feat(aes): add substitute bytes

Cargo.toml

@@ -23,4 +23,4 @@ zeroize = { version = "1.8", features = ["derive"] }
 [workspace.lints.clippy]
 pedantic = "warn"
 nursery = "warn"
-unwrap_used = "warn"
+# unwrap_used = "warn"

aes/src/aes.rs

@@ -2,6 +2,7 @@ use crate::{
     Block128,
     block::Block32,
     key::{Key, Subkey, Subkeys},
+    sbox::SboxLookup,
 };
 use cipher_core::{BlockCipher, CipherError};
 
@@ -34,18 +35,28 @@ impl BlockCipher for Aes {
 
         let block128 = Block128::from_be_bytes(block_arr);
 
-        let round_key = add_round_key(
-            *self.subkeys.first().unwrap(),
-            *block128.as_block32_array().first().unwrap(),
-        );
+        let mut subkey_iter = self.subkeys.chunks();
+        dbg!(&subkey_iter.count());
+
+        // let foo = *subkey_iter.next().unwrap();
+        // let round_key = add_round_key(
+        //     *block128.as_block32_array().first().unwrap(),
+        //     *subkey_iter.next().unwrap(),
+        // );
+
+        // for i in subkey_iter {}
         todo!()
     }
 }
 
-fn add_round_key(subkey: Subkey, block: Block32) -> Block32 {
+fn add_round_key(block: Block32, subkey: Subkey) -> Block32 {
     block ^ subkey
 }
 
+fn substitute_bytes(block: Block128) -> Block128 {
+    block.sbox_lookup()
+}
+
 #[cfg(test)]
 mod tests {
     use rstest::rstest;
@@ -63,8 +74,20 @@ mod tests {
         let block = Block32::new(block);
         let subkey = Subkey::from_u32(subkey);
 
-        let result = add_round_key(subkey, block);
+        let result = add_round_key(block, subkey);
 
         assert_eq!(result.as_u32(), expected);
     }
+
+    #[rstest]
+    #[case(
+        0x0E36_34AE_CE72_25B6_E26B_174E_D92B_5588,
+        0xAB05_18E4_8B40_3F4E_987F_F02F_35F1_FCC4
+    )]
+    fn byte_substitution(#[case] block: u128, #[case] expected: u128) {
+        let block = Block128::new(block);
+
+        let result = substitute_bytes(block);
+        assert_eq!(result.as_u128(), expected);
+    }
 }

aes/src/block/block128.rs

@@ -1,7 +1,10 @@
-use crate::block::{Block32, secret_block};
+use crate::{
+    block::{Block32, secret_block},
+    sbox::SboxLookup,
+};
 use cipher_core::{BlockError, InputBlock};
 use std::{
-    slice::{from_raw_parts, from_raw_parts_mut},
+    slice::{ChunksExact, from_raw_parts, from_raw_parts_mut},
     str::FromStr,
 };
 
@@ -19,6 +22,12 @@ impl InputBlock for Block128 {
     }
 }
 
+impl SboxLookup for Block128 {
+    fn sbox_lookup(self) -> Self {
+        Self(self.0.sbox_lookup())
+    }
+}
+
 impl Block128 {
     #[inline]
     #[must_use]
@@ -40,6 +49,7 @@ impl Block128 {
 
     #[inline]
     #[must_use]
+    #[allow(clippy::cast_possible_truncation)]
     pub const fn as_block32_array(self) -> [Block32; 4] {
         let val = self.0;
         [

aes/src/key/mod.rs

@@ -5,4 +5,8 @@ mod subkey;
 mod subkeys;
 
 use crate::secret_key;
-pub use {aes_key::Key, subkey::Subkey, subkeys::Subkeys};
+pub use {
+    aes_key::Key,
+    subkey::Subkey,
+    subkeys::{SubkeyChunks, Subkeys},
+};

aes/src/key/subkeys.rs

@@ -1,12 +1,13 @@
 use crate::{
-    constants::{RCON, S_BOXES},
+    constants::RCON,
     key::{Key, expanded::ExpandedKey, subkey::Subkey},
+    sbox::SboxLookup,
 };
 use std::{
     fmt::Debug,
     iter::Rev,
     ops::Index,
-    slice::{Iter, IterMut},
+    slice::{ChunksExact, Iter, IterMut},
 };
 
 const SUBKEY_COUNT: usize = 44;
@@ -53,9 +54,15 @@ impl Subkeys {
     }
 
     /// Returns the first element of the slice, or `None` if it is empty.
-    pub fn first(&self) -> Option<&Subkey> {
+    pub const fn first(&self) -> Option<&Subkey> {
         self.0.first()
     }
+
+    #[inline]
+    #[must_use]
+    pub fn chunks(&self) -> SubkeyChunks<'_> {
+        SubkeyChunks(self.0.chunks_exact(4))
+    }
 }
 
 impl<'a> IntoIterator for &'a Subkeys {
@@ -87,7 +94,18 @@ impl Debug for Subkeys {
     }
 }
 
-const fn expand(subkey: Subkey, rcon: u32) -> ExpandedKey {
+pub struct SubkeyChunks<'a>(ChunksExact<'a, Subkey>);
+
+impl<'a> Iterator for SubkeyChunks<'a> {
+    type Item = &'a [Subkey; 4];
+    fn next(&mut self) -> Option<Self::Item> {
+        self.0
+            .next()
+            .map(|chunk| <&[Subkey; 4]>::try_from(chunk).unwrap())
+    }
+}
+
+fn expand(subkey: Subkey, rcon: u32) -> ExpandedKey {
     let word = subkey.rotate_left(8).as_u32();
 
     let b0 = sbox_lookup(word >> 24);
@@ -98,12 +116,8 @@ const fn expand(subkey: Subkey, rcon: u32) -> ExpandedKey {
     ExpandedKey::from_u32(substituted ^ rcon)
 }
 
-const fn sbox_lookup(byte: u32) -> u32 {
-    const MASK: u32 = 0xFF;
-    let row = ((byte & MASK) as usize) >> 4;
-    let col = ((byte & MASK) as usize) & 0xF;
-
-    S_BOXES[row][col] as u32
+fn sbox_lookup<T: SboxLookup>(val: T) -> T {
+    val.sbox_lookup()
 }
 
 #[cfg(test)]

aes/src/lib.rs

@@ -2,6 +2,7 @@ mod aes;
 mod block;
 mod constants;
 mod key;
+mod sbox;
 mod utils;
 
 pub use {aes::Aes, block::Block128};

aes/src/sbox.rs

@@ -0,0 +1,27 @@
+use crate::constants::S_BOXES;
+
+pub trait SboxLookup: Sized {
+    fn sbox_lookup(self) -> Self;
+}
+
+macro_rules! impl_sbox_lookup {
+    ($ty:ty, $bytes:expr) => {
+        impl SboxLookup for $ty {
+            fn sbox_lookup(self) -> Self {
+                (0..$bytes).fold(0, |acc, idx| {
+                    let shift = ($bytes - 1 - idx) * 8;
+                    let byte = ((self >> shift) & 0xFF) as u8;
+                    let row = (byte >> 4) as usize;
+                    let col = (byte & 0xF) as usize;
+                    acc | Self::from(S_BOXES[row][col]) << shift
+                })
+            }
+        }
+    };
+}
+
+impl_sbox_lookup!(u8, 1);
+impl_sbox_lookup!(u16, 2);
+impl_sbox_lookup!(u32, 4);
+impl_sbox_lookup!(u64, 8);
+impl_sbox_lookup!(u128, 16);