将chacha20的初始化放到结构体里面,以减少开销
This commit is contained in:
parent
1de337f3d6
commit
d0ec1e29c6
@ -1,14 +1,16 @@
|
|||||||
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 sdlan_sn_rs::utils::{aes_decrypt, aes_encrypt, Result, SDLanError};
|
||||||
|
|
||||||
use chacha20poly1305::{KeyInit, aead::Aead};
|
const COUNTER_MASK: u32 = (1 << 24) - 1;
|
||||||
use sdlan_sn_rs::utils::{Result, SDLanError, aes_decrypt, aes_encrypt};
|
|
||||||
|
|
||||||
const COUNTER_MASK: u32 = (1<<24) - 1;
|
|
||||||
|
|
||||||
pub trait Encryptor {
|
pub trait Encryptor {
|
||||||
fn is_setted(&self) -> bool;
|
fn is_setted(&self) -> bool;
|
||||||
fn set_key(&mut self, region_id: u32, key:Vec<u8>);
|
fn set_key(&mut self, region_id: u32, key: Vec<u8>);
|
||||||
fn encrypt(&self, data: &[u8]) -> Result<Vec<u8>>;
|
fn encrypt(&self, data: &[u8]) -> Result<Vec<u8>>;
|
||||||
fn decrypt(&self, ciphered: &[u8]) -> Result<Vec<u8>>;
|
fn decrypt(&self, ciphered: &[u8]) -> Result<Vec<u8>>;
|
||||||
}
|
}
|
||||||
@ -27,16 +29,12 @@ 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()
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
pub fn set_key(&mut self, region_id: u32, key:Vec<u8>) {
|
pub fn set_key(&mut self, region_id: u32, key: Vec<u8>) {
|
||||||
match self {
|
match self {
|
||||||
Self::Invalid => {}
|
Self::Invalid => {}
|
||||||
Self::Aes(aes) => {
|
Self::Aes(aes) => {
|
||||||
@ -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),
|
||||||
@ -95,47 +83,55 @@ 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();
|
||||||
nonce.extend_from_slice(®ion_id);
|
nonce.extend_from_slice(®ion_id);
|
||||||
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()
|
||||||
|
))),
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@ -171,9 +167,8 @@ impl Encryptor for AesEncryptor {
|
|||||||
self.is_setted
|
self.is_setted
|
||||||
}
|
}
|
||||||
|
|
||||||
fn set_key(&mut self, _region_id: u32, key:Vec<u8>) {
|
fn set_key(&mut self, _region_id: u32, key: Vec<u8>) {
|
||||||
self.key = key;
|
self.key = key;
|
||||||
self.is_setted = true;
|
self.is_setted = true;
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
Loading…
x
Reference in New Issue
Block a user