sdlan-lib-rs/src/utils/encrypter.rs

use std::{
    sync::atomic::{AtomicU32, Ordering},
    time::{SystemTime, UNIX_EPOCH},
};

use bytes::BytesMut;
use chacha20poly1305::{aead::AeadInPlace, ChaCha20Poly1305, KeyInit};
use sdlan_sn_rs::utils::{aes_decrypt, aes_encrypt, Result, SDLanError};

const COUNTER_MASK: u32 = (1 << 24) - 1;

pub trait Encryptor {
    fn is_setted(&self) -> bool;
    fn set_key(&mut self, region_id: u32, key: Vec<u8>);
    fn encrypt(&self, data: &mut BytesMut) -> Result<()>;
    fn decrypt(&self, data: &mut BytesMut) -> Result<()>;
}

pub enum MyEncryptor {
    Invalid,
    ChaChao20(Chacha20Encryptor),
    Aes(AesEncryptor),
}

impl MyEncryptor {
    pub fn new() -> Self {
        Self::Invalid
    }

    pub fn is_setted(&self) -> bool {
        match self {
            Self::Invalid => false,
            Self::Aes(aes) => aes.is_setted(),
            Self::ChaChao20(cha) => cha.is_setted(),
        }
    }

    pub fn set_key(&mut self, region_id: u32, key: Vec<u8>) {
        match self {
            Self::Invalid => {}
            Self::Aes(aes) => {
                aes.set_key(region_id, key);
            }
            Self::ChaChao20(cha) => {
                cha.set_key(region_id, key);
            }
        }
    }

    pub fn encrypt(&self, data: &mut BytesMut) -> Result<()> {
        match self {
            Self::Invalid => Err(SDLanError::EncryptError("invalid encryptor".to_owned())),
            Self::Aes(aes) => aes.encrypt(data),
            Self::ChaChao20(cha) => cha.encrypt(data),
        }
    }

    pub fn decrypt(&self, data: &mut BytesMut) -> Result<()> {
        match self {
            Self::Invalid => Err(SDLanError::EncryptError("invalid encryptor".to_owned())),
            Self::Aes(aes) => aes.decrypt(data),
            Self::ChaChao20(cha) => cha.decrypt(data),
        }
    }
}

pub struct Chacha20Encryptor {
    cipher: ChaCha20Poly1305,
    key: Vec<u8>,
    is_setted: bool,
    next_counter: AtomicU32,
    region_id: u32,
}

impl Chacha20Encryptor {
    pub fn new(key: Vec<u8>, region_id: u32) -> Self {
        Self {
            cipher: chacha20poly1305::ChaCha20Poly1305::new(key.as_slice().into()),
            key,
            is_setted: true,
            next_counter: AtomicU32::new(0),
            region_id,
        }
    }
}

impl Encryptor for Chacha20Encryptor {
    fn set_key(&mut self, region_id: u32, key: Vec<u8>) {
        self.cipher = chacha20poly1305::ChaCha20Poly1305::new(key.as_slice().into());
        self.key = key;
        self.region_id = region_id;
    }

    fn encrypt(&self, data: &mut BytesMut) -> Result<()> {
        let plaintext_len = data.len();
        
        // Prepare nonce
        let now = SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .unwrap()
            .as_millis() as u64;

        let next_counter = self
            .next_counter
            .fetch_update(Ordering::Release, Ordering::Acquire, |current| {
                Some((current + 1) & COUNTER_MASK)
            })
            .unwrap() as u64;

        let mut nonce_bytes = [0u8; 12];
        let region_id_bytes = self.region_id.to_be_bytes();
        nonce_bytes[0..4].copy_from_slice(&region_id_bytes);
        let next_data = (now << 24) | next_counter;
        nonce_bytes[4..12].copy_from_slice(&next_data.to_be_bytes());
        let nonce = chacha20poly1305::Nonce::from_slice(&nonce_bytes);

        // Make room for tag (16 bytes) at the end
        data.resize(plaintext_len + 16, 0);
        let (payload, tag_space) = data.split_at_mut(plaintext_len);

        // Encrypt payload in place and get tag
        let tag = self
            .cipher
            .encrypt_in_place_detached(&nonce, &[], payload)
            .map_err(|e| SDLanError::EncryptError(e.to_string()))?;

        tag_space[..16].copy_from_slice(&tag);

        // Prepend nonce (12 bytes)
        let mut final_buf = BytesMut::with_capacity(12 + data.len());
        final_buf.extend_from_slice(&nonce_bytes);
        final_buf.unsplit(data.split_off(0));
        *data = final_buf;

        Ok(())
    }

    fn decrypt(&self, data: &mut BytesMut) -> Result<()> {
        if data.len() < 28 {
            return Err(SDLanError::EncryptError(
                "ciphered text size error".to_owned(),
            ));
        }

        // Split off the 12-byte Nonce
        let mut payload = data.split_off(12);
        let mut nonce_bytes = [0u8; 12];
        nonce_bytes.copy_from_slice(&data[0..12]);
        let nonce = chacha20poly1305::Nonce::from_slice(&nonce_bytes);

        // Split off the 16-byte Tag
        let ciphertext_len = payload.len() - 16;
        let (ciphertext, tag_space) = payload.split_at_mut(ciphertext_len);
        let tag = chacha20poly1305::Tag::from_slice(&tag_space[..16]);

        // Decrypt in place
        self.cipher
            .decrypt_in_place_detached(&nonce, &[], ciphertext, tag)
            .map_err(|e| SDLanError::EncryptError(format!("failed to decrypt: {}", e.to_string())))?;

        payload.truncate(ciphertext_len);
        *data = payload;

        Ok(())
    }

    fn is_setted(&self) -> bool {
        self.is_setted
    }
}

pub struct AesEncryptor {
    key: Vec<u8>,
    is_setted: bool,
}

impl AesEncryptor {
    pub fn new(key: Vec<u8>) -> Self {
        Self {
            key,
            is_setted: true,
        }
    }
}

impl Encryptor for AesEncryptor {
    fn decrypt(&self, data: &mut BytesMut) -> Result<()> {
        let res = aes_decrypt(&self.key, data)?;
        *data = BytesMut::from(res.as_slice());
        Ok(())
    }

    fn encrypt(&self, data: &mut BytesMut) -> Result<()> {
        let res = aes_encrypt(&self.key, data)?;
        *data = BytesMut::from(res.as_slice());
        Ok(())
    }

    fn is_setted(&self) -> bool {
        self.is_setted
    }

    fn set_key(&mut self, _region_id: u32, key: Vec<u8>) {
        self.key = key;
        self.is_setted = true;
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_chacha20_encrypt_decrypt() {
        let key = vec![0u8; 32];
        let encryptor = Chacha20Encryptor::new(key, 1);
        let mut data = BytesMut::from(&b"hello world"[..]);
        encryptor.encrypt(&mut data).unwrap();
        encryptor.decrypt(&mut data).unwrap();
        assert_eq!(&data[..], b"hello world");
    }
}