core/common/src/utils/crypto.rs

// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License.  You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.

use crate::IggyError;
use crate::text;
use aes_gcm::aead::{Aead, Generate};
use aes_gcm::{Aes256Gcm, KeyInit, Nonce};
use std::fmt::Debug;

#[derive(Debug, Clone)]
pub enum EncryptorKind {
    Aes256Gcm(Aes256GcmEncryptor),
}

impl EncryptorKind {
    pub fn encrypt(&self, data: &[u8]) -> Result<Vec<u8>, IggyError> {
        match self {
            EncryptorKind::Aes256Gcm(e) => e.encrypt(data),
        }
    }
    pub fn decrypt(&self, data: &[u8]) -> Result<Vec<u8>, IggyError> {
        match self {
            EncryptorKind::Aes256Gcm(e) => e.decrypt(data),
        }
    }
}

pub trait Encryptor {
    fn encrypt(&self, data: &[u8]) -> Result<Vec<u8>, IggyError>;
    fn decrypt(&self, data: &[u8]) -> Result<Vec<u8>, IggyError>;
}

#[derive(Clone)]
pub struct Aes256GcmEncryptor {
    cipher: Aes256Gcm,
}

impl Debug for Aes256GcmEncryptor {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("Encryptor").finish()
    }
}

impl Aes256GcmEncryptor {
    pub fn new(key: &[u8]) -> Result<Self, IggyError> {
        Ok(Self {
            cipher: Aes256Gcm::new_from_slice(key).map_err(|_| IggyError::InvalidEncryptionKey)?,
        })
    }

    pub fn from_base64_key(key: &str) -> Result<Self, IggyError> {
        Self::new(&text::from_base64_as_bytes(key)?)
    }
}

impl Encryptor for Aes256GcmEncryptor {
    fn encrypt(&self, data: &[u8]) -> Result<Vec<u8>, IggyError> {
        let nonce = Nonce::generate();
        let encrypted_data = self.cipher.encrypt(&nonce, data);
        if encrypted_data.is_err() {
            return Err(IggyError::CannotEncryptData);
        }
        let payload = [&nonce, encrypted_data.unwrap().as_slice()].concat();
        Ok(payload)
    }

    fn decrypt(&self, data: &[u8]) -> Result<Vec<u8>, IggyError> {
        let nonce = (&data[0..12])
            .try_into()
            .map_err(|_| IggyError::CannotDecryptData)?;
        let payload = self.cipher.decrypt(nonce, &data[12..]);
        if payload.is_err() {
            return Err(IggyError::CannotDecryptData);
        }
        Ok(payload.unwrap())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::{HeaderKey, HeaderValue, IggyMessage};

    #[test]
    fn given_the_same_key_data_should_be_encrypted_and_decrypted_correctly() {
        let key = [1; 32];
        let encryptor = Aes256GcmEncryptor::new(&key).unwrap();
        let data = b"Hello World!";
        let encrypted_data = encryptor.encrypt(data);
        assert!(encrypted_data.is_ok());
        let encrypted_data = encrypted_data.unwrap();
        let decrypted_data = encryptor.decrypt(&encrypted_data);
        assert!(decrypted_data.is_ok());
        let decrypted_data = decrypted_data.unwrap();
        assert_eq!(data, decrypted_data.as_slice());
    }

    #[test]
    fn given_the_invalid_key_data_should_not_be_decrypted_correctly() {
        let first_key = [1; 32];
        let second_key = [2; 32];
        let first_encryptor = Aes256GcmEncryptor::new(&first_key).unwrap();
        let second_encryptor = Aes256GcmEncryptor::new(&second_key).unwrap();
        let data = b"Hello World!";
        let encrypted_data = first_encryptor.encrypt(data);
        assert!(encrypted_data.is_ok());
        let encrypted_data = encrypted_data.unwrap();
        let decrypted_data = second_encryptor.decrypt(&encrypted_data);
        assert!(decrypted_data.is_err());
        let error = decrypted_data.err().unwrap();
        assert_eq!(error.as_code(), IggyError::CannotDecryptData.as_code());
    }

    #[test]
    fn message_payload_and_headers_should_encrypt_and_decrypt_correctly() {
        use bytes::Bytes;
        use std::collections::BTreeMap;

        let key = [1; 32];
        let encryptor = Aes256GcmEncryptor::new(&key).unwrap();

        let mut headers = BTreeMap::new();
        headers.insert(
            HeaderKey::try_from("batch").unwrap(),
            HeaderValue::from(1u64),
        );
        headers.insert(
            HeaderKey::try_from("type").unwrap(),
            HeaderValue::try_from("test-message").unwrap(),
        );

        let mut message = IggyMessage::builder()
            .payload(Bytes::from("test payload data"))
            .user_headers(headers)
            .build()
            .unwrap();

        let original_payload = message.payload.clone();
        let original_headers = message.user_headers.clone().unwrap();

        message.payload = Bytes::from(encryptor.encrypt(&message.payload).unwrap());
        message.header.payload_length = message.payload.len() as u32;

        let encrypted_headers = encryptor.encrypt(&original_headers).unwrap();
        message.header.user_headers_length = encrypted_headers.len() as u32;
        message.user_headers = Some(Bytes::from(encrypted_headers));

        assert_ne!(message.payload, original_payload);
        assert_ne!(message.user_headers.as_ref().unwrap(), &original_headers);

        let decrypted_payload = encryptor.decrypt(&message.payload).unwrap();
        message.payload = Bytes::from(decrypted_payload);
        message.header.payload_length = message.payload.len() as u32;

        let decrypted_headers = encryptor
            .decrypt(message.user_headers.as_ref().unwrap())
            .unwrap();
        message.header.user_headers_length = decrypted_headers.len() as u32;
        message.user_headers = Some(Bytes::from(decrypted_headers));

        assert_eq!(message.payload, original_payload);
        assert_eq!(message.user_headers.as_ref().unwrap(), &original_headers);

        let parsed = message.user_headers_map().unwrap().unwrap();
        assert_eq!(
            parsed
                .get(&HeaderKey::try_from("batch").unwrap())
                .unwrap()
                .as_uint64()
                .unwrap(),
            1
        );
        assert_eq!(
            parsed
                .get(&HeaderKey::try_from("type").unwrap())
                .unwrap()
                .as_str()
                .unwrap(),
            "test-message"
        );
    }
}