将chacha20的初始化放到结构体里面,以减少开销

This commit is contained in:
alex 2026-06-30 09:48:59 +08:00
parent 1de337f3d6
commit d0ec1e29c6

View File

@ -1,8 +1,10 @@
use std::{sync::atomic::{AtomicU32, Ordering}, time::{SystemTime, UNIX_EPOCH}}; use std::{
sync::atomic::{AtomicU32, Ordering},
time::{SystemTime, UNIX_EPOCH},
};
use chacha20poly1305::{aead::Aead, ChaCha20Poly1305, KeyInit};
use chacha20poly1305::{KeyInit, aead::Aead}; use sdlan_sn_rs::utils::{aes_decrypt, aes_encrypt, Result, SDLanError};
use sdlan_sn_rs::utils::{Result, SDLanError, aes_decrypt, aes_encrypt};
const COUNTER_MASK: u32 = (1 << 24) - 1; const COUNTER_MASK: u32 = (1 << 24) - 1;
@ -27,12 +29,8 @@ impl MyEncryptor {
pub fn is_setted(&self) -> bool { pub fn is_setted(&self) -> bool {
match self { match self {
Self::Invalid => false, Self::Invalid => false,
Self::Aes(aes) => { Self::Aes(aes) => aes.is_setted(),
aes.is_setted() Self::ChaChao20(cha) => cha.is_setted(),
}
Self::ChaChao20(cha) => {
cha.is_setted()
}
} }
} }
@ -50,33 +48,22 @@ impl MyEncryptor {
pub fn encrypt(&self, data: &[u8]) -> Result<Vec<u8>> { pub fn encrypt(&self, data: &[u8]) -> Result<Vec<u8>> {
match self { match self {
Self::Invalid => { Self::Invalid => Err(SDLanError::EncryptError("invalid encryptor".to_owned())),
Err(SDLanError::EncryptError("invalid encryptor".to_owned())) Self::Aes(aes) => aes.encrypt(data),
} Self::ChaChao20(cha) => cha.encrypt(data),
Self::Aes(aes) => {
aes.encrypt(data)
}
Self::ChaChao20(cha) => {
cha.encrypt(data)
}
} }
} }
pub fn decrypt(&self, ciphered: &[u8]) -> Result<Vec<u8>> { pub fn decrypt(&self, ciphered: &[u8]) -> Result<Vec<u8>> {
match self { match self {
Self::Invalid => { Self::Invalid => Err(SDLanError::EncryptError("invalid encryptor".to_owned())),
Err(SDLanError::EncryptError("invalid encryptor".to_owned())) Self::Aes(aes) => aes.decrypt(ciphered),
} Self::ChaChao20(cha) => cha.decrypt(ciphered),
Self::Aes(aes) => {
aes.decrypt(ciphered)
}
Self::ChaChao20(cha) => {
cha.decrypt(ciphered)
}
} }
} }
} }
pub struct Chacha20Encryptor { pub struct Chacha20Encryptor {
cipher: ChaCha20Poly1305,
key: Vec<u8>, key: Vec<u8>,
is_setted: bool, is_setted: bool,
next_counter: AtomicU32, next_counter: AtomicU32,
@ -86,6 +73,7 @@ pub struct Chacha20Encryptor {
impl Chacha20Encryptor { impl Chacha20Encryptor {
pub fn new(key: Vec<u8>, region_id: u32) -> Self { pub fn new(key: Vec<u8>, region_id: u32) -> Self {
Self { Self {
cipher: chacha20poly1305::ChaCha20Poly1305::new(key.as_slice().into()),
key, key,
is_setted: true, is_setted: true,
next_counter: AtomicU32::new(0), next_counter: AtomicU32::new(0),
@ -96,17 +84,24 @@ impl Chacha20Encryptor {
impl Encryptor for Chacha20Encryptor { impl Encryptor for Chacha20Encryptor {
fn set_key(&mut self, region_id: u32, key: Vec<u8>) { fn set_key(&mut self, region_id: u32, key: Vec<u8>) {
self.cipher = chacha20poly1305::ChaCha20Poly1305::new(key.as_slice().into());
self.key = key; self.key = key;
self.region_id = region_id; self.region_id = region_id;
} }
fn encrypt(&self, data: &[u8]) -> Result<Vec<u8>> { fn encrypt(&self, data: &[u8]) -> Result<Vec<u8>> {
let cipher = chacha20poly1305::ChaCha20Poly1305::new(self.key.as_slice().into()); // let cipher = chacha20poly1305::ChaCha20Poly1305::new(self.key.as_slice().into());
let now = SystemTime::now().duration_since(UNIX_EPOCH).unwrap().as_millis() as u64; 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| { let next_counter = self
.next_counter
.fetch_update(Ordering::Release, Ordering::Acquire, |current| {
Some((current + 1) & COUNTER_MASK) Some((current + 1) & COUNTER_MASK)
}).unwrap() as u64; })
.unwrap() as u64;
let mut nonce = Vec::new(); let mut nonce = Vec::new();
let region_id = self.region_id.to_be_bytes(); let region_id = self.region_id.to_be_bytes();
@ -114,28 +109,29 @@ impl Encryptor for Chacha20Encryptor {
let next_data = (now << 24) | next_counter; let next_data = (now << 24) | next_counter;
nonce.extend_from_slice(&next_data.to_be_bytes()); nonce.extend_from_slice(&next_data.to_be_bytes());
match cipher.encrypt(nonce.as_slice().into(), data) { match self.cipher.encrypt(nonce.as_slice().into(), data) {
Ok(data) => { Ok(data) => {
nonce.extend_from_slice(&data); nonce.extend_from_slice(&data);
Ok(nonce) Ok(nonce)
},
Err(e) => {
Err(SDLanError::EncryptError(e.to_string()))
} }
Err(e) => Err(SDLanError::EncryptError(e.to_string())),
} }
} }
fn decrypt(&self, ciphered: &[u8]) -> Result<Vec<u8>> { fn decrypt(&self, ciphered: &[u8]) -> Result<Vec<u8>> {
if ciphered.len() < 12 { if ciphered.len() < 12 {
return Err(SDLanError::EncryptError("ciphered text size error".to_owned())) return Err(SDLanError::EncryptError(
"ciphered text size error".to_owned(),
));
} }
let cipher = chacha20poly1305::ChaCha20Poly1305::new(self.key.as_slice().into()); // let cipher = chacha20poly1305::ChaCha20Poly1305::new(self.key.as_slice().into());
let nonce = &ciphered[0..12]; let nonce = &ciphered[0..12];
match cipher.decrypt(nonce.into(), &ciphered[12..]) { match self.cipher.decrypt(nonce.into(), &ciphered[12..]) {
Ok(data) => Ok(data), Ok(data) => Ok(data),
Err(e) => { Err(e) => Err(SDLanError::EncryptError(format!(
Err(SDLanError::EncryptError(format!("failed to decyrpt: {}", e.to_string()))) "failed to decyrpt: {}",
} e.to_string()
))),
} }
} }
@ -176,4 +172,3 @@ impl Encryptor for AesEncryptor {
self.is_setted = true; self.is_setted = true;
} }
} }