diff --git a/aes/src/aes.rs b/aes/src/aes.rs index ad182cc..8c18a42 100644 --- a/aes/src/aes.rs +++ b/aes/src/aes.rs @@ -26,7 +26,7 @@ impl Aes { &self.subkeys } - fn encryot_block(&self, mut state: Block128) -> Block128 { + fn encrypt_block(&self, mut state: Block128) -> Block128 { let mut keys = self.subkeys.chunks(); state = add_round_key(state, keys.next().expect("Round key 0")); @@ -45,8 +45,8 @@ impl Aes { state } - fn decryot_block(&self, mut state: Block128) -> Block128 { - let mut keys = self.subkeys.chunks(); + fn decrypt_block(&self, mut state: Block128) -> Block128 { + let mut keys = self.subkeys.chunks_rev(); state = add_round_key(state, keys.next().expect("Final round key")); for _ in 1..10 { @@ -83,8 +83,8 @@ impl BlockCipher for Aes { let block128 = Block128::from_be_bytes(block_arr); let result = match action { - CipherAction::Encrypt => self.encryot_block(block128), - CipherAction::Decrypt => self.decryot_block(block128), + CipherAction::Encrypt => self.encrypt_block(block128), + CipherAction::Decrypt => self.decrypt_block(block128), }; Ok(result.into()) diff --git a/aes/src/block/block128.rs b/aes/src/block/block128.rs index d9535b1..814c416 100644 --- a/aes/src/block/block128.rs +++ b/aes/src/block/block128.rs @@ -4,6 +4,7 @@ use crate::{ }; use cipher_core::{BlockError, InputBlock}; use std::{ + ops::BitXor, slice::{from_raw_parts, from_raw_parts_mut}, str::FromStr, }; @@ -147,3 +148,10 @@ impl From<&Block128> for Vec { value.to_be_bytes().to_vec() } } + +impl BitXor for Block128 { + type Output = Self; + fn bitxor(self, rhs: u128) -> Self::Output { + Self(self.0 ^ rhs) + } +} diff --git a/aes/src/key/subkey.rs b/aes/src/key/subkey.rs index c1cd791..7cf97a2 100644 --- a/aes/src/key/subkey.rs +++ b/aes/src/key/subkey.rs @@ -1,5 +1,5 @@ use crate::key::secret_key; -use std::ops::BitXor; +use std::ops::{BitOr, BitXor, Shl}; secret_key! { /// A single AES round subkey @@ -47,3 +47,17 @@ impl BitXor for Subkey { Self(self.0 ^ rhs.as_u32()) } } + +impl Shl for Subkey { + type Output = u128; + fn shl(self, rhs: i32) -> Self::Output { + self.as_u128() << rhs + } +} + +impl BitOr for u128 { + type Output = Self; + fn bitor(self, rhs: Subkey) -> Self::Output { + self | rhs.as_u128() + } +} diff --git a/aes/src/operations/round_key.rs b/aes/src/operations/round_key.rs index ccbc6ca..bdd1fca 100644 --- a/aes/src/operations/round_key.rs +++ b/aes/src/operations/round_key.rs @@ -1,12 +1,9 @@ use crate::{Block128, key::Subkey}; -pub const fn add_round_key(state: Block128, subkeys: &[Subkey; 4]) -> Block128 { - let k0 = subkeys[0].as_u128(); - let k1 = subkeys[1].as_u128(); - let k2 = subkeys[2].as_u128(); - let k3 = subkeys[3].as_u128(); +pub fn add_round_key(state: Block128, subkeys: &[Subkey; 4]) -> Block128 { + let [k0, k1, k2, k3] = [subkeys[0], subkeys[1], subkeys[2], subkeys[3]]; let key_block = (k0 << 96) | (k1 << 64) | (k2 << 32) | k3; - Block128::new(state.as_u128() ^ key_block) + state ^ key_block } #[cfg(test)] @@ -15,11 +12,28 @@ mod tests { use rstest::rstest; const TEST_KEY: u128 = 0x0F15_71C9_47D9_E859_1CB7_ADD6_AF7F_6798; + const TEST_MESSAGE: u128 = 0x0123_4567_89AB_CDEF_FEDC_BA98_7654_3210; #[rstest] + #[case( + 0x0123_4567_89AB_CDEF_FEDC_BA98_7654_3210, + [0x0F15_71C9, 0x47D9_E859, 0x1CB7_ADD6, 0xAF7F_6798], + 0x0E36_34AE_CE72_25B6_E26B_174E_D92B_5588 + )] + fn round_key_addition(#[case] input: u128, #[case] subkeys: [u32; 4], #[case] expected: u128) { + let block = Block128::new(input); + let subkeys = subkeys.map(Subkey::from_u32); + let result = add_round_key(block, &subkeys).as_u128(); + assert_eq!( + result, expected, + "Adding Round Key failed. Expected 0x{expected:032X}, got 0x{result:032X}", + ); + } + + #[rstest] + #[case(TEST_MESSAGE)] #[case(0x0000_0000_0000_0000_0000_0000_0000_0000)] #[case(0xFFFF_FFFF_FFFF_FFFF_FFFF_FFFF_FFFF_FFFF)] - #[case(0x1234_5678_9ABC_DEF0_1234_5678_9ABC_DEF0)] fn add_round_key_roundtrip(#[case] plaintext: u128) { use crate::Aes; diff --git a/aes/tests/aes.rs b/aes/tests/aes.rs index b36ea00..670c14a 100644 --- a/aes/tests/aes.rs +++ b/aes/tests/aes.rs @@ -2,50 +2,36 @@ use aes::Aes; use cipher_core::BlockCipher; use rstest::rstest; -const TEST_KEY: u128 = 0x2B7E_1516_28AE_D2A6_ABF7_1588_09CF_4F3C; -const TEST_PLAINTEXT: u128 = 0x6BC1_BEE2_2E40_9F96_E93D_7E11_7393_1728; -const TEST_CIPHERTEXT: u128 = 0x79BD_98A6_CB0F_D3AE_3D7D_C1A3_3CD3_6E2F; +const TEST_KEY: u128 = 0x0F15_71C9_47D9_E859_1CB7_ADD6_AF7F_6798; +const TEST_PLAINTEXT: u128 = 0x0123_4567_89AB_CDEF_FEDC_BA98_7654_3210; +const TEST_CIPHERTEXT: u128 = 0x34D3_F0EE_CB4D_FA16_CB8B_F07F_29A0_CB79; -// #[rstest] -// #[case(TEST_PLAINTEXT, TEST_CIPHERTEXT, TEST_KEY)] -// #[case( -// 0xAE2D_8A57_1E03_AC9C_9EB7_6FAC_45AF_8E51, -// 0xBC14_003D_01C9_B46C_AC63_D481_5210_E80B, -// 0xF5D3_D585_03B9_699D_E785_895A_96FD_BAAF -// )] -// #[case( -// 0x30C8_1C46_A35C_E411_E5FB_C119_1A0A_52EF, -// 0xB9AF_FEE2_98CD_0F4A_6708_44A6_D6CE_EF87, -// 0x43B1_CD7F_598E_CE23_881B_00E3_ED03_0688 -// )] -// #[case( -// 0xF69F_2445_DF4F_9B17_AD2B_417B_E66C_3710, -// 0xA279_FA71_A91B_9FA9_213C_E13E_659D_5C3B, -// 0x7B0C_785E_27E8_AD3F_8223_2071_0472_5DD4 -// )] -// fn encrypt_decrypt_roundtrip( -// #[case] plaintext: u128, -// #[case] expected_ciphertext: u128, -// #[case] key: u128, -// ) { -// let aes = Aes::new(key); -// let pt_bytes = plaintext.to_be_bytes(); -// -// // Test Encrypt -// let ciphertext = aes.encrypt(&pt_bytes).expect("Encryption failed"); -// let ciphertext_u128 = u128::from_be_bytes(ciphertext.as_slice().try_into().unwrap()); -// -// assert_eq!( -// ciphertext_u128, expected_ciphertext, -// "Encryption mismatch. Expected 0x{expected_ciphertext:032X}, got 0x{ciphertext_u128:032X}" -// ); -// -// // Test Decrypt -// let decrypted = aes.decrypt(&ciphertext).expect("Decryption failed"); -// let decrypted_u128 = u128::from_be_bytes(decrypted.as_slice().try_into().unwrap()); -// -// assert_eq!( -// decrypted_u128, plaintext, -// "Decryption mismatch. Expected 0x{plaintext:032X}, got 0x{decrypted_u128:032X}" -// ); -// } +#[rstest] +#[case(TEST_PLAINTEXT, TEST_CIPHERTEXT, TEST_KEY)] +fn encrypt_decrypt_roundtrip( + #[case] plaintext: u128, + #[case] expected_ciphertext: u128, + #[case] key: u128, +) { + let aes = Aes::new(key); + + // Encrypt + let ciphertext = aes + .encrypt(&plaintext.to_be_bytes()) + .expect("Encryption failed"); + let ciphertext_u128 = u128::from_be_bytes(ciphertext.as_slice().try_into().unwrap()); + + assert_eq!( + ciphertext_u128, expected_ciphertext, + "Encryption failed. Expected 0x{expected_ciphertext:032X}, got 0x{ciphertext_u128:032X}" + ); + + // Decrypt + let decrypted = aes.decrypt(&ciphertext).expect("Decryption failed"); + let decrypted_u128 = u128::from_be_bytes(decrypted.as_slice().try_into().unwrap()); + + assert_eq!( + decrypted_u128, plaintext, + "Decryption failed. Expected 0x{plaintext:032X}, got 0x{decrypted_u128:032X}" + ); +}