trelent-hyok 0.1.12

//! AES-GCM-256 encryption implementation.
//!
//! This module provides an implementation of the AES-GCM-256 authenticated encryption algorithm.
//! It includes functionality for encryption and decryption of data using a 256-bit key,
//! along with Additional Authenticated Data (AAD) for enhanced security.

use std::time::{SystemTime, UNIX_EPOCH};
use crate::dek::DEK;
use crate::error::encryption::EncryptionError;
use crate::held_data::encryption::EncryptionStrategy;
use aes_gcm::aead::generic_array::GenericArray;
use aes_gcm::aead::{Aead, OsRng, Payload};
use aes_gcm::{AeadCore, Aes256Gcm, Key, KeyInit};
use async_trait::async_trait;

/// Size of the Additional Authenticated Data (AAD) in bytes
pub const AAD_SIZE: usize = 16;
/// Size of the nonce (number used once) in bytes
pub const NONCE_SIZE: usize = 12;
/// Size of the encryption key in bytes
pub const KEY_SIZE: usize = 32;

/// Implementation of the AES-GCM-256 encryption strategy
#[derive(Debug, Clone)]
pub struct AesGcm256Strategy;

/// Contains the necessary data for AES-GCM encryption/decryption operations
#[derive(Debug, Clone)]
pub struct AesGcmEncryptionData {
    /// Additional Authenticated Data used for integrity checking
    pub aad: [u8; AAD_SIZE],
    /// Nonce (number used once) for encryption
    pub nonce: [u8; NONCE_SIZE],
}

#[async_trait]
impl EncryptionStrategy for AesGcm256Strategy {
    type EncryptionData = AesGcmEncryptionData;

    /// Encrypts data using AES-GCM-256
    ///
    /// # Arguments
    /// * `dek` - Data Encryption Key
    /// * `plaintext` - Data to encrypt
    /// * `encryption_data` - Contains AAD and nonce for encryption
    ///
    /// # Returns
    /// * `Result<Vec<u8>, EncryptionError>` - Encrypted data or error
    async fn encrypt(
        &self,
        dek: DEK,
        plaintext: Vec<u8>,
        encryption_data: Self::EncryptionData
    ) -> Result<Vec<u8>, EncryptionError> {
        let cipher = get_cipher(dek)?;
        let AesGcmEncryptionData { aad, nonce } = encryption_data;
        let nonce = GenericArray::from(nonce);
        let payload = Payload {
            msg: plaintext.as_slice(),
            aad: &aad,
        };
        let encrypted = cipher
            .encrypt(&nonce.clone(), payload)
            .map_err(|e| EncryptionError::new(format!("Could not encrypt! Error: {:?}", e)))?;
        let ciphertext = merge_aad_and_ciphertext(encrypted, encryption_data);
        Ok(ciphertext)
    }

    /// Decrypts data using AES-GCM-256
    ///
    /// # Arguments
    /// * `dek` - Data Encryption Key
    /// * `ciphertext` - Encrypted data
    /// * `encryption_data` - Contains AAD and nonce for decryption
    ///
    /// # Returns
    /// * `Result<Vec<u8>, EncryptionError>` - Decrypted data or error
    async fn decrypt(
        &self,
        dek: DEK,
        ciphertext: Vec<u8>,
        encryption_data: Self::EncryptionData
    ) -> Result<Vec<u8>, EncryptionError> {
        let cipher = get_cipher(dek)?;
        let AesGcmEncryptionData { aad, nonce } = encryption_data;
        let nonce = GenericArray::from(nonce);
        let payload = Payload {
            msg: ciphertext.as_slice(),
            aad: &aad,
        };
        cipher
            .decrypt(&nonce, payload)
            .map_err(|e| EncryptionError::new(format!("Error Decrypting, Error: {:?}", e)))
    }
}

/// Creates an AES-GCM cipher instance from a DEK
///
/// # Arguments
/// * `dek` - Data Encryption Key to use for cipher creation
///
/// # Returns
/// * `Result<Aes256Gcm, EncryptionError>` - Cipher instance or error
fn get_cipher(dek: DEK) -> Result<Aes256Gcm, EncryptionError> {
    let key: Vec<u8> = dek.into();
    if key.len() != KEY_SIZE {
        return Err(
            EncryptionError::new(
                format!(
                    "DEK must have a length of 32 bits to use AES GCM 256, instead has length of {}",
                    key.len()
                )
            )
        );
    }
    let key: &[u8; KEY_SIZE] = key
        .as_slice()
        .try_into()
        .map_err(|_| EncryptionError::new("Could not convert DEK into array".to_string()))?;
    let key: &Key<Aes256Gcm> = key.into();
    Ok(Aes256Gcm::new(key))
}

/// Splits combined encryption data into ciphertext and metadata
///
/// # Arguments
/// * `inp` - Combined input data containing AAD, nonce, and ciphertext
///
/// # Returns
/// * `Result<(Vec<u8>, AesGcmEncryptionData), EncryptionError>` - Tuple of ciphertext and encryption data
pub fn split_encryption_data(
    inp: Vec<u8>
) -> Result<(Vec<u8>, AesGcmEncryptionData), EncryptionError> {
    if inp.len() < AAD_SIZE + NONCE_SIZE {
        return Err(
            EncryptionError::new(
                format!(
                    "Size of input needs to be larger than {}, but was {}",
                    AAD_SIZE + NONCE_SIZE,
                    inp.len()
                )
            )
        );
    }
    let aad: &[u8] = &inp[0..AAD_SIZE];
    let nonce: &[u8] = &inp[AAD_SIZE..AAD_SIZE + NONCE_SIZE];
    let ciphertext: &[u8] = &inp[AAD_SIZE + NONCE_SIZE..];
    let result = (
        ciphertext.to_vec(),
        AesGcmEncryptionData {
            aad: aad
                .try_into()
                .map_err(|_| EncryptionError::new("Could not properly resize aad".to_string()))?,
            nonce: nonce
                .try_into()
                .map_err(|_| EncryptionError::new("Could not properly resize nonce".to_string()))?,
        },
    );
    Ok(result)
}

/// Combines ciphertext with encryption metadata
///
/// # Arguments
/// * `ciphertext` - Encrypted data
/// * `aad` - Encryption metadata including AAD and nonce
///
/// # Returns
/// * `Vec<u8>` - Combined data
pub fn merge_aad_and_ciphertext(ciphertext: Vec<u8>, aad: AesGcmEncryptionData) -> Vec<u8> {
    let mut new_vec: Vec<u8> = vec![];
    new_vec.append(&mut aad.aad.clone().to_vec());
    new_vec.append(&mut aad.nonce.clone().to_vec());
    new_vec.append(&mut ciphertext.clone());
    new_vec
}

/// Generates a new nonce using system time and random bits
///
/// The nonce is composed of:
/// - 4 bytes of system time (milliseconds)
/// - 8 bytes of random data
///
/// # Returns
/// * `Result<[u8; NONCE_SIZE], EncryptionError>` - Generated nonce or error
pub fn generate_nonce() -> Result<[u8; NONCE_SIZE], EncryptionError> {
    let mut nonce_vec = Vec::new();
    let mut time_millis = SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .map_err(|e| EncryptionError::new(format!("System Time Error, Error: {:?}", e)))?
        .as_millis()
        .to_be_bytes()
        .to_vec();
    time_millis.resize(4, 0);
    let nonce_random_bits_vec = Aes256Gcm::generate_nonce(&mut OsRng)
        .as_slice()
        .split_at(8)
        .0.to_vec();
    nonce_vec.append(&mut time_millis.clone());
    nonce_vec.append(&mut nonce_random_bits_vec.clone());
    nonce_vec.resize(NONCE_SIZE, 0);
    nonce_vec
        .as_slice()
        .try_into()
        .map_err(|e| EncryptionError(format!("Could not convert nonce into array, Error: {:?}", e)))
}