lxl66566 · June 11, 2025 02:59
diff --git a/README.md b/README.md
diff --git a/aes_cfb_bench.rs b/aes_cfb_bench.rs
 // benches/aes_cfb_bench.rs

 use aes_simd_rust::*; // Your crate name
 use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion, Throughput};
 use openssl::symm::{Cipher, Crypter, Mode};
 use rand::rngs::StdRng;
 use rand::{RngCore, SeedableRng};
 use std::hint::black_box;
 use std::mem;

 // Re-define M128i if not public or use crate::M128i
 #[cfg(target_arch = "x86")]
 use std::arch::x86::__m128i as M128iBench;
 #[cfg(target_arch = "x86_64")]
 use std::arch::x86_64::__m128i as M128iBench;

 fn openssl_encrypt_routine(
    plaintext: &[u8],
    key: &[u8; 32],
    iv: &[u8; 16],
    output_buffer: &mut Vec<u8>, // Re-use buffer
 ) {
    output_buffer.clear();
    output_buffer.resize(
        plaintext.len() + Cipher::aes_256_cfb128().block_size() - 1,
        0,
    );

    let cipher = Cipher::aes_256_cfb128();
    let mut crypter = Crypter::new(cipher, Mode::Encrypt, key, Some(iv)).unwrap();
    crypter.pad(false);

    let mut count = crypter.update(plaintext, output_buffer).unwrap();
    count += crypter.finalize(&mut output_buffer[count..]).unwrap();
    output_buffer.truncate(count);
 }

 fn bench_aes_cfb_encryption(c: &mut Criterion) {
    if !check_aesni_support() {
        println!("Skipping AES-NI benchmark: CPU does not support required features.");
        // Fallback or skip: For now, we just print and skip.
        // Alternatively, one could implement a non-SIMD version for comparison
        // or error out if AES-NI is essential.
        return;
    }

    let mut rng = StdRng::seed_from_u64(12345); // Consistent seed for reproducible benchmarks
    let mut key = [0u8; 32];
    let mut iv = [0u8; 16];
    rng.fill_bytes(&mut key);
    rng.fill_bytes(&mut iv);

    // Pre-compute round keys for custom SIMD version
    let mut round_key_custom: [M128iBench; 15] = unsafe { mem::zeroed() };
    unsafe {
        // Assuming aes256_key_expansion is in the crate root or imported correctly
        aes_simd_rust::aes256_key_expansion(&key, &mut round_key_custom);
    }

    let mut group = c.benchmark_group("AES-256-CFB Encryption");

    // Test with various data sizes
    for size in [64, 256, 1024, 4096, 16384, 65536].iter() {
        let mut plaintext = vec![0u8; *size];
        rng.fill_bytes(&mut plaintext);

        // Buffer for custom SIMD output
        let mut ciphertext_custom_out = vec![0u8; *size];
        // Buffer for OpenSSL output (re-used)
        let mut ciphertext_openssl_out = Vec::with_capacity(*size + 32);

        group.throughput(Throughput::Bytes(*size as u64));

        group.bench_with_input(BenchmarkId::new("Custom SIMD", size), &plaintext, |b, p| {
            b.iter(|| unsafe {
                // Assuming aes256_cfb_encrypt is in the crate root or imported correctly
                aes_simd_rust::aes256_cfb_encrypt(
                    black_box(&mut ciphertext_custom_out),
                    black_box(p),
                    black_box(&iv),
                    black_box(&round_key_custom),
                )
            });
        });

        group.bench_with_input(BenchmarkId::new("OpenSSL", size), &plaintext, |b, p| {
            b.iter(|| {
                openssl_encrypt_routine(
                    black_box(p),
                    black_box(&key),
                    black_box(&iv),
                    black_box(&mut ciphertext_openssl_out),
                )
            });
        });
    }
    group.finish();
 }

 criterion_group!(benches, bench_aes_cfb_encryption);
 criterion_main!(benches);
diff --git a/Cargo.toml b/Cargo.toml
 [package]
 name = "aes_simd_rust"
 version = "0.1.0"
 edition = "2021"

 [dependencies]
 openssl = "*"
 rand = "0.9"  # For generating test data

 [dev-dependencies]
 criterion = { version = "0.6", features = ["html_reports"] }

 [[bench]]
 name = "aes_cfb_bench"
 harness = false
diff --git a/lib.rs b/lib.rs
 #![feature(portable_simd)] // For some intrinsics, though many are stable.
                           // Not strictly needed if relying on std::arch and target_feature.

 // Import necessary SIMD intrinsics
 #[cfg(target_arch = "x86")]
 use std::arch::x86::*;
 #[cfg(target_arch = "x86_64")]
 use std::arch::x86_64::*;

 use std::mem;

 // Type alias for clarity, similar to C++
 pub type M128i = __m128i;

 // Helper macro for AES-256 key expansion (from Intel's whitepaper/common implementations)
 // This macro assists in generating one 128-bit word of the key schedule.
 // It corresponds to the sequence:
 // temp = _mm_aeskeygenassist_si128(key2, rcon);
 // temp = _mm_shuffle_epi32(temp, 0xFF); // or 0xAA for the other part
 // key1 = _mm_xor_si128(key1, _mm_slli_si128(key1, 4));
 // key1 = _mm_xor_si128(key1, _mm_slli_si128(key1, 8));
 // key1 = _mm_xor_si128(key1, temp);
 #[macro_export]
 macro_rules! aes256_assist {
    ($key1:expr, $key2:expr, $rcon:expr, $shuffle_mask:expr) => {{
        let mut temp = _mm_aeskeygenassist_si128($key2, $rcon);
        temp = _mm_shuffle_epi32(temp, $shuffle_mask);
        let mut k1_mod = $key1;
        k1_mod = _mm_xor_si128(k1_mod, _mm_slli_si128(k1_mod, 4));
        k1_mod = _mm_xor_si128(k1_mod, _mm_slli_si128(k1_mod, 8));
        // k1_mod = _mm_xor_si128(k1_mod, _mm_slli_si128(k1_mod, 12)); // This is not in the C++ code, it uses two 4/8 shifts
        k1_mod = _mm_xor_si128(k1_mod, temp);
        k1_mod
    }};
 }

 /// AES-256 Key Expansion
 ///
 /// Expands a 32-byte key into 15 128-bit round keys.
 ///
 /// # Safety
 ///
 /// This function is unsafe because it uses CPU intrinsics that require specific CPU features (AES-NI, SSE2).
 /// The caller must ensure these features are available.
 /// `key` must point to 32 valid bytes.
 /// `round_key` must be a mutable slice of 15 `M128i` elements.
 #[target_feature(enable = "aes,sse2")]
 pub unsafe fn aes256_key_expansion(key: &[u8; 32], round_key: &mut [M128i; 15]) {
    let mut k1 = _mm_loadu_si128(key.as_ptr() as *const M128i);
    let mut k2 = _mm_loadu_si128(key.as_ptr().add(16) as *const M128i);

    round_key[0] = k1;
    round_key[1] = k2;

    // RCON values for AES-256 (only 7 are needed for the 14 rounds of key expansion)
    // The C++ code uses specific RCONs for _mm_aeskeygenassist_si128
    // 0x01, 0x00, 0x02, 0x00, 0x04, 0x00, 0x08, 0x00, 0x10, 0x00, 0x20, 0x00, 0x40
    // The shuffle masks are 0xFF and 0xAA

    // Round 1 (generates round_key[2])
    k1 = aes256_assist!(k1, k2, 0x01, 0xff); // shuffle FF for [3,3,3,3]
    round_key[2] = k1;

    // Round 2 (generates round_key[3])
    k2 = aes256_assist!(k2, k1, 0x00, 0xaa); // shuffle AA for [2,2,2,2]
    round_key[3] = k2;

    // Round 3 (generates round_key[4])
    k1 = aes256_assist!(k1, k2, 0x02, 0xff);
    round_key[4] = k1;

    // Round 4 (generates round_key[5])
    k2 = aes256_assist!(k2, k1, 0x00, 0xaa);
    round_key[5] = k2;

    // Round 5 (generates round_key[6])
    k1 = aes256_assist!(k1, k2, 0x04, 0xff);
    round_key[6] = k1;

    // Round 6 (generates round_key[7])
    k2 = aes256_assist!(k2, k1, 0x00, 0xaa);
    round_key[7] = k2;

    // Round 7 (generates round_key[8])
    k1 = aes256_assist!(k1, k2, 0x08, 0xff);
    round_key[8] = k1;

    // Round 8 (generates round_key[9])
    k2 = aes256_assist!(k2, k1, 0x00, 0xaa);
    round_key[9] = k2;

    // Round 9 (generates round_key[10])
    k1 = aes256_assist!(k1, k2, 0x10, 0xff);
    round_key[10] = k1;

    // Round 10 (generates round_key[11])
    k2 = aes256_assist!(k2, k1, 0x00, 0xaa);
    round_key[11] = k2;

    // Round 11 (generates round_key[12])
    k1 = aes256_assist!(k1, k2, 0x20, 0xff);
    round_key[12] = k1;

    // Round 12 (generates round_key[13])
    k2 = aes256_assist!(k2, k1, 0x00, 0xaa);
    round_key[13] = k2;

    // Round 13 (generates round_key[14])
    k1 = aes256_assist!(k1, k2, 0x40, 0xff);
    round_key[14] = k1;
 }

 /// AES-256 Encrypt a single 128-bit block
 ///
 /// # Safety
 ///
 /// This function is unsafe because it uses CPU intrinsics that require specific CPU features (AES-NI).
 /// The caller must ensure these features are available.
 /// `round_key` must contain 15 valid round keys.
 #[target_feature(enable = "aes")]
 pub unsafe fn aes256_encrypt_block(mut block: M128i, round_key: &[M128i; 15]) -> M128i {
    block = _mm_xor_si128(block, round_key[0]); // Initial AddRoundKey

    // 13 full rounds
    for i in 1..14 {
        block = _mm_aesenc_si128(block, round_key[i]);
    }

    // Final round (no MixColumns)
    block = _mm_aesenclast_si128(block, round_key[14]);
    block
 }

 /// AES-256-CFB Encryption
 ///
 /// Encrypts `plaintext` into `ciphertext` using CFB mode.
 /// `plaintext` and `ciphertext` must have the same length.
 /// `iv` must be 16 bytes.
 /// `round_key` must be the expanded AES-256 key schedule (15 round keys).
 ///
 /// # Safety
 ///
 /// This function is unsafe due to the use of SIMD intrinsics and pointer operations.
 /// Caller must ensure CPU features (AES-NI, SSE2) are available.
 /// Input slices must be valid and correctly sized.
 #[target_feature(enable = "aes,sse2")]
 pub unsafe fn aes256_cfb_encrypt(
    ciphertext: &mut [u8],
    plaintext: &[u8],
    iv: &[u8; 16],
    round_key: &[M128i; 15],
 ) {
    assert_eq!(
        plaintext.len(),
        ciphertext.len(),
        "Plaintext and ciphertext buffers must have the same length."
    );
    let len = plaintext.len();
    if len == 0 {
        return;
    }

    let mut feedback = _mm_loadu_si128(iv.as_ptr() as *const M128i);

    let mut processed_bytes = 0;
    let block_size = 16;

    // Process full blocks
    while processed_bytes + block_size <= len {
        let keystream = aes256_encrypt_block(feedback, round_key);

        let plain_block_ptr = plaintext.as_ptr().add(processed_bytes) as *const M128i;
        let plain_block = _mm_loadu_si128(plain_block_ptr);

        let cipher_block = _mm_xor_si128(plain_block, keystream);

        let cipher_out_ptr = ciphertext.as_mut_ptr().add(processed_bytes) as *mut M128i;
        _mm_storeu_si128(cipher_out_ptr, cipher_block);

        feedback = cipher_block; // Next feedback is the current ciphertext block
        processed_bytes += block_size;
    }

    // Process remaining partial block
    if processed_bytes < len {
        let keystream = aes256_encrypt_block(feedback, round_key);
        let keystream_bytes: [u8; 16] = mem::transmute(keystream); // Get bytes from M128i

        for i in 0..(len - processed_bytes) {
            ciphertext[processed_bytes + i] = plaintext[processed_bytes + i] ^ keystream_bytes[i];
        }
    }
 }

 /// AES-256-CFB Decryption
 ///
 /// Decrypts `ciphertext` into `plaintext` using CFB mode.
 /// `plaintext` and `ciphertext` must have the same length.
 /// `iv` must be 16 bytes.
 /// `round_key` must be the expanded AES-256 key schedule (15 round keys).
 ///
 /// # Safety
 ///
 /// This function is unsafe due to the use of SIMD intrinsics and pointer operations.
 /// Caller must ensure CPU features (AES-NI, SSE2) are available.
 /// Input slices must be valid and correctly sized.
 #[target_feature(enable = "aes,sse2")]
 pub unsafe fn aes256_cfb_decrypt(
    plaintext: &mut [u8],
    ciphertext: &[u8],
    iv: &[u8; 16],
    round_key: &[M128i; 15],
 ) {
    assert_eq!(
        plaintext.len(),
        ciphertext.len(),
        "Plaintext and ciphertext buffers must have the same length."
    );
    let len = ciphertext.len();
    if len == 0 {
        return;
    }

    let mut feedback = _mm_loadu_si128(iv.as_ptr() as *const M128i);

    let mut processed_bytes = 0;
    let block_size = 16;

    // Process full blocks
    while processed_bytes + block_size <= len {
        let keystream = aes256_encrypt_block(feedback, round_key);

        let cipher_block_ptr = ciphertext.as_ptr().add(processed_bytes) as *const M128i;
        let cipher_block = _mm_loadu_si128(cipher_block_ptr);

        let plain_block = _mm_xor_si128(cipher_block, keystream);

        let plain_out_ptr = plaintext.as_mut_ptr().add(processed_bytes) as *mut M128i;
        _mm_storeu_si128(plain_out_ptr, plain_block);

        feedback = cipher_block; // Next feedback is the current ciphertext block
        processed_bytes += block_size;
    }

    // Process remaining partial block
    if processed_bytes < len {
        let keystream = aes256_encrypt_block(feedback, round_key);
        let keystream_bytes: [u8; 16] = mem::transmute(keystream); // Get bytes from M128i

        for i in 0..(len - processed_bytes) {
            plaintext[processed_bytes + i] = ciphertext[processed_bytes + i] ^ keystream_bytes[i];
        }
    }
 }

 // Helper for tests and benchmarks
 pub fn check_aesni_support() -> bool {
    is_x86_feature_detected!("aes") && is_x86_feature_detected!("sse2")
 }

 #[cfg(test)]
 mod tests {
    use super::*;
    use openssl::symm::{Cipher, Crypter, Mode};
    use rand::rngs::StdRng;
    use rand::{RngCore, SeedableRng};

    fn print_hex(label: &str, data: &[u8]) {
        print!("{} ({} bytes): ", label, data.len());
        for byte in data.iter().take(32) {
            // Print first 32 bytes
            print!("{:02x}", byte);
        }
        if data.len() > 32 {
            print!("...");
        }
        println!();
    }

    fn openssl_aes256_cfb128_encrypt(
        plaintext: &[u8],
        key: &[u8; 32],
        iv: &[u8; 16],
    ) -> Result<Vec<u8>, openssl::error::ErrorStack> {
        let cipher = Cipher::aes_256_cfb128();
        let mut crypter = Crypter::new(cipher, Mode::Encrypt, key, Some(iv))?;
        crypter.pad(false); // CFB mode typically doesn't use padding

        let mut ciphertext = vec![0; plaintext.len() + cipher.block_size() - 1]; // Max possible size for update
        let mut count = crypter.update(plaintext, &mut ciphertext)?;
        count += crypter.finalize(&mut ciphertext[count..])?;
        ciphertext.truncate(count);
        Ok(ciphertext)
    }

    fn openssl_aes256_cfb128_decrypt(
        ciphertext: &[u8],
        key: &[u8; 32],
        iv: &[u8; 16],
    ) -> Result<Vec<u8>, openssl::error::ErrorStack> {
        let cipher = Cipher::aes_256_cfb128();
        let mut crypter = Crypter::new(cipher, Mode::Decrypt, key, Some(iv))?;
        crypter.pad(false);

        let mut plaintext = vec![0; ciphertext.len() + cipher.block_size() - 1];
        let mut count = crypter.update(ciphertext, &mut plaintext)?;
        count += crypter.finalize(&mut plaintext[count..])?;
        plaintext.truncate(count);
        Ok(plaintext)
    }

    #[test]
    fn test_aes256_cfb_correctness() {
        if !check_aesni_support() {
            println!("Skipping AES-NI test: CPU does not support required features.");
            return;
        }

        let key: [u8; 32] = [
            0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0a, 0x0b, 0x0c, 0x0d,
            0x0e, 0x0f, 0x10, 0x11, 0x12, 0x13, 0x14, 0x15, 0x16, 0x17, 0x18, 0x19, 0x1a, 0x1b,
            0x1c, 0x1d, 0x1e, 0x1f,
        ];

        let iv: [u8; 16] = [
            0x00, 0x11, 0x22, 0x33, 0x44, 0x55, 0x66, 0x77, 0x88, 0x99, 0xaa, 0xbb, 0xcc, 0xdd,
            0xee, 0xff,
        ];

        let plaintext_str = "Hello, AES-256-CFB! This is a test message for encryption and decryption. \
                             This additional text is to make the test data longer for more meaningful performance testing. \
                             AES (Advanced Encryption Standard) is a specification for the encryption of electronic data \
                             established by the U.S. National Institute of Standards and Technology (NIST) in 2001.";
        let plaintext = plaintext_str.as_bytes();
        let len = plaintext.len();

        // Initialize round keys (unsafe block as it calls target_feature enabled function)
        let mut round_key: [M128i; 15] = unsafe { mem::zeroed() };
        unsafe {
            aes256_key_expansion(&key, &mut round_key);
        }

        // Custom SIMD Encryption
        let mut ciphertext_custom = vec![0u8; len];
        unsafe {
            aes256_cfb_encrypt(&mut ciphertext_custom, plaintext, &iv, &round_key);
        }

        // OpenSSL Encryption
        let ciphertext_openssl = openssl_aes256_cfb128_encrypt(plaintext, &key, &iv).unwrap();

        print_hex("Plaintext      ", plaintext);
        print_hex("Custom Ciphertext", &ciphertext_custom);
        print_hex("OpenSSL Ciphertext", &ciphertext_openssl);

        assert_eq!(
            ciphertext_custom.len(),
            ciphertext_openssl.len(),
            "Ciphertext lengths differ"
        );
        assert_eq!(
            ciphertext_custom, ciphertext_openssl,
            "Encryption results differ between custom and OpenSSL"
        );

        // Custom SIMD Decryption
        let mut decrypted_custom = vec![0u8; len];
        unsafe {
            aes256_cfb_decrypt(&mut decrypted_custom, &ciphertext_custom, &iv, &round_key);
        }
        print_hex("Custom Decrypted ", &decrypted_custom);
        assert_eq!(
            decrypted_custom, plaintext,
            "Custom decryption failed to recover original plaintext"
        );

        // OpenSSL Decryption (of OpenSSL's ciphertext for sanity)
        let decrypted_openssl =
            openssl_aes256_cfb128_decrypt(&ciphertext_openssl, &key, &iv).unwrap();
        assert_eq!(decrypted_openssl, plaintext, "OpenSSL decryption failed");

        println!("AES-256-CFB correctness test passed!");
    }

    #[test]
    fn test_various_lengths() {
        if !check_aesni_support() {
            println!("Skipping AES-NI length test: CPU does not support required features.");
            return;
        }
        let mut rng = StdRng::seed_from_u64(42);
        let mut key = [0u8; 32];
        let mut iv = [0u8; 16];
        rng.fill_bytes(&mut key);
        rng.fill_bytes(&mut iv);

        let mut round_key: [M128i; 15] = unsafe { mem::zeroed() };
        unsafe {
            aes256_key_expansion(&key, &mut round_key);
        }

        for len in [0, 1, 15, 16, 17, 31, 32, 33, 100, 1024].iter() {
            let mut plaintext = vec![0u8; *len];
            rng.fill_bytes(&mut plaintext);

            let mut ciphertext_custom = vec![0u8; *len];
            unsafe {
                aes256_cfb_encrypt(&mut ciphertext_custom, &plaintext, &iv, &round_key);
            }

            let ciphertext_openssl = openssl_aes256_cfb128_encrypt(&plaintext, &key, &iv).unwrap();
            assert_eq!(
                ciphertext_custom, ciphertext_openssl,
                "Mismatch for length {}",
                len
            );

            let mut decrypted_custom = vec![0u8; *len];
            unsafe {
                aes256_cfb_decrypt(&mut decrypted_custom, &ciphertext_custom, &iv, &round_key);
            }
            assert_eq!(
                decrypted_custom, plaintext,
                "Decryption mismatch for length {}",
                len
            );
        }
        println!("AES-256-CFB various lengths test passed!");
    }
 }
	// benches/aes_cfb_bench.rs

	use aes_simd_rust::*; // Your crate name
	use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion, Throughput};
	use openssl::symm::{Cipher, Crypter, Mode};
	use rand::rngs::StdRng;
	use rand::{RngCore, SeedableRng};
	use std::hint::black_box;
	use std::mem;

	// Re-define M128i if not public or use crate::M128i
	#[cfg(target_arch = "x86")]
	use std::arch::x86::__m128i as M128iBench;
	#[cfg(target_arch = "x86_64")]
	use std::arch::x86_64::__m128i as M128iBench;

	fn openssl_encrypt_routine(
	plaintext: &[u8],
	key: &[u8; 32],
	iv: &[u8; 16],
	output_buffer: &mut Vec<u8>, // Re-use buffer
	) {
	output_buffer.clear();
	output_buffer.resize(
	plaintext.len() + Cipher::aes_256_cfb128().block_size() - 1,
	0,
	);

	let cipher = Cipher::aes_256_cfb128();
	let mut crypter = Crypter::new(cipher, Mode::Encrypt, key, Some(iv)).unwrap();
	crypter.pad(false);

	let mut count = crypter.update(plaintext, output_buffer).unwrap();
	count += crypter.finalize(&mut output_buffer[count..]).unwrap();
	output_buffer.truncate(count);
	}

	fn bench_aes_cfb_encryption(c: &mut Criterion) {
	if !check_aesni_support() {
	println!("Skipping AES-NI benchmark: CPU does not support required features.");
	// Fallback or skip: For now, we just print and skip.
	// Alternatively, one could implement a non-SIMD version for comparison
	// or error out if AES-NI is essential.
	return;
	}

	let mut rng = StdRng::seed_from_u64(12345); // Consistent seed for reproducible benchmarks
	let mut key = [0u8; 32];
	let mut iv = [0u8; 16];
	rng.fill_bytes(&mut key);
	rng.fill_bytes(&mut iv);

	// Pre-compute round keys for custom SIMD version
	let mut round_key_custom: [M128iBench; 15] = unsafe { mem::zeroed() };
	unsafe {
	// Assuming aes256_key_expansion is in the crate root or imported correctly
	aes_simd_rust::aes256_key_expansion(&key, &mut round_key_custom);
	}

	let mut group = c.benchmark_group("AES-256-CFB Encryption");

	// Test with various data sizes
	for size in [64, 256, 1024, 4096, 16384, 65536].iter() {
	let mut plaintext = vec![0u8; *size];
	rng.fill_bytes(&mut plaintext);

	// Buffer for custom SIMD output
	let mut ciphertext_custom_out = vec![0u8; *size];
	// Buffer for OpenSSL output (re-used)
	let mut ciphertext_openssl_out = Vec::with_capacity(*size + 32);

	group.throughput(Throughput::Bytes(*size as u64));

	group.bench_with_input(BenchmarkId::new("Custom SIMD", size), &plaintext, \|b, p\| {
	b.iter(\|\| unsafe {
	// Assuming aes256_cfb_encrypt is in the crate root or imported correctly
	aes_simd_rust::aes256_cfb_encrypt(
	black_box(&mut ciphertext_custom_out),
	black_box(p),
	black_box(&iv),
	black_box(&round_key_custom),
	)
	});
	});

	group.bench_with_input(BenchmarkId::new("OpenSSL", size), &plaintext, \|b, p\| {
	b.iter(\|\| {
	openssl_encrypt_routine(
	black_box(p),
	black_box(&key),
	black_box(&iv),
	black_box(&mut ciphertext_openssl_out),
	)
	});
	});
	}
	group.finish();
	}

	criterion_group!(benches, bench_aes_cfb_encryption);
	criterion_main!(benches);
	[package]
	name = "aes_simd_rust"
	version = "0.1.0"
	edition = "2021"

	[dependencies]
	openssl = "*"
	rand = "0.9" # For generating test data

	[dev-dependencies]
	criterion = { version = "0.6", features = ["html_reports"] }

	[[bench]]
	name = "aes_cfb_bench"
	harness = false
	#![feature(portable_simd)] // For some intrinsics, though many are stable.
	// Not strictly needed if relying on std::arch and target_feature.

	// Import necessary SIMD intrinsics
	#[cfg(target_arch = "x86")]
	use std::arch::x86::*;
	#[cfg(target_arch = "x86_64")]
	use std::arch::x86_64::*;

	use std::mem;

	// Type alias for clarity, similar to C++
	pub type M128i = __m128i;

	// Helper macro for AES-256 key expansion (from Intel's whitepaper/common implementations)
	// This macro assists in generating one 128-bit word of the key schedule.
	// It corresponds to the sequence:
	// temp = _mm_aeskeygenassist_si128(key2, rcon);
	// temp = _mm_shuffle_epi32(temp, 0xFF); // or 0xAA for the other part
	// key1 = _mm_xor_si128(key1, _mm_slli_si128(key1, 4));
	// key1 = _mm_xor_si128(key1, _mm_slli_si128(key1, 8));
	// key1 = _mm_xor_si128(key1, temp);
	#[macro_export]
	macro_rules! aes256_assist {
	($key1:expr, $key2:expr, $rcon:expr, $shuffle_mask:expr) => {{
	let mut temp = _mm_aeskeygenassist_si128($key2, $rcon);
	temp = _mm_shuffle_epi32(temp, $shuffle_mask);
	let mut k1_mod = $key1;
	k1_mod = _mm_xor_si128(k1_mod, _mm_slli_si128(k1_mod, 4));
	k1_mod = _mm_xor_si128(k1_mod, _mm_slli_si128(k1_mod, 8));
	// k1_mod = _mm_xor_si128(k1_mod, _mm_slli_si128(k1_mod, 12)); // This is not in the C++ code, it uses two 4/8 shifts
	k1_mod = _mm_xor_si128(k1_mod, temp);
	k1_mod
	}};
	}

	/// AES-256 Key Expansion
	///
	/// Expands a 32-byte key into 15 128-bit round keys.
	///
	/// # Safety
	///
	/// This function is unsafe because it uses CPU intrinsics that require specific CPU features (AES-NI, SSE2).
	/// The caller must ensure these features are available.
	/// `key` must point to 32 valid bytes.
	/// `round_key` must be a mutable slice of 15 `M128i` elements.
	#[target_feature(enable = "aes,sse2")]
	pub unsafe fn aes256_key_expansion(key: &[u8; 32], round_key: &mut [M128i; 15]) {
	let mut k1 = _mm_loadu_si128(key.as_ptr() as *const M128i);
	let mut k2 = _mm_loadu_si128(key.as_ptr().add(16) as *const M128i);

	round_key[0] = k1;
	round_key[1] = k2;

	// RCON values for AES-256 (only 7 are needed for the 14 rounds of key expansion)
	// The C++ code uses specific RCONs for _mm_aeskeygenassist_si128
	// 0x01, 0x00, 0x02, 0x00, 0x04, 0x00, 0x08, 0x00, 0x10, 0x00, 0x20, 0x00, 0x40
	// The shuffle masks are 0xFF and 0xAA

	// Round 1 (generates round_key[2])
	k1 = aes256_assist!(k1, k2, 0x01, 0xff); // shuffle FF for [3,3,3,3]
	round_key[2] = k1;

	// Round 2 (generates round_key[3])
	k2 = aes256_assist!(k2, k1, 0x00, 0xaa); // shuffle AA for [2,2,2,2]
	round_key[3] = k2;

	// Round 3 (generates round_key[4])
	k1 = aes256_assist!(k1, k2, 0x02, 0xff);
	round_key[4] = k1;

	// Round 4 (generates round_key[5])
	k2 = aes256_assist!(k2, k1, 0x00, 0xaa);
	round_key[5] = k2;

	// Round 5 (generates round_key[6])
	k1 = aes256_assist!(k1, k2, 0x04, 0xff);
	round_key[6] = k1;

	// Round 6 (generates round_key[7])
	k2 = aes256_assist!(k2, k1, 0x00, 0xaa);
	round_key[7] = k2;

	// Round 7 (generates round_key[8])
	k1 = aes256_assist!(k1, k2, 0x08, 0xff);
	round_key[8] = k1;

	// Round 8 (generates round_key[9])
	k2 = aes256_assist!(k2, k1, 0x00, 0xaa);
	round_key[9] = k2;

	// Round 9 (generates round_key[10])
	k1 = aes256_assist!(k1, k2, 0x10, 0xff);
	round_key[10] = k1;

	// Round 10 (generates round_key[11])
	k2 = aes256_assist!(k2, k1, 0x00, 0xaa);
	round_key[11] = k2;

	// Round 11 (generates round_key[12])
	k1 = aes256_assist!(k1, k2, 0x20, 0xff);
	round_key[12] = k1;

	// Round 12 (generates round_key[13])
	k2 = aes256_assist!(k2, k1, 0x00, 0xaa);
	round_key[13] = k2;

	// Round 13 (generates round_key[14])
	k1 = aes256_assist!(k1, k2, 0x40, 0xff);
	round_key[14] = k1;
	}

	/// AES-256 Encrypt a single 128-bit block
	///
	/// # Safety
	///
	/// This function is unsafe because it uses CPU intrinsics that require specific CPU features (AES-NI).
	/// The caller must ensure these features are available.
	/// `round_key` must contain 15 valid round keys.
	#[target_feature(enable = "aes")]
	pub unsafe fn aes256_encrypt_block(mut block: M128i, round_key: &[M128i; 15]) -> M128i {
	block = _mm_xor_si128(block, round_key[0]); // Initial AddRoundKey

	// 13 full rounds
	for i in 1..14 {
	block = _mm_aesenc_si128(block, round_key[i]);
	}

	// Final round (no MixColumns)
	block = _mm_aesenclast_si128(block, round_key[14]);
	block
	}

	/// AES-256-CFB Encryption
	///
	/// Encrypts `plaintext` into `ciphertext` using CFB mode.
	/// `plaintext` and `ciphertext` must have the same length.
	/// `iv` must be 16 bytes.
	/// `round_key` must be the expanded AES-256 key schedule (15 round keys).
	///
	/// # Safety
	///
	/// This function is unsafe due to the use of SIMD intrinsics and pointer operations.
	/// Caller must ensure CPU features (AES-NI, SSE2) are available.
	/// Input slices must be valid and correctly sized.
	#[target_feature(enable = "aes,sse2")]
	pub unsafe fn aes256_cfb_encrypt(
	ciphertext: &mut [u8],
	plaintext: &[u8],
	iv: &[u8; 16],
	round_key: &[M128i; 15],
	) {
	assert_eq!(
	plaintext.len(),
	ciphertext.len(),
	"Plaintext and ciphertext buffers must have the same length."
	);
	let len = plaintext.len();
	if len == 0 {
	return;
	}

	let mut feedback = _mm_loadu_si128(iv.as_ptr() as *const M128i);

	let mut processed_bytes = 0;
	let block_size = 16;

	// Process full blocks
	while processed_bytes + block_size <= len {
	let keystream = aes256_encrypt_block(feedback, round_key);

	let plain_block_ptr = plaintext.as_ptr().add(processed_bytes) as *const M128i;
	let plain_block = _mm_loadu_si128(plain_block_ptr);

	let cipher_block = _mm_xor_si128(plain_block, keystream);

	let cipher_out_ptr = ciphertext.as_mut_ptr().add(processed_bytes) as *mut M128i;
	_mm_storeu_si128(cipher_out_ptr, cipher_block);

	feedback = cipher_block; // Next feedback is the current ciphertext block
	processed_bytes += block_size;
	}

	// Process remaining partial block
	if processed_bytes < len {
	let keystream = aes256_encrypt_block(feedback, round_key);
	let keystream_bytes: [u8; 16] = mem::transmute(keystream); // Get bytes from M128i

	for i in 0..(len - processed_bytes) {
	ciphertext[processed_bytes + i] = plaintext[processed_bytes + i] ^ keystream_bytes[i];
	}
	}
	}

	/// AES-256-CFB Decryption
	///
	/// Decrypts `ciphertext` into `plaintext` using CFB mode.
	/// `plaintext` and `ciphertext` must have the same length.
	/// `iv` must be 16 bytes.
	/// `round_key` must be the expanded AES-256 key schedule (15 round keys).
	///
	/// # Safety
	///
	/// This function is unsafe due to the use of SIMD intrinsics and pointer operations.
	/// Caller must ensure CPU features (AES-NI, SSE2) are available.
	/// Input slices must be valid and correctly sized.
	#[target_feature(enable = "aes,sse2")]
	pub unsafe fn aes256_cfb_decrypt(
	plaintext: &mut [u8],
	ciphertext: &[u8],
	iv: &[u8; 16],
	round_key: &[M128i; 15],
	) {
	assert_eq!(
	plaintext.len(),
	ciphertext.len(),
	"Plaintext and ciphertext buffers must have the same length."
	);
	let len = ciphertext.len();
	if len == 0 {
	return;
	}

	let mut feedback = _mm_loadu_si128(iv.as_ptr() as *const M128i);

	let mut processed_bytes = 0;
	let block_size = 16;

	// Process full blocks
	while processed_bytes + block_size <= len {
	let keystream = aes256_encrypt_block(feedback, round_key);

	let cipher_block_ptr = ciphertext.as_ptr().add(processed_bytes) as *const M128i;
	let cipher_block = _mm_loadu_si128(cipher_block_ptr);

	let plain_block = _mm_xor_si128(cipher_block, keystream);

	let plain_out_ptr = plaintext.as_mut_ptr().add(processed_bytes) as *mut M128i;
	_mm_storeu_si128(plain_out_ptr, plain_block);

	feedback = cipher_block; // Next feedback is the current ciphertext block
	processed_bytes += block_size;
	}

	// Process remaining partial block
	if processed_bytes < len {
	let keystream = aes256_encrypt_block(feedback, round_key);
	let keystream_bytes: [u8; 16] = mem::transmute(keystream); // Get bytes from M128i

	for i in 0..(len - processed_bytes) {
	plaintext[processed_bytes + i] = ciphertext[processed_bytes + i] ^ keystream_bytes[i];
	}
	}
	}

	// Helper for tests and benchmarks
	pub fn check_aesni_support() -> bool {
	is_x86_feature_detected!("aes") && is_x86_feature_detected!("sse2")
	}

	#[cfg(test)]
	mod tests {
	use super::*;
	use openssl::symm::{Cipher, Crypter, Mode};
	use rand::rngs::StdRng;
	use rand::{RngCore, SeedableRng};

	fn print_hex(label: &str, data: &[u8]) {
	print!("{} ({} bytes): ", label, data.len());
	for byte in data.iter().take(32) {
	// Print first 32 bytes
	print!("{:02x}", byte);
	}
	if data.len() > 32 {
	print!("...");
	}
	println!();
	}

	fn openssl_aes256_cfb128_encrypt(
	plaintext: &[u8],
	key: &[u8; 32],
	iv: &[u8; 16],
	) -> Result<Vec<u8>, openssl::error::ErrorStack> {
	let cipher = Cipher::aes_256_cfb128();
	let mut crypter = Crypter::new(cipher, Mode::Encrypt, key, Some(iv))?;
	crypter.pad(false); // CFB mode typically doesn't use padding

	let mut ciphertext = vec![0; plaintext.len() + cipher.block_size() - 1]; // Max possible size for update
	let mut count = crypter.update(plaintext, &mut ciphertext)?;
	count += crypter.finalize(&mut ciphertext[count..])?;
	ciphertext.truncate(count);
	Ok(ciphertext)
	}

	fn openssl_aes256_cfb128_decrypt(
	ciphertext: &[u8],
	key: &[u8; 32],
	iv: &[u8; 16],
	) -> Result<Vec<u8>, openssl::error::ErrorStack> {
	let cipher = Cipher::aes_256_cfb128();
	let mut crypter = Crypter::new(cipher, Mode::Decrypt, key, Some(iv))?;
	crypter.pad(false);

	let mut plaintext = vec![0; ciphertext.len() + cipher.block_size() - 1];
	let mut count = crypter.update(ciphertext, &mut plaintext)?;
	count += crypter.finalize(&mut plaintext[count..])?;
	plaintext.truncate(count);
	Ok(plaintext)
	}

	#[test]
	fn test_aes256_cfb_correctness() {
	if !check_aesni_support() {
	println!("Skipping AES-NI test: CPU does not support required features.");
	return;
	}

	let key: [u8; 32] = [
	0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0a, 0x0b, 0x0c, 0x0d,
	0x0e, 0x0f, 0x10, 0x11, 0x12, 0x13, 0x14, 0x15, 0x16, 0x17, 0x18, 0x19, 0x1a, 0x1b,
	0x1c, 0x1d, 0x1e, 0x1f,
	];

	let iv: [u8; 16] = [
	0x00, 0x11, 0x22, 0x33, 0x44, 0x55, 0x66, 0x77, 0x88, 0x99, 0xaa, 0xbb, 0xcc, 0xdd,
	0xee, 0xff,
	];

	let plaintext_str = "Hello, AES-256-CFB! This is a test message for encryption and decryption. \
	This additional text is to make the test data longer for more meaningful performance testing. \
	AES (Advanced Encryption Standard) is a specification for the encryption of electronic data \
	established by the U.S. National Institute of Standards and Technology (NIST) in 2001.";
	let plaintext = plaintext_str.as_bytes();
	let len = plaintext.len();

	// Initialize round keys (unsafe block as it calls target_feature enabled function)
	let mut round_key: [M128i; 15] = unsafe { mem::zeroed() };
	unsafe {
	aes256_key_expansion(&key, &mut round_key);
	}

	// Custom SIMD Encryption
	let mut ciphertext_custom = vec![0u8; len];
	unsafe {
	aes256_cfb_encrypt(&mut ciphertext_custom, plaintext, &iv, &round_key);
	}

	// OpenSSL Encryption
	let ciphertext_openssl = openssl_aes256_cfb128_encrypt(plaintext, &key, &iv).unwrap();

	print_hex("Plaintext ", plaintext);
	print_hex("Custom Ciphertext", &ciphertext_custom);
	print_hex("OpenSSL Ciphertext", &ciphertext_openssl);

	assert_eq!(
	ciphertext_custom.len(),
	ciphertext_openssl.len(),
	"Ciphertext lengths differ"
	);
	assert_eq!(
	ciphertext_custom, ciphertext_openssl,
	"Encryption results differ between custom and OpenSSL"
	);

	// Custom SIMD Decryption
	let mut decrypted_custom = vec![0u8; len];
	unsafe {
	aes256_cfb_decrypt(&mut decrypted_custom, &ciphertext_custom, &iv, &round_key);
	}
	print_hex("Custom Decrypted ", &decrypted_custom);
	assert_eq!(
	decrypted_custom, plaintext,
	"Custom decryption failed to recover original plaintext"
	);

	// OpenSSL Decryption (of OpenSSL's ciphertext for sanity)
	let decrypted_openssl =
	openssl_aes256_cfb128_decrypt(&ciphertext_openssl, &key, &iv).unwrap();
	assert_eq!(decrypted_openssl, plaintext, "OpenSSL decryption failed");

	println!("AES-256-CFB correctness test passed!");
	}

	#[test]
	fn test_various_lengths() {
	if !check_aesni_support() {
	println!("Skipping AES-NI length test: CPU does not support required features.");
	return;
	}
	let mut rng = StdRng::seed_from_u64(42);
	let mut key = [0u8; 32];
	let mut iv = [0u8; 16];
	rng.fill_bytes(&mut key);
	rng.fill_bytes(&mut iv);

	let mut round_key: [M128i; 15] = unsafe { mem::zeroed() };
	unsafe {
	aes256_key_expansion(&key, &mut round_key);
	}

	for len in [0, 1, 15, 16, 17, 31, 32, 33, 100, 1024].iter() {
	let mut plaintext = vec![0u8; *len];
	rng.fill_bytes(&mut plaintext);

	let mut ciphertext_custom = vec![0u8; *len];
	unsafe {
	aes256_cfb_encrypt(&mut ciphertext_custom, &plaintext, &iv, &round_key);
	}

	let ciphertext_openssl = openssl_aes256_cfb128_encrypt(&plaintext, &key, &iv).unwrap();
	assert_eq!(
	ciphertext_custom, ciphertext_openssl,
	"Mismatch for length {}",
	len
	);

	let mut decrypted_custom = vec![0u8; *len];
	unsafe {
	aes256_cfb_decrypt(&mut decrypted_custom, &ciphertext_custom, &iv, &round_key);
	}
	assert_eq!(
	decrypted_custom, plaintext,
	"Decryption mismatch for length {}",
	len
	);
	}
	println!("AES-256-CFB various lengths test passed!");
	}
	}