refactor: use u64 instead of u8[8] buffer; allow user to override salt.

This commit is contained in:
Jixun Wu
2024-09-11 23:54:52 +01:00
parent 22422fa268
commit 8fcc07199b
2 changed files with 48 additions and 53 deletions

View File

@@ -1,42 +1,33 @@
use super::{ecb, TcTeaError}; use super::{ecb_impl, TcTeaError};
use rand::prelude::*; use byteorder::{ByteOrder, BE};
use std::cmp::min; use std::cmp::min;
pub(crate) const SALT_LEN: usize = 2; pub(crate) const SALT_LEN: usize = 2;
pub(crate) const ZERO_LEN: usize = 7; pub(crate) const ZERO_LEN: usize = 7;
pub(crate) const FIXED_PADDING_LEN: usize = 1 + SALT_LEN + ZERO_LEN; pub(crate) const FIXED_PADDING_LEN: usize = 1 + SALT_LEN + ZERO_LEN;
fn xor_tea_block(a: &[u8; 8], b: &[u8; 8]) -> [u8; 8] {
let mut dest = *a;
dest.iter_mut().zip(b).for_each(|(a, b)| *a ^= *b);
dest
}
#[inline(always)]
fn encrypt_round( fn encrypt_round(
cipher: &mut [u8], cipher: &mut [u8],
plain: &[u8], plain: &[u8],
key: &[u32; 4], key: &[u32; 4],
iv1: &mut [u8; 8], iv1: u64,
iv2: &mut [u8; 8], iv2: u64,
) { ) -> (u64, u64) {
let mut plain_block = [0u8; 8]; let plain_block = BE::read_u64(plain);
plain_block.copy_from_slice(plain);
let iv2_next = xor_tea_block(&plain_block, iv1); let iv2_next = plain_block ^ iv1;
let mut result = iv2_next; let result = ecb_impl::encrypt(iv2_next, key);
ecb::encrypt(&mut result, key); let cipher_block = result ^ iv2;
BE::write_u64(cipher, cipher_block);
let cipher_block = xor_tea_block(&result, iv2); (cipher_block, iv2_next)
*iv1 = cipher_block;
*iv2 = iv2_next;
cipher[..8].copy_from_slice(&cipher_block);
} }
pub fn encrypt<'a>( pub fn encrypt<'a>(
cipher: &'a mut [u8], cipher: &'a mut [u8],
plain: &[u8], plain: &[u8],
key: &[u32; 4], key: &[u32; 4],
salt: &[u8; 10],
) -> Result<&'a [u8], TcTeaError> { ) -> Result<&'a [u8], TcTeaError> {
// buffer size calculation // buffer size calculation
let len = FIXED_PADDING_LEN + plain.len(); let len = FIXED_PADDING_LEN + plain.len();
@@ -56,10 +47,7 @@ pub fn encrypt<'a>(
let mut header = [0u8; 16]; let mut header = [0u8; 16];
// Set up a header with random padding/salt // Set up a header with random padding/salt
#[cfg(feature = "secure_random")] header[..header_len].copy_from_slice(&salt[..header_len]);
rand_chacha::ChaCha20Rng::from_entropy().fill_bytes(&mut header[0..header_len]);
#[cfg(not(feature = "secure_random"))]
rand_pcg::Pcg32::from_entropy().fill_bytes(&mut header[0..header_len]);
// Build header // Build header
let copy_to_header_len = min(16 - header_len, plain.len()); let copy_to_header_len = min(16 - header_len, plain.len());
@@ -70,8 +58,8 @@ pub fn encrypt<'a>(
// Access to slice of "cipher" from inner scope // Access to slice of "cipher" from inner scope
{ {
let mut iv1 = [0u8; 8]; let mut iv1 = 0u64;
let mut iv2 = [0u8; 8]; let mut iv2 = 0u64;
// Process whole blocks // Process whole blocks
let plain_last_block_len = plain.len() % 8; let plain_last_block_len = plain.len() % 8;
@@ -79,15 +67,15 @@ pub fn encrypt<'a>(
// Encrypt first 2 blocks from the header, then whole blocks // Encrypt first 2 blocks from the header, then whole blocks
// cbc_encrypt_round(cipher, &header, key, &mut iv1, &mut iv2); // cbc_encrypt_round(cipher, &header, key, &mut iv1, &mut iv2);
encrypt_round(cipher, &header[..8], key, &mut iv1, &mut iv2); (iv1, iv2) = encrypt_round(cipher, &header[..8], key, iv1, iv2);
let cipher = &mut cipher[8..]; let cipher = &mut cipher[8..];
encrypt_round(cipher, &header[8..], key, &mut iv1, &mut iv2); (iv1, iv2) = encrypt_round(cipher, &header[8..], key, iv1, iv2);
let mut cipher = &mut cipher[8..]; let mut cipher = &mut cipher[8..];
// Handle whole blocks // 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); (iv1, iv2) = encrypt_round(cipher, plain, key, iv1, iv2);
} }
cipher = &mut cipher[plain.len()..]; cipher = &mut cipher[plain.len()..];
@@ -95,7 +83,7 @@ pub fn encrypt<'a>(
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);
encrypt_round(cipher, &last_block, key, &mut iv1, &mut iv2); encrypt_round(cipher, &last_block, key, iv1, iv2);
} }
} }
@@ -103,25 +91,20 @@ pub fn encrypt<'a>(
Ok(cipher) Ok(cipher)
} }
#[inline(always)]
fn decrypt_round( fn decrypt_round(
plain: &mut [u8], plain: &mut [u8],
cipher: &[u8], cipher: &[u8],
key: &[u32; 4], key: &[u32; 4],
iv1: &mut [u8; 8], iv1: u64,
iv2: &mut [u8; 8], iv2: u64,
) { ) -> (u64, u64) {
let mut cipher_block = [0u8; 8]; let cipher_block = BE::read_u64(cipher);
cipher_block.copy_from_slice(cipher); let result = cipher_block ^ iv2;
let next_iv2 = ecb_impl::decrypt(result, key);
let plain_block = next_iv2 ^ iv1;
let mut result = xor_tea_block(&cipher_block, iv2); BE::write_u64(plain, plain_block);
ecb::decrypt(&mut result, key); (cipher_block, next_iv2)
let plain_block = xor_tea_block(&result, iv1);
*iv1 = cipher_block;
*iv2 = result;
plain[..8].copy_from_slice(&plain_block);
} }
pub fn decrypt<'a>( pub fn decrypt<'a>(
@@ -139,10 +122,11 @@ pub fn decrypt<'a>(
} }
let plain = &mut plain[..input_len]; let plain = &mut plain[..input_len];
let mut iv1 = [0u8; 8];
let mut iv2 = [0u8; 8]; let mut iv1 = 0u64;
let mut iv2 = 0u64;
for (cipher, plain) in cipher.chunks_exact(8).zip(plain.chunks_exact_mut(8)) { for (cipher, plain) in cipher.chunks_exact(8).zip(plain.chunks_exact_mut(8)) {
decrypt_round(plain, cipher, key, &mut iv1, &mut iv2); (iv1, iv2) = decrypt_round(plain, cipher, key, iv1, iv2);
} }
let pad_size = usize::from(plain[0] & 0b111); let pad_size = usize::from(plain[0] & 0b111);
@@ -163,6 +147,8 @@ pub fn decrypt<'a>(
mod tests { mod tests {
use super::*; use super::*;
const TEST_SALT: [u8; 10] = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9];
// Known good data, generated from its C++ implementation // Known good data, generated from its C++ implementation
const GOOD_ENCRYPTED_DATA: [u8; 24] = [ const GOOD_ENCRYPTED_DATA: [u8; 24] = [
0x91, 0x09, 0x51, 0x62, 0xe3, 0xf5, 0xb6, 0xdc, // 0x91, 0x09, 0x51, 0x62, 0xe3, 0xf5, 0xb6, 0xdc, //
@@ -196,7 +182,7 @@ mod tests {
#[test] #[test]
fn tc_tea_encrypt_empty() -> Result<(), TcTeaError> { fn tc_tea_encrypt_empty() -> Result<(), TcTeaError> {
let mut cipher_buffer = [0xffu8; 100]; let mut cipher_buffer = [0xffu8; 100];
let cipher = encrypt(&mut cipher_buffer, b"", &ENCRYPTION_KEY)?; let cipher = encrypt(&mut cipher_buffer, b"", &ENCRYPTION_KEY, &TEST_SALT)?;
assert_eq!(cipher.len(), 16); assert_eq!(cipher.len(), 16);
let mut plain = vec![0xffu8; 24]; let mut plain = vec![0xffu8; 24];
@@ -210,7 +196,7 @@ mod tests {
#[test] #[test]
fn tc_tea_basic_encryption() -> Result<(), TcTeaError> { fn tc_tea_basic_encryption() -> Result<(), TcTeaError> {
let mut cipher_buffer = [0xffu8; 100]; let mut cipher_buffer = [0xffu8; 100];
let cipher = encrypt(&mut cipher_buffer, &EXPECTED_PLAIN_TEXT, &ENCRYPTION_KEY)?; let cipher = encrypt(&mut cipher_buffer, &EXPECTED_PLAIN_TEXT, &ENCRYPTION_KEY, &TEST_SALT)?;
assert_eq!(cipher.len(), 24); assert_eq!(cipher.len(), 24);
let mut plain = vec![0xffu8; 24]; let mut plain = vec![0xffu8; 24];
@@ -226,7 +212,7 @@ mod tests {
let mut cipher_buffer = [0xffu8; 100]; let mut cipher_buffer = [0xffu8; 100];
let input = b"...test data by Jixun ... ... test hello aaa"; let input = b"...test data by Jixun ... ... test hello aaa";
for _ in 0..16 { for _ in 0..16 {
let cipher = encrypt(&mut cipher_buffer, input, &ENCRYPTION_KEY)?; let cipher = encrypt(&mut cipher_buffer, input, &ENCRYPTION_KEY, &TEST_SALT)?;
assert_eq!(cipher.len() % 8, 0); assert_eq!(cipher.len() % 8, 0);
assert!(cipher.len() > input.len()); assert!(cipher.len() > input.len());
@@ -247,7 +233,7 @@ mod tests {
let input = b"...test data by Jixun ... ... test hello aaa"; let input = b"...test data by Jixun ... ... test hello aaa";
for test_len in 0usize..input.len() { for test_len in 0usize..input.len() {
let input = &input[..test_len]; let input = &input[..test_len];
let cipher = encrypt(&mut cipher_buffer, input, &ENCRYPTION_KEY)?; let cipher = encrypt(&mut cipher_buffer, input, &ENCRYPTION_KEY, &TEST_SALT)?;
let decrypted = decrypt(&mut plain_buffer, cipher, &ENCRYPTION_KEY)?; let decrypted = decrypt(&mut plain_buffer, cipher, &ENCRYPTION_KEY)?;
assert_eq!(decrypted, input); assert_eq!(decrypted, input);
} }

View File

@@ -3,11 +3,13 @@
//! 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 byteorder::{ByteOrder, BE}; use byteorder::{ByteOrder, BE};
use rand::RngCore;
use thiserror::Error; use thiserror::Error;
pub mod cbc; pub mod cbc;
pub mod ecb; pub mod ecb;
mod ecb_impl; mod ecb_impl;
use rand::prelude::*;
#[derive(Error, Debug, PartialEq)] #[derive(Error, Debug, PartialEq)]
pub enum TcTeaError { pub enum TcTeaError {
@@ -68,7 +70,14 @@ pub fn encrypt<T: AsRef<[u8]>>(plaintext: T, key: &[u8]) -> Result<Vec<u8>, TcTe
let plaintext = plaintext.as_ref(); let plaintext = plaintext.as_ref();
let cipher_len = get_encrypted_size(plaintext.len()); let cipher_len = get_encrypted_size(plaintext.len());
let mut cipher = vec![0u8; cipher_len]; let mut cipher = vec![0u8; cipher_len];
cbc::encrypt(&mut cipher, plaintext, &key)?;
let mut salt = [0u8; 10];
#[cfg(feature = "secure_random")]
rand_chacha::ChaCha20Rng::from_entropy().fill_bytes(&mut salt);
#[cfg(not(feature = "secure_random"))]
rand_pcg::Pcg32::from_entropy().fill_bytes(&mut salt);
cbc::encrypt(&mut cipher, plaintext, &key, &salt)?;
Ok(cipher) Ok(cipher)
} }