use core::fmt;
use tc_block_cipher::{
BlockCipher, BlockCipherInit, BlockError, CipherDirection, InitError, KeyParams,
};
use tc_zeroize::Zeroize;
use crate::cipher::ROUNDS;
use crate::{ALGO_NAME, BLOCK_BYTES, KEY_BYTES, cipher};
pub struct Sm4Engine {
round_keys: [u32; ROUNDS],
initialised: bool,
}
impl Sm4Engine {
pub const fn new() -> Self {
Self {
round_keys: [0; ROUNDS],
initialised: false,
}
}
}
impl Default for Sm4Engine {
fn default() -> Self {
Self::new()
}
}
impl fmt::Display for Sm4Engine {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.write_str(ALGO_NAME)
}
}
impl Drop for Sm4Engine {
fn drop(&mut self) {
self.round_keys.zeroize();
}
}
impl BlockCipher for Sm4Engine {
type Error = BlockError;
fn block_size(&self) -> usize {
BLOCK_BYTES
}
fn process_block(&mut self, input: &[u8], output: &mut [u8]) -> Result<usize, BlockError> {
if !self.initialised {
return Err(BlockError::NotInitialised);
}
let (Some(input), Some(output)) = (
input.first_chunk::<BLOCK_BYTES>(),
output.first_chunk_mut::<BLOCK_BYTES>(),
) else {
return Err(BlockError::BufferTooShort);
};
cipher::process_block(&self.round_keys, input, output);
Ok(BLOCK_BYTES)
}
}
impl<P: KeyParams + ?Sized> BlockCipherInit<P> for Sm4Engine {
type Error = InitError;
fn init(&mut self, direction: CipherDirection, params: &P) -> Result<(), InitError> {
let key = params.key();
let key: &[u8; KEY_BYTES] = key
.try_into()
.map_err(|_| InitError::InvalidKeyLength(key.len()))?;
let for_encryption = direction == CipherDirection::Encrypt;
self.round_keys.zeroize();
cipher::expand_key(for_encryption, key, &mut self.round_keys);
self.initialised = true;
Ok(())
}
}