refactor: housekeeping changes

This commit is contained in:
Jixun Wu
2024-09-11 22:18:25 +01:00
parent 8fe3fe8849
commit d88ad94a82
2 changed files with 45 additions and 40 deletions

View File

@@ -1,20 +1,10 @@
use super::{ecb, TcTeaError};
use rand::prelude::*; use rand::prelude::*;
use std::cmp::min; use std::cmp::min;
use super::{ecb, TcTeaError}; pub(crate) const SALT_LEN: usize = 2;
pub(crate) const ZERO_LEN: usize = 7;
const SALT_LEN: usize = 2; pub(crate) const FIXED_PADDING_LEN: usize = 1 + SALT_LEN + ZERO_LEN;
const ZERO_LEN: usize = 7;
const FIXED_PADDING_LEN: usize = 1 + SALT_LEN + ZERO_LEN;
/// Calculate expected size of encrypted data.
///
/// `body_size` is the size of data you'd like to encrypt.
pub fn get_encrypted_size(body_size: usize) -> usize {
let len = FIXED_PADDING_LEN + body_size;
let pad_len = (8 - (len & 0b0111)) & 0b0111;
len + pad_len
}
fn xor_tea_block(a: &[u8; 8], b: &[u8; 8]) -> [u8; 8] { fn xor_tea_block(a: &[u8; 8], b: &[u8; 8]) -> [u8; 8] {
let mut dest = *a; let mut dest = *a;
@@ -67,7 +57,7 @@ pub fn encrypt<'a>(
// Set up a header with random padding/salt // Set up a header with random padding/salt
#[cfg(feature = "secure_random")] #[cfg(feature = "secure_random")]
ChaCha20Rng::from_entropy().fill_bytes(&mut header[0..header_len]); rand_chacha::ChaCha20Rng::from_entropy().fill_bytes(&mut header[0..header_len]);
#[cfg(not(feature = "secure_random"))] #[cfg(not(feature = "secure_random"))]
rand_pcg::Pcg32::from_entropy().fill_bytes(&mut header[0..header_len]); rand_pcg::Pcg32::from_entropy().fill_bytes(&mut header[0..header_len]);
@@ -78,10 +68,12 @@ pub fn encrypt<'a>(
header[0] = (header[0] & 0b1111_1000) | ((pad_len as u8) & 0b0000_0111); header[0] = (header[0] & 0b1111_1000) | ((pad_len as u8) & 0b0000_0111);
header[header_len..header_len + copy_to_header_len].copy_from_slice(plain_header); header[header_len..header_len + copy_to_header_len].copy_from_slice(plain_header);
// Access to slice of "cipher" from inner scope
{ {
let mut iv1 = [0u8; 8]; let mut iv1 = [0u8; 8];
let mut iv2 = [0u8; 8]; let mut iv2 = [0u8; 8];
// Process whole blocks
let plain_last_block_len = plain.len() % 8; let plain_last_block_len = plain.len() % 8;
let (plain, plain_last_block) = plain.split_at(plain.len() - plain_last_block_len); let (plain, plain_last_block) = plain.split_at(plain.len() - plain_last_block_len);
@@ -93,13 +85,13 @@ pub fn encrypt<'a>(
encrypt_round(cipher, &header[8..], key, &mut iv1, &mut iv2); encrypt_round(cipher, &header[8..], key, &mut iv1, &mut iv2);
let mut cipher = &mut cipher[8..]; let mut cipher = &mut cipher[8..];
if !plain.is_empty() { // Handle whole blocks
for (plain, cipher) in plain.chunks_exact(8).zip(cipher.chunks_exact_mut(8)) { for (plain, cipher) in plain.chunks_exact(8).zip(cipher.chunks_exact_mut(8)) {
encrypt_round(cipher, plain, key, &mut iv1, &mut iv2); encrypt_round(cipher, plain, key, &mut iv1, &mut iv2);
} }
cipher = &mut cipher[plain.len()..]; cipher = &mut cipher[plain.len()..];
}
// Handle last block, if there's any
if plain_last_block_len != 0 { if plain_last_block_len != 0 {
let mut last_block = [0u8; 8]; let mut last_block = [0u8; 8];
last_block[..plain_last_block_len].copy_from_slice(plain_last_block); last_block[..plain_last_block_len].copy_from_slice(plain_last_block);
@@ -262,15 +254,4 @@ mod tests {
Ok(()) Ok(())
} }
#[test]
fn test_calc_encrypted_size() {
assert_eq!(get_encrypted_size(0), 16);
assert_eq!(get_encrypted_size(1), 16);
assert_eq!(get_encrypted_size(6), 16);
assert_eq!(get_encrypted_size(7), 24);
assert_eq!(get_encrypted_size(14), 24);
assert_eq!(get_encrypted_size(15), 32);
}
} }

View File

@@ -2,7 +2,6 @@
//! //!
//! Notably, it uses a different round number and uses a "tweaked" CBC mode. //! Notably, it uses a different round number and uses a "tweaked" CBC mode.
use crate::cbc::get_encrypted_size;
use byteorder::{ByteOrder, BE}; use byteorder::{ByteOrder, BE};
use thiserror::Error; use thiserror::Error;
@@ -25,6 +24,15 @@ pub enum TcTeaError {
SliceError, SliceError,
} }
/// Calculate expected size of encrypted data.
///
/// `body_size` is the size of data you'd like to encrypt.
pub fn get_encrypted_size(body_size: usize) -> usize {
let len = cbc::FIXED_PADDING_LEN + body_size;
let pad_len = (8 - (len & 0b0111)) & 0b0111;
len + pad_len
}
/// Parse key to u32 array /// Parse key to u32 array
pub fn parse_key(key: &[u8]) -> Result<[u32; 4], TcTeaError> { pub fn parse_key(key: &[u8]) -> Result<[u32; 4], TcTeaError> {
let key_chunks = match key.len() { let key_chunks = match key.len() {
@@ -80,8 +88,23 @@ pub fn decrypt<T: AsRef<[u8]>>(encrypted: T, key: &[u8]) -> Result<Vec<u8>, TcTe
Ok(Vec::from(result)) Ok(Vec::from(result))
} }
#[test] #[cfg(test)]
fn test_sanity_test() -> Result<(), TcTeaError> { mod tests {
use super::*;
#[test]
fn test_calc_encrypted_size() {
assert_eq!(get_encrypted_size(0), 16);
assert_eq!(get_encrypted_size(1), 16);
assert_eq!(get_encrypted_size(6), 16);
assert_eq!(get_encrypted_size(7), 24);
assert_eq!(get_encrypted_size(14), 24);
assert_eq!(get_encrypted_size(15), 32);
}
#[test]
fn test_sanity_test() -> Result<(), TcTeaError> {
let key = b"43218765dcbahgfe"; let key = b"43218765dcbahgfe";
let message = b"this is a test message."; let message = b"this is a test message.";
let cipher = encrypt(message, key)?; let cipher = encrypt(message, key)?;
@@ -89,4 +112,5 @@ fn test_sanity_test() -> Result<(), TcTeaError> {
assert_eq!(message, plain.as_slice()); assert_eq!(message, plain.as_slice());
Ok(()) Ok(())
}
} }