chacha: Add zeroize feature and runtime detection of CPU features

Author: bloom

.gitignore

@@ -48,3 +48,5 @@ node_modules/
 
 # Rust
 target/
+
+#chacha12/Cargo.lock

Cargo.lock

@@ -47,7 +47,7 @@ chacha12-blake3 = "0.9"
 use chacha12_blake3::ChaCha12Blake3;
 
 fn main() {
-    // OD NOT USE A ALL-ZERO KEY / NONCE, THIS CODE IS FOR DEMONSTRATION ONLY
+    // DO NOT USE A ALL-ZERO KEY / NONCE, THIS CODE IS FOR DEMONSTRATION ONLY
     let key = [0u8; 32];
     let nonce = [0u8; 32];
     // or with an u64 counter to encrypt up to 2^64 messages with a single key:

README.md

@@ -3,7 +3,7 @@ use chacha20::{ChaCha12, KeyIvInit, cipher::StreamCipher};
 use criterion::*;
 
 fn bench(c: &mut Criterion) {
-    for n in [64, 2000, 64 * 1000, 1000 * 1000, 10 * 1000 * 1000] {
+    for n in [64, 2048, 64 * 1024, 1024 * 1024, 10 * 1024 * 1024] {
         let mut group = c.benchmark_group(format!("{}", n));
         let mut plaintext = vec![0u8; n];
 

benchmarks/benches/chacha.rs

@@ -4,7 +4,7 @@ use chacha20poly1305::{ChaCha20Poly1305, KeyInit};
 use criterion::*;
 
 fn bench(c: &mut Criterion) {
-    for n in [64, 2000, 64 * 1000, 1000 * 1000, 10 * 1000 * 1000] {
+    for n in [64, 2048, 64 * 1024, 1024 * 1024, 10 * 1024 * 1024] {
         let mut group = c.benchmark_group(format!("{}", n));
         let mut plaintext = vec![0u8; n];
 

benchmarks/benches/encrypt.rs

@@ -5,10 +5,21 @@ edition = "2024"
 description = "Pure-Rust, SIMD-accelerated ChaCha20 / ChaCha12 / ChaCha8 for any platform"
 repository = "https://github.com/bloom42/chacha12-blake3"
 license = "MIT"
-keywords = ["cryptography", "chacha20", "blake3", "encryption", "chacha12", "simd", "wasm", "chacha", "webassembly"]
-categories = ["cryptography", "wasm"]
+keywords = ["cryptography", "chacha20", "encryption", "chacha12", "simd", "wasm", "chacha", "webassembly"]
+categories = ["cryptography", "wasm", "simd"]
+
+[features]
+default = ["std", "zeroize"]
+
+# enable runtime detection of CPU features (SIMD instructions)
+std = []
+
+# erase sensitive secrets from memory
+zeroize = ["dep:zeroize"]
 
 [dependencies]
+zeroize = { version = "1", features = ["simd", "derive"], optional = true }
+
 
 [dev-dependencies]
 hex = "0.4"

chacha12/Cargo.lock

@@ -1,5 +1,21 @@
-# ChaCha12
+# ChaCha
 
-Pure-Rust, SIMD-accelerated ChaCha12 for any platform.
+Pure-Rust, SIMD-accelerated ChaCha20 / ChaCha12 / ChaCha8 for any platform.
 
-Coming soon.
+> **⚠️ Warning ⚠️:** This is a preliminary release, DO NOT USE IN PRODUCTION.
+
+
+```rust
+use chacha12::ChaCha;
+
+fn main() {
+    // DO NOT USE A ALL-ZERO KEY / NONCE, THIS CODE IS FOR DEMONSTRATION ONLY
+    let key = [0u8; 32];
+    let nonce = [0u8; 8];
+
+    let mut message = b"Hello World!".to_vec();
+
+    let mut cipher = ChaCha::<12>::new(&key, &nonce);
+    cipher.xor_keystream(&mut message);
+}
+```

chacha12/Cargo.toml

@@ -1,20 +1,34 @@
+#![doc = include_str!(concat!(env!("CARGO_MANIFEST_DIR"), "/README.md"))]
+
 // aarch64 assumes that NEON instructions are always present
 #[cfg(target_arch = "aarch64")]
 use crate::chacha_neon::chacha_neon;
 
 #[cfg(target_arch = "aarch64")]
 mod chacha_neon;
 
-#[cfg(all(any(target_arch = "x86", target_arch = "x86_64"), target_feature = "avx2"))]
+#[cfg(any(
+    all(target_arch = "x86_64", feature = "std"),
+    all(target_arch = "x86_64", target_feature = "avx2")
+))]
 mod chacha_avx2;
 
-#[cfg(all(any(target_arch = "x86", target_arch = "x86_64"), target_feature = "avx2"))]
+#[cfg(any(
+    all(target_arch = "x86_64", feature = "std"),
+    all(target_arch = "x86_64", target_feature = "avx2")
+))]
 use chacha_avx2::chacha_avx2;
 
-#[cfg(all(target_arch = "x86_64", target_feature = "avx512f"))]
+#[cfg(any(
+    all(target_arch = "x86_64", feature = "std"),
+    all(target_arch = "x86_64", target_feature = "avx512f")
+))]
 mod chacha_avx512;
 
-#[cfg(all(target_arch = "x86_64", target_feature = "avx512f"))]
+#[cfg(any(
+    all(target_arch = "x86_64", feature = "std"),
+    all(target_arch = "x86_64", target_feature = "avx512f")
+))]
 use chacha_avx512::chacha_avx512;
 
 #[cfg(all(target_arch = "wasm32", target_feature = "simd128"))]
@@ -23,6 +37,9 @@ mod chacha_wasm_simd128;
 #[cfg(all(target_arch = "wasm32", target_feature = "simd128"))]
 use chacha_wasm_simd128::chacha_wasm_simd128;
 
+#[cfg(feature = "zeroize")]
+use zeroize::{Zeroize, ZeroizeOnDrop};
+
 const CONSTANT: [u32; 4] = [
     0x61707865, // "expa"
     0x3320646e, // "nd 3"
@@ -33,6 +50,7 @@ const CONSTANT: [u32; 4] = [
 /// The number of 32-bit words that compose ChaCha's state
 const STATE_WORDS: usize = 16;
 
+#[cfg_attr(feature = "zeroize", derive(Zeroize, ZeroizeOnDrop))]
 pub struct ChaCha<const ROUNDS: usize> {
     state: [u32; STATE_WORDS],
     counter: u64,
@@ -47,7 +65,8 @@ pub struct ChaCha<const ROUNDS: usize> {
     /// Then, when calling `xor_keystream` again, we first check if there is sone leftover form the last
     /// keystream.
     /// NOTE: the `last_keystream_block` is valid only if the previous call to `xor_keystream` had
-    /// an input.len() % 64 != 0
+    /// an input.len() % 64 != 0.
+    /// Otherwise there is no need to preserve the last keystream block.
     last_keystream_block: [u8; 64],
     last_keystream_block_index: usize,
 }
@@ -68,19 +87,21 @@ impl<const ROUNDS: usize> ChaCha<ROUNDS> {
         state[14] = u32::from_le_bytes(nonce[0..4].try_into().unwrap());
         state[15] = u32::from_le_bytes(nonce[4..8].try_into().unwrap());
 
-        ChaCha {
+        return ChaCha {
             state,
             counter: 0,
             last_keystream_block: [0u8; 64],
             last_keystream_block_index: 0,
-        }
+        };
     }
 
+    /// XOR `plaintext` with the ChaCha keystream.
     pub fn xor_keystream(&mut self, mut plaintext: &mut [u8]) {
         if plaintext.len() == 0 {
             return;
         }
 
+        // first, consume the keystream leftover, if any
         if self.last_keystream_block_index != 0 {
             let remaining_keystream = &self.last_keystream_block[self.last_keystream_block_index..];
 
@@ -102,36 +123,63 @@ impl<const ROUNDS: usize> ChaCha<ROUNDS> {
         }
         self.last_keystream_block_index = plaintext.len() % 64;
 
+        // aarch64 assumes that NEON is always available
         #[cfg(target_arch = "aarch64")]
         if plaintext.len() >= 128 {
             self.counter = chacha_neon::<ROUNDS>(self.state, self.counter, plaintext, &mut self.last_keystream_block);
             return;
         }
 
-        #[cfg(all(target_arch = "x86_64", target_feature = "avx512f"))]
+        #[cfg(all(target_arch = "wasm32", target_feature = "simd128"))]
         if plaintext.len() >= 128 {
-            self.counter = chacha_avx512::<ROUNDS>(self.state, self.counter, plaintext, &mut self.last_keystream_block);
+            self.counter =
+                chacha_wasm_simd128::<ROUNDS>(self.state, self.counter, plaintext, &mut self.last_keystream_block);
             return;
         }
 
-        #[cfg(all(any(target_arch = "x86", target_arch = "x86_64"), target_feature = "avx2"))]
+        // runtime detection of CPU features for x86 and x86_64 when the "std" feature is enabled
+        #[cfg(feature = "std")]
+        {
+            #[cfg(target_arch = "x86_64")]
+            if is_x86_feature_detected!("avx512f") {
+                self.counter =
+                    chacha_avx512::<ROUNDS>(self.state, self.counter, plaintext, &mut self.last_keystream_block);
+                return;
+            }
+
+            #[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
+            if is_x86_feature_detected!("avx2") {
+                self.counter =
+                    chacha_avx2::<ROUNDS>(self.state, self.counter, plaintext, &mut self.last_keystream_block);
+                return;
+            }
+        }
+
+        // compile-time CPU detection
+        #[cfg(all(target_arch = "x86_64", target_feature = "avx512f", not(feature = "std")))]
         if plaintext.len() >= 128 {
-            self.counter = chacha_avx2::<ROUNDS>(self.state, self.counter, plaintext, &mut self.last_keystream_block);
+            self.counter = chacha_avx512::<ROUNDS>(self.state, self.counter, plaintext, &mut self.last_keystream_block);
             return;
         }
 
-        #[cfg(all(target_arch = "wasm32", target_feature = "simd128"))]
+        #[cfg(all(
+            any(target_arch = "x86", target_arch = "x86_64"),
+            target_feature = "avx2",
+            not(feature = "std")
+        ))]
         if plaintext.len() >= 128 {
-            self.counter =
-                chacha_wasm_simd128::<ROUNDS>(self.state, self.counter, plaintext, &mut self.last_keystream_block);
+            self.counter = chacha_avx2::<ROUNDS>(self.state, self.counter, plaintext, &mut self.last_keystream_block);
             return;
         }
 
         self.counter = chacha_generic::<ROUNDS>(self.state, self.counter, plaintext, &mut self.last_keystream_block);
     }
 
+    /// Set the ChaCha counter (words 12 and 13). It can be used to move forward and backward in the
+    /// keystream.
     pub fn set_counter(&mut self, counter: u64) {
         self.counter = counter;
+        // setting the counter "realigns" the keystream to the beginning of a block.
         self.last_keystream_block_index = 0;
     }
 }
@@ -197,7 +245,9 @@ fn chacha_generic<const ROUNDS: usize>(
         counter = counter.wrapping_add(1);
     }
 
-    last_keystream_block.copy_from_slice(&keystream);
+    if plaintext.len() % 64 != 0 {
+        last_keystream_block.copy_from_slice(&keystream);
+    }
 
     return counter;
 }