diff --git a/Cargo.toml b/Cargo.toml index 4ad50ca..6062f7e 100644 --- a/Cargo.toml +++ b/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" diff --git a/aes/src/aes.rs b/aes/src/aes.rs index c76e586..6be43e6 100644 --- a/aes/src/aes.rs +++ b/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); + } } diff --git a/aes/src/block/block128.rs b/aes/src/block/block128.rs index be6199e..a65bebc 100644 --- a/aes/src/block/block128.rs +++ b/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; [ diff --git a/aes/src/key/mod.rs b/aes/src/key/mod.rs index 5df3051..737c843 100644 --- a/aes/src/key/mod.rs +++ b/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}, +}; diff --git a/aes/src/key/subkeys.rs b/aes/src/key/subkeys.rs index 8a42103..f20f563 100644 --- a/aes/src/key/subkeys.rs +++ b/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.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(val: T) -> T { + val.sbox_lookup() } #[cfg(test)] diff --git a/aes/src/lib.rs b/aes/src/lib.rs index c390a0e..ed31ae7 100644 --- a/aes/src/lib.rs +++ b/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}; diff --git a/aes/src/sbox.rs b/aes/src/sbox.rs new file mode 100644 index 0000000..7a51f0d --- /dev/null +++ b/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);