Merge pull request #5 from 0xphen/feat/block-modes

Feat/block modes

aes/Cargo.toml

@@ -7,6 +7,7 @@ edition = "2021"
 
 [dependencies]
 lazy_static = "1.4.0"
+rand = "0.8.5"
 rayon = "1.8.0"
 serial_test = "2.0.0"
 thiserror = "1.0.50"

aes/src/aes_ops.rs

@@ -0,0 +1,214 @@
+use super::{
+    constants::{AES_S_BOX, TRANSFORMATION_MATRIX},
+    key_schedule::KeySchedule,
+    util::{galois_mul, xor_matrices},
+};
+
+pub struct AesOps;
+
+impl AesOps {
+    /// Performs AES encryption on the given state.
+    ///
+    /// The `state` is a mutable reference to a vector of 4-byte arrays,
+    /// which is encrypted using the provided key schedule. The encryption
+    /// modifies the `state` in place, resulting in the ciphertext.
+    ///
+    /// The AES encryption process consists of:
+    /// - An initial AddRoundKey step.
+    /// - Several rounds (number specified by `keys.rounds`) of:
+    ///   - SubBytes: A non-linear substitution step.
+    ///   - ShiftRows: A transposition step.
+    ///   - MixColumns: A mixing operation applied to each column.
+    ///   - AddRoundKey: Combined with a round key derived from the encryption key.
+    /// - A final round that includes SubBytes, ShiftRows, and AddRoundKey,
+    ///   but omits the MixColumns step.
+    ///
+    /// # Arguments
+    /// * `state` - A mutable reference to the AES state to be encrypted.
+    /// * `keys` - A reference to the `KeySchedule` used for the encryption.
+    ///
+    /// # Notes
+    /// The final encrypted state, or ciphertext, is stored in the `state`
+    /// after the completion of this method. As the encryption is done in place,
+    /// the input `state` is overwritten with the encrypted data.
+    pub fn encrypt(state: &mut [[u8; 4]; 4], keys: &KeySchedule) {
+        let rounds = keys.rounds;
+        // Add initial round key
+        Self::add_round_key(state, keys.round_key(0));
+
+        // Main encryption rounds
+        for round in 1..(rounds) {
+            Self::sub_bytes(state);
+            Self::shift_rows(state);
+            Self::mix_columns(state);
+            Self::add_round_key(state, keys.round_key(round as usize));
+        }
+
+        //Final round without mixing columns
+        Self::sub_bytes(state);
+        Self::shift_rows(state);
+        Self::add_round_key(state, keys.round_key(rounds as usize));
+    }
+
+    /// Performs the AddRoundKey step, a crucial part of the AES encryption algorithm.
+    ///
+    /// This method XORs each byte of the AES state with the corresponding byte of the given round key.
+    ///
+    /// # Arguments
+    /// * `key` - The round key to be XORed with the AES state.
+    fn add_round_key(state: &mut [[u8; 4]; 4], key: [[u8; 4]; 4]) {
+        *state = xor_matrices(*state, key);
+    }
+
+    /// Performs the SubBytes transformation on the AES state.
+    /// This is a non-linear byte substitution step where each byte is replaced
+    /// with another according to a lookup table (S-box).
+    /// It mutates the current AES state by updating each byte with its substituted value.
+    fn sub_bytes(state: &mut [[u8; 4]; 4]) {
+        // Iterate over each byte of the state matrix
+        for (i, row) in state.iter_mut().enumerate() {
+            for (j, e) in row.iter_mut().enumerate() {
+                // Apply the S-box transformation and store in `new_state`
+                *e = AES_S_BOX[*e as usize];
+            }
+        }
+    }
+
+    /// Performs the "ShiftRows" step in the AES encryption process.
+    /// This function shifts the rows of the state matrix as per AES specification:
+    /// - The first row is not shifted.
+    /// - Each subsequent row is shifted to the left by an offset equal to its row index.
+    /// - The state matrix is assumed to be column-major, i.e., each inner array represents a column.
+    ///
+    /// # Arguments
+    /// * `&mut self` - A mutable reference to the current instance of the AES struct.
+    fn shift_rows(state: &mut [[u8; 4]; 4]) {
+        // Temporary variable to hold the values for row shifting
+        let mut temp: [u8; 4] = [0; 4];
+
+        for i in 1..4 {
+            for j in 0..4 {
+                // Store the shifted row in a temporary variable
+                temp[j] = state[(j + i) % 4][i];
+            }
+            for j in 0..4 {
+                // Update the state with the shifted values
+                state[j][i] = temp[j];
+            }
+        }
+    }
+
+    /// Performs the MixColumns transformation on the AES state.
+    ///
+    /// This function applies the MixColumns step to each column of the AES state matrix.
+    /// It uses the Galois Field multiplication (`galois_mul`) for the transformation.
+    ///
+    /// # Arguments
+    /// * `&mut self` - A mutable reference to the AES structure, containing the state matrix.
+    fn mix_columns(state: &mut [[u8; 4]; 4]) {
+        for col in 0..4 {
+            // Temporary storage for the column being processed
+            let mut temp_column = [0u8; 4];
+
+            // Transform the current column using Galois Field multiplication
+            for i in 0..4 {
+                temp_column[i] = galois_mul(TRANSFORMATION_MATRIX[i][0], state[col][0])
+                    ^ galois_mul(TRANSFORMATION_MATRIX[i][1], state[col][1])
+                    ^ galois_mul(TRANSFORMATION_MATRIX[i][2], state[col][2])
+                    ^ galois_mul(TRANSFORMATION_MATRIX[i][3], state[col][3]);
+            }
+
+            // Update the state matrix with the transformed column
+            for i in 0..4 {
+                state[col][i] = temp_column[i];
+            }
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn aes_ops_encrypt_test() {
+        let mut state: [[u8; 4]; 4] = [
+            [0, 17, 34, 51],
+            [68, 85, 102, 119],
+            [136, 153, 170, 187],
+            [204, 221, 238, 255],
+        ];
+
+        let key_schedule =
+            KeySchedule::new(&[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15]).unwrap();
+
+        AesOps::encrypt(&mut state, &key_schedule);
+
+        assert_eq!(
+            state,
+            [
+                [105, 196, 224, 216],
+                [106, 123, 4, 48],
+                [216, 205, 183, 128],
+                [112, 180, 197, 90]
+            ]
+        );
+    }
+
+    #[test]
+    fn initial_round_key_and_one_round_test() {
+        let mut state: [[u8; 4]; 4] = [
+            [0, 17, 34, 51],
+            [68, 85, 102, 119],
+            [136, 153, 170, 187],
+            [204, 221, 238, 255],
+        ];
+
+        let key_schedule =
+            KeySchedule::new(&[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15]).unwrap();
+
+        AesOps::add_round_key(&mut state, key_schedule.round_key(0));
+        assert_eq!(
+            state,
+            [
+                [0, 16, 32, 48],
+                [64, 80, 96, 112],
+                [128, 144, 160, 176],
+                [192, 208, 224, 240]
+            ]
+        );
+
+        AesOps::sub_bytes(&mut state);
+        assert_eq!(
+            state,
+            [
+                [99, 202, 183, 4],
+                [9, 83, 208, 81],
+                [205, 96, 224, 231],
+                [186, 112, 225, 140]
+            ]
+        );
+
+        AesOps::shift_rows(&mut state);
+        assert_eq!(
+            state,
+            [
+                [99, 83, 224, 140],
+                [9, 96, 225, 4],
+                [205, 112, 183, 81],
+                [186, 202, 208, 231]
+            ]
+        );
+
+        AesOps::mix_columns(&mut state);
+        assert_eq!(
+            state,
+            [
+                [95, 114, 100, 21],
+                [87, 245, 188, 146],
+                [247, 190, 59, 41],
+                [29, 185, 249, 26]
+            ]
+        );
+    }
+}

aes/src/block_modes.rs

@@ -0,0 +1,123 @@
+use rand::{rngs::OsRng, RngCore};
+
+use super::{
+    aes_ops::AesOps,
+    definitions::{AesEncryptor, PaddingProcessor},
+    error::AesError,
+    key_schedule::KeySchedule,
+    pkcs_padding::PkcsPadding,
+    util::*,
+};
+
+pub struct CbcEncryptor {
+    pub state: Option<Vec<u8>>,
+    pub padding_processor: Box<dyn PaddingProcessor>,
+    pub iv: [[u8; 4]; 4],
+    keys: KeySchedule,
+}
+
+impl CbcEncryptor {
+    /// Generates a 16-byte initialization vector (IV) for AES encryption.
+    ///
+    /// This function uses a cryptographically secure random number generator (OsRng)
+    /// to fill a 16-byte array with random data, which serves as the IV.
+    ///
+    /// Returns:
+    /// A 16-byte array `[u8; 16]` representing the IV.
+    fn gen_iv() -> [u8; 16] {
+        let mut iv = [0u8; 16];
+        OsRng.fill_bytes(&mut iv);
+
+        iv
+    }
+
+    /// Creates a new instance of an AES encryption structure with CBC mode and padding.
+    ///
+    /// Parameters:
+    /// * `pk`: A byte slice representing the encryption key.
+    /// * `padding_processor`: An instance of a type that implements `PaddingProcessor`.
+    ///   This type must have a `'static` lifetime.
+    ///
+    /// Returns:
+    /// A `Result` containing the new instance or an `AesError` on failure.
+    ///
+    /// The function initializes the key schedule for AES based on `pk`,
+    /// sets the initial state and IV, and stores the padding processor.
+    pub fn new<T: PaddingProcessor + 'static>(
+        pk: &[u8],
+        padding_processor: T,
+    ) -> Result<Self, AesError> {
+        Ok(Self {
+            keys: KeySchedule::new(pk)?,
+            state: None,
+            iv: gen_matrix(&Self::gen_iv()),
+            padding_processor: Box::new(padding_processor),
+        })
+    }
+}
+
+impl AesEncryptor for CbcEncryptor {
+    /// Encrypts a message using AES with CBC mode and PKCS padding.
+    ///
+    /// This function encrypts the given message using the AES encryption algorithm in CBC mode.
+    /// PKCS padding is applied to the message to ensure proper block sizing.
+    ///
+    /// # Arguments
+    /// * `message` - A slice of bytes representing the plaintext message to be encrypted.
+    ///
+    /// # Returns
+    /// A `Result` containing a vector of encrypted 4x4 byte matrices (`Vec<[[u8; 4]; 4]>`)
+    /// on success, or an `AesError` on failure.
+    fn encrypt(&mut self, message: &[u8]) -> Result<Vec<[[u8; 4]; 4]>, AesError> {
+        // Convert the message to a byte vector and apply PKCS padding
+        let mut plain_bytes = message.to_vec();
+        PkcsPadding.pad_input(&mut plain_bytes);
+
+        // Chunk the padded message into 4x4 byte matrices
+        let input_blocks = chunk_bytes_into_4x4_matrices(&plain_bytes);
+
+        // Initialize the working state by XORing the first block with the IV
+        let mut working_state = xor_matrices(input_blocks[0], self.iv);
+
+        let mut encrypted_blocks = Vec::with_capacity(input_blocks.len());
+
+        for block in input_blocks {
+            AesOps::encrypt(&mut working_state, &self.keys);
+            encrypted_blocks.push(working_state);
+            working_state = xor_matrices(working_state, block);
+        }
+
+        Ok(encrypted_blocks)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    const INPUT: [u8; 16] = [
+        0, 17, 34, 51, 68, 85, 102, 119, 136, 153, 170, 187, 204, 221, 238, 255,
+    ];
+
+    const PK: [u8; 16] = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15];
+
+    const IV: [u8; 16] = [
+        102, 71, 120, 83, 87, 100, 53, 57, 65, 89, 100, 105, 81, 88, 90, 83,
+    ];
+
+    #[test]
+    fn test_cbc_encryption() {
+        let mut cbc_ops = CbcEncryptor::new(&PK, PkcsPadding).unwrap();
+        cbc_ops.iv =  gen_matrix(&IV);
+
+        let start_cipher_bytes: Vec<[[u8; 4]; 4]> = vec![[
+            [59, 67, 136, 134],
+            [79, 78, 189, 114],
+            [137, 150, 207, 148],
+            [186, 117, 130, 178],
+        ]];
+
+        let result = cbc_ops.encrypt(&INPUT).unwrap();
+        assert!(result.as_slice().starts_with(&start_cipher_bytes));
+    }
+}

aes/src/definitions.rs

@@ -1,8 +1,26 @@
-pub trait Encryptor {
-    fn encrypt(&mut self) -> Vec<u8>;
+use super::error::AesError;
+
+pub trait AesEncryptor {
+    fn encrypt(&mut self, message: &[u8]) -> Result<Vec<[[u8; 4]; 4]>, AesError>;
 }
 
-pub trait PaddingScheme {
-    fn pad_input(input_buffer: &mut Vec<u8>);
-    fn strip_output(output_buffer: &mut Vec<u8>);
-}
+/// Trait for padding processing in cryptographic operations.
+pub trait PaddingProcessor {
+    /// Adds padding to the given input buffer.
+    ///
+    /// # Arguments
+    /// * `input_buffer` - A mutable reference to a vector of bytes representing the input data.
+    fn pad_input(&self, input_buffer: &mut Vec<u8>);
+
+    /// Removes padding from the given output buffer.
+    ///
+    /// # Arguments
+    /// * `output_buffer` - A mutable reference to a vector of bytes representing the output data.
+    fn strip_output(&self, output_buffer: &mut Vec<u8>);
+}
+
+/// Enum representing different padding schemes.
+pub enum PaddingScheme {
+    /// Represents the PKSC padding scheme.
+    PKSC,
+}

aes/src/error.rs

@@ -8,6 +8,12 @@ pub enum AesError {
     #[error("Invalid bits size. Expected 128 got `{0}`")]
     InvalidBitsSize(usize),
 
-    #[error("Fauled to convert matrix to fixed size")]
+    #[error("Failed to convert matrix to fixed size")]
     KeyMatrixConversionError,
+
+    #[error("Failed to generate IV")]
+    IVGenerationError,
+
+    #[error("Failed to parse slice to matrix: {0}")]
+    FailedToParseSliceToMatrix(String),
 }