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

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


use chacha20poly1305::{KeyInit, aead::Aead};
use sdlan_sn_rs::utils::{Result, SDLanError, aes_decrypt, aes_encrypt};

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: &[u8]) -> Result<Vec<u8>>;
    fn decrypt(&self, ciphered: &[u8]) -> Result<Vec<u8>>;
}

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: &[u8]) -> Result<Vec<u8>> {
        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, ciphered: &[u8]) -> Result<Vec<u8>> {
        match self {
            Self::Invalid => {
                Err(SDLanError::EncryptError("invalid encryptor".to_owned()))
            }
            Self::Aes(aes) => {
                aes.decrypt(ciphered)
            }
            Self::ChaChao20(cha) => {
                cha.decrypt(ciphered)
            }
        }
    }
}

pub struct Chacha20Encryptor {
    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 {
            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.key = key;
        self.region_id = region_id;
    }

    fn encrypt(&self, data: &[u8]) -> Result<Vec<u8>> {
        let cipher = chacha20poly1305::ChaCha20Poly1305::new(self.key.as_slice().into());
        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 = Vec::new();
        let region_id = self.region_id.to_be_bytes();
        nonce.extend_from_slice(&region_id);
        let next_data = (now<<24) | next_counter;
        nonce.extend_from_slice(&next_data.to_be_bytes());

        match cipher.encrypt(nonce.as_slice().into(), data) {
            Ok(data) => {
                nonce.extend_from_slice(&data);
                Ok(nonce)
            },
            Err(e) => {
                Err(SDLanError::EncryptError(e.to_string()))
            }
        }
    }

    fn decrypt(&self, ciphered: &[u8]) -> Result<Vec<u8>> {
        if ciphered.len() < 12 {
            return Err(SDLanError::EncryptError("ciphered text size error".to_owned()))
        }
        let cipher = chacha20poly1305::ChaCha20Poly1305::new(self.key.as_slice().into());
        let nonce = &ciphered[0..12];
        match cipher.decrypt(nonce.into(), &ciphered[12..]) {
            Ok(data) => Ok(data),
            Err(e) => {
                Err(SDLanError::EncryptError(format!("failed to decyrpt: {}", e.to_string())))
            }
        }
    }

    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, ciphered: &[u8]) -> Result<Vec<u8>> {
        aes_decrypt(&self.key, ciphered)
    }

    fn encrypt(&self, data: &[u8]) -> Result<Vec<u8>> {
        aes_encrypt(&self.key, data)
    }

    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;
    }
}