Skip to main content

crypto/chacha/
chacha.rs

1#[cfg(feature = "zeroize")]
2use zeroize::{Zeroize, ZeroizeOnDrop};
3
4use crate::StreamCipher;
5
6/// The number of 32-bit words that compose ChaCha's state.
7pub(crate) const STATE_WORDS: usize = 16;
8
9/// The size of a ChaCha block in bytes which is the size of the state in bytes
10pub(crate) const BLOCK_SIZE: usize = 64;
11
12/// Maximum number of bytes that may be encrypted under a single `(key, nonce)` pair with the
13/// IETF (RFC 8439) layout.
14///
15/// RFC 8439 Section 2.8 caps the IETF variant at `2^32 - 1` blocks of 64 bytes because the block
16/// counter is 32 bits wide. This constant keeps one block of margin (`2^32 - 2` blocks) so the
17/// counter can never wrap and reuse keystream.
18///
19/// Exceeding this limit makes the ChaCha20-Poly1305 / XChaCha20-Poly1305 / ChaCha8-Poly1305 AEADs
20/// panic on encryption and return [`AeadError::InvalidCiphertext`](crate::AeadError) on decryption,
21/// and makes [`StreamCipher::xor_keystream`](crate::StreamCipher::xor_keystream) panic.
22pub const CHACHA20_IETF_MAX_LEN: u64 = (u32::MAX as u64 - 1) * BLOCK_SIZE as u64;
23
24/// The "sigma" constant which is the value of the first row of ChaCha's state.
25pub(crate) const CONSTANT: [u32; 4] = [
26    0x61707865, // "expa"
27    0x3320646e, // "nd 3"
28    0x79622d32, // "2-by"
29    0x6b206574, // "te k"
30];
31
32pub type ChaCha8Djb = ChaCha<8, false>;
33pub type ChaCha12Djb = ChaCha<12, false>;
34pub type ChaCha20Djb = ChaCha<20, false>;
35pub type ChaCha20Ietf = ChaCha<20, true>;
36pub type XChaCha20 = XChaCha<20>;
37
38/// ChaCha stream cipher.
39///
40/// `IS_IETF` selects the nonce/counter layout:
41/// - `false` (DJB original): 64-bit counter at words 12–13, 64-bit nonce at words 14–15.
42/// - `true` (IETF / RFC 8439): 32-bit counter at word 12, 96-bit nonce at words 13–15.
43///
44/// The IETF layout can only encrypt up to [`CHACHA20_IETF_MAX_LEN`] bytes per `(key, nonce)` pair;
45/// [`xor_keystream`](StreamCipher::xor_keystream) panics instead of wrapping the counter.
46#[cfg_attr(feature = "zeroize", derive(Zeroize, ZeroizeOnDrop))]
47pub struct ChaCha<const ROUNDS: usize, const IS_IETF: bool> {
48    state: [u32; STATE_WORDS],
49    /// ChaCha is a stream cipher that works with 64-byte blocks.
50    /// It means that consumers of this packages should be able to call `xor_keystream` multiple
51    /// times even if there input is not aligned with ChaCha blocks.
52    /// Thus calling multiple times `xor_keystream`:
53    /// xor_keystream(plaintext[0..3]), xor_keystream(plaintext[3..50]), xor_keystream(plaintext[50..150]);
54    /// Should be equal to calling it only once:
55    /// xor_keystream(plaintext[0..150]);
56    /// For that, we keep the last computed keystream block, as well as an offset indicating where
57    /// the unconsumed tail starts.
58    /// The full leftover block is stored in `keystream_leftover` minus 1 byte because if there is leftover
59    /// it means that the leftover is <= (BLOCK_SIZE - 1).
60    /// When `keystream_leftover_offset == (BLOCK_SIZE - 1)`, there is no leftover (empty slice).
61    /// NOTE: the `keystream_leftover` buffer is valid only if the previous call to `xor_keystream` had
62    /// an `input.len() % 64 != 0`, Otherwise there is no need to preserve the last keystream block.
63    keystream_leftover: [u8; BLOCK_SIZE - 1],
64    keystream_leftover_offset: u8,
65}
66
67impl<const ROUNDS: usize, const IS_IETF: bool> ChaCha<ROUNDS, IS_IETF> {
68    #[inline(always)]
69    fn extract_counter(state: &[u32; STATE_WORDS]) -> u64 {
70        if IS_IETF {
71            state[12] as u64
72        } else {
73            ((state[13] as u64) << 32) | (state[12] as u64)
74        }
75    }
76
77    #[inline(always)]
78    fn inject_counter(state: &mut [u32; STATE_WORDS], counter: u64) {
79        state[12] = counter as u32;
80        if !IS_IETF {
81            state[13] = (counter >> 32) as u32;
82        }
83    }
84}
85
86impl<const ROUNDS: usize> ChaCha<ROUNDS, false> {
87    /// Create a new ChaCha instance with the DJB nonce layout (8-byte nonce, 64-bit counter).
88    pub fn new(key: &[u8; 32], nonce: &[u8; 8]) -> ChaCha<ROUNDS, false> {
89        let mut state = [0u32; STATE_WORDS];
90
91        state[..4].copy_from_slice(&CONSTANT);
92
93        for (state_word, key_chunk) in state[4..12].iter_mut().zip(key.as_chunks::<4>().0) {
94            *state_word = u32::from_le_bytes(*key_chunk);
95        }
96
97        state[14] = u32::from_le_bytes(nonce[0..4].try_into().unwrap());
98        state[15] = u32::from_le_bytes(nonce[4..8].try_into().unwrap());
99
100        return ChaCha {
101            state,
102            keystream_leftover: [0u8; BLOCK_SIZE - 1],
103            keystream_leftover_offset: (BLOCK_SIZE - 1) as u8,
104        };
105    }
106
107    /// Set the ChaCha counter (words 12 and 13). It can be used to move forward and backward in the
108    /// keystream.
109    #[inline(always)]
110    pub fn set_counter(&mut self, counter: u64) {
111        Self::inject_counter(&mut self.state, counter);
112        self.keystream_leftover_offset = (BLOCK_SIZE - 1) as u8;
113    }
114}
115
116impl<const ROUNDS: usize> ChaCha<ROUNDS, true> {
117    /// Create a new ChaCha instance with the IETF (RFC 8439) nonce layout (12-byte nonce, 32-bit counter).
118    pub fn new(key: &[u8; 32], nonce: &[u8; 12]) -> ChaCha<ROUNDS, true> {
119        let mut state = [0u32; STATE_WORDS];
120
121        state[..4].copy_from_slice(&CONSTANT);
122
123        for (state_word, key_chunk) in state[4..12].iter_mut().zip(key.as_chunks::<4>().0) {
124            *state_word = u32::from_le_bytes(*key_chunk);
125        }
126
127        state[13] = u32::from_le_bytes(nonce[0..4].try_into().unwrap());
128        state[14] = u32::from_le_bytes(nonce[4..8].try_into().unwrap());
129        state[15] = u32::from_le_bytes(nonce[8..12].try_into().unwrap());
130
131        return ChaCha {
132            state,
133            keystream_leftover: [0u8; BLOCK_SIZE - 1],
134            keystream_leftover_offset: (BLOCK_SIZE - 1) as u8,
135        };
136    }
137
138    /// Set the ChaCha counter (word 12). The counter is a u32. It can be used to move forward
139    /// and backward in the keystream.
140    ///
141    /// The `(key, nonce, counter)` triple must never be reused across distinct messages. Setting
142    /// the counter near `u32::MAX` leaves little room before the counter space is exhausted;
143    /// streaming past it will panic rather than reuse keystream (see [`CHACHA20_IETF_MAX_LEN`]).
144    #[inline(always)]
145    pub fn set_counter(&mut self, counter: u32) {
146        Self::inject_counter(&mut self.state, counter as u64);
147        self.keystream_leftover_offset = (BLOCK_SIZE - 1) as u8;
148    }
149}
150
151impl<const ROUNDS: usize, const IS_IETF: bool> StreamCipher for ChaCha<ROUNDS, IS_IETF> {
152    /// XOR `plaintext` with the ChaCha keystream.
153    ///
154    /// # Panics
155    ///
156    /// For the IETF layout, panics if this call would advance the 32-bit block counter past
157    /// `u32::MAX` and silently reuse keystream. This happens when encrypting more than
158    /// [`CHACHA20_IETF_MAX_LEN`] bytes under one key/nonce, or after
159    /// [`set_counter`](ChaCha::set_counter) jumps near the end of the counter space.
160    fn xor_keystream(&mut self, mut in_out: &mut [u8]) {
161        if in_out.len() == 0 {
162            return;
163        }
164
165        // For the IETF layout, reject any call that would advance the 32-bit block counter past
166        // u32::MAX and silently reuse keystream. Compute how many new blocks this call needs,
167        // taking into account the leftover keystream buffered by a previous partial call.
168        if IS_IETF {
169            let leftover_len = if self.keystream_leftover_offset < (BLOCK_SIZE - 1) as u8 {
170                (BLOCK_SIZE - 1) - self.keystream_leftover_offset as usize
171            } else {
172                0
173            };
174            let new_blocks = in_out.len().saturating_sub(leftover_len).div_ceil(BLOCK_SIZE) as u64;
175            let counter = Self::extract_counter(&self.state);
176
177            assert!(
178                counter + new_blocks <= u32::MAX as u64,
179                "input exceeds maximum allowed length ((2^32 - 1) * 64 bytes)"
180            );
181        }
182
183        // first, consume the keystream leftover, if any
184        if self.keystream_leftover_offset < (BLOCK_SIZE - 1) as u8 {
185            let keystream_leftover = &self.keystream_leftover[(self.keystream_leftover_offset as usize)..];
186
187            in_out
188                .iter_mut()
189                .zip(keystream_leftover)
190                .for_each(|(plaintext, keystream)| *plaintext ^= *keystream);
191
192            if in_out.len() > keystream_leftover.len() {
193                in_out = &mut in_out[keystream_leftover.len()..];
194            } else if in_out.len() < keystream_leftover.len() {
195                self.keystream_leftover_offset += in_out.len() as u8;
196                return;
197            } else {
198                // in_out.len() == keystream_leftover.len() -> in_out has consumed exactly all the
199                // leftover keystream
200                self.keystream_leftover_offset = (BLOCK_SIZE - 1) as u8;
201                return;
202            }
203        }
204        // at this point, we already know how many bytes of leftover there will be
205        self.keystream_leftover_offset = ((in_out.len() + BLOCK_SIZE - 1) % BLOCK_SIZE) as u8;
206
207        // runtime detection of CPU features for x86 and x86_64 when the "std" feature is enabled
208        #[cfg(all(feature = "std", target_arch = "x86_64"))]
209        if in_out.len() >= 128 && is_x86_feature_detected!("avx512f") {
210            use super::chacha_avx512::chacha_avx512;
211            unsafe { chacha_avx512::<ROUNDS, IS_IETF>(&mut self.state, in_out, &mut self.keystream_leftover) };
212            return;
213        }
214
215        #[cfg(all(feature = "std", target_arch = "x86_64"))]
216        if in_out.len() >= 128 && is_x86_feature_detected!("avx2") {
217            use super::chacha_avx2::chacha_avx2;
218            unsafe { chacha_avx2::<ROUNDS, IS_IETF>(&mut self.state, in_out, &mut self.keystream_leftover) };
219            return;
220        }
221
222        // compile-time dispatch
223        #[cfg(all(not(feature = "std"), target_arch = "x86_64", target_feature = "avx512f"))]
224        if in_out.len() >= 128 {
225            use super::chacha_avx512::chacha_avx512;
226            unsafe { chacha_avx512::<ROUNDS, IS_IETF>(&mut self.state, in_out, &mut self.keystream_leftover) };
227            return;
228        }
229
230        #[cfg(all(target_arch = "x86_64", target_feature = "avx2"))]
231        if in_out.len() >= 128 {
232            use super::chacha_avx2::chacha_avx2;
233            unsafe { chacha_avx2::<ROUNDS, IS_IETF>(&mut self.state, in_out, &mut self.keystream_leftover) };
234            return;
235        }
236
237        // aarch64 assumes that NEON is always available so compile-time dispatch is enough
238        #[cfg(all(target_arch = "aarch64", target_feature = "neon"))]
239        if in_out.len() >= 128 {
240            use super::chacha_neon::chacha_neon;
241            unsafe { chacha_neon::<ROUNDS, IS_IETF>(&mut self.state, in_out, &mut self.keystream_leftover) };
242            return;
243        }
244
245        // wasm32 only supports compile-time features detection
246        #[cfg(all(target_arch = "wasm32", target_feature = "simd128"))]
247        if in_out.len() >= 128 {
248            use super::chacha_wasm_simd128::chacha_wasm_simd128;
249            chacha_wasm_simd128::<ROUNDS, IS_IETF>(&mut self.state, in_out, &mut self.keystream_leftover);
250            return;
251        }
252
253        // fallback for when SIMD acceleration is not available
254        chacha_generic::<ROUNDS, IS_IETF>(&mut self.state, &mut self.keystream_leftover, in_out);
255    }
256}
257
258#[inline]
259fn chacha_generic<const ROUNDS: usize, const IS_IETF: bool>(
260    state: &mut [u32; STATE_WORDS],
261    keystream_leftover: &mut [u8; BLOCK_SIZE - 1],
262    plaintext: &mut [u8],
263) {
264    let mut keystream = [0u8; BLOCK_SIZE];
265    let keystream_ptr = keystream.as_mut_ptr();
266    let mut counter = ChaCha::<ROUNDS, IS_IETF>::extract_counter(state);
267
268    // process the input by blocks of 64 bytes
269    for plaintext_block in plaintext.chunks_mut(BLOCK_SIZE) {
270        ChaCha::<ROUNDS, IS_IETF>::inject_counter(state, counter);
271
272        // prepare temporary (working) state
273        let mut tmp_state = *state;
274
275        // perform the ROUNDS / 2 double rounds e.g. 10 double rounds for ChaCha20
276        for _ in 0..(ROUNDS / 2) {
277            // column rounds
278            quarter_round(&mut tmp_state, 0, 4, 8, 12);
279            quarter_round(&mut tmp_state, 1, 5, 9, 13);
280            quarter_round(&mut tmp_state, 2, 6, 10, 14);
281            quarter_round(&mut tmp_state, 3, 7, 11, 15);
282
283            // diagonal rounds
284            quarter_round(&mut tmp_state, 0, 5, 10, 15);
285            quarter_round(&mut tmp_state, 1, 6, 11, 12);
286            quarter_round(&mut tmp_state, 2, 7, 8, 13);
287            quarter_round(&mut tmp_state, 3, 4, 9, 14);
288        }
289
290        // add initial state to tmp_state to generate the keystream and "serialize" it to little endian
291        // for (tmp_word, state_word) in tmp_state.iter_mut().zip(state.iter()) {
292        //     *tmp_word = tmp_word.wrapping_add(*state_word).to_le();
293        // }
294        for word_index in 0..STATE_WORDS {
295            // first we add the initial state to the working state to get the keystream
296            tmp_state[word_index] = tmp_state[word_index].wrapping_add(state[word_index]);
297
298            // then we serialize the keystream
299            // SAFETY: this is safe because `tmp_state` and `keystream` both have fixed, known size.
300            // We are just merely converting [u32; STATE_WORDS] to [u8; BLOCK_SIZE] with the correct
301            // endianness.
302            unsafe {
303                core::ptr::copy_nonoverlapping(
304                    tmp_state[word_index].to_le_bytes().as_ptr(),
305                    keystream_ptr.add(word_index * 4),
306                    4,
307                );
308            }
309        }
310
311        // XOR plaintext with keystream
312        plaintext_block
313            .iter_mut()
314            .zip(keystream)
315            .for_each(|(plaintext, keystream)| *plaintext ^= keystream);
316
317        counter = counter.wrapping_add(1);
318    }
319
320    ChaCha::<ROUNDS, IS_IETF>::inject_counter(state, counter);
321
322    if !plaintext.len().is_multiple_of(BLOCK_SIZE) {
323        // copy the last 63 bytes of the leftover keystream block
324        keystream_leftover.copy_from_slice(&keystream[1..]);
325    }
326}
327
328#[inline(always)]
329pub(crate) const fn quarter_round(state: &mut [u32; 16], a: usize, b: usize, c: usize, d: usize) {
330    // a += b; d ^= a; d <<<= 16
331    state[a] = state[a].wrapping_add(state[b]);
332    state[d] ^= state[a];
333    state[d] = state[d].rotate_left(16);
334
335    // c += d; b ^= c; b <<<= 12
336    state[c] = state[c].wrapping_add(state[d]);
337    state[b] ^= state[c];
338    state[b] = state[b].rotate_left(12);
339
340    // a += b; d ^= a; d <<<= 8
341    state[a] = state[a].wrapping_add(state[b]);
342    state[d] ^= state[a];
343    state[d] = state[d].rotate_left(8);
344
345    // c += d; b ^= c; b <<<= 7
346    state[c] = state[c].wrapping_add(state[d]);
347    state[b] ^= state[c];
348    state[b] = state[b].rotate_left(7);
349}
350
351/// XChaCha20 stream cipher with 24-byte (192-bit) nonce (draft-irtf-cfrg-xchacha-03).
352///
353/// Internally uses HChaCha20 to derive a subkey from the first 16 bytes of the nonce,
354/// then uses the IETF ChaCha20 variant with the remaining 8 nonce bytes.
355#[cfg_attr(feature = "zeroize", derive(Zeroize, ZeroizeOnDrop))]
356pub struct XChaCha<const ROUNDS: usize> {
357    inner: ChaCha<ROUNDS, true>,
358}
359
360impl<const ROUNDS: usize> XChaCha<ROUNDS> {
361    /// Creates a new XChaCha stream cipher from a 32-byte key and a 24-byte nonce.
362    ///
363    /// The first 16 bytes of the nonce are used with HChaCha20 to derive a subkey.
364    /// The last 8 bytes of the nonce become the ChaCha20 IETF nonce (prefixed with 4 zero bytes).
365    pub fn new(key: &[u8; 32], nonce: &[u8; 24]) -> XChaCha<ROUNDS> {
366        let subkey = super::hchacha20(key, nonce[..16].try_into().unwrap());
367        let mut ietf_nonce = [0u8; 12];
368        ietf_nonce[4..12].copy_from_slice(&nonce[16..24]);
369        return XChaCha {
370            inner: ChaCha::<ROUNDS, true>::new(&subkey, &ietf_nonce),
371        };
372    }
373
374    /// Sets the ChaCha20 block counter. Counter is a u32.
375    pub fn set_counter(&mut self, counter: u32) {
376        self.inner.set_counter(counter);
377    }
378}
379
380impl<const ROUNDS: usize> StreamCipher for XChaCha<ROUNDS> {
381    fn xor_keystream(&mut self, in_out: &mut [u8]) {
382        self.inner.xor_keystream(in_out);
383    }
384}
385
386#[cfg(test)]
387mod test {
388    use super::{CHACHA20_IETF_MAX_LEN, ChaCha8Djb, ChaCha12Djb, ChaCha20Djb, ChaCha20Ietf, XChaCha20};
389    use crate::StreamCipher;
390
391    struct TestDjb {
392        key: [u8; 32],
393        nonce: [u8; 8],
394        initial_counter: u64,
395        plaintext: Vec<u8>,
396        expected_ciphertext: Vec<u8>,
397    }
398
399    #[test]
400    fn chacha20_test_vectors_djb() {
401        let tests = vec![
402            // https://www.rfc-editor.org/rfc/rfc8439#section-2.4.2
403            TestDjb {
404                key: hex::decode("000102030405060708090a0b0c0d0e0f101112131415161718191a1b1c1d1e1f")
405                    .unwrap()
406                    .try_into()
407                    .unwrap(),
408                nonce: hex::decode("0000004a00000000").unwrap().try_into().unwrap(),
409                initial_counter: 1,
410                plaintext: hex::decode(
411                    "4c616469657320616e642047656e746c\
412656d656e206f662074686520636c6173\
41373206f66202739393a20496620492063\
4146f756c64206f6666657220796f75206f\
4156e6c79206f6e652074697020666f7220\
416746865206675747572652c2073756e73\
417637265656e20776f756c642062652069\
418742e",
419                )
420                .unwrap(),
421                expected_ciphertext: hex::decode(
422                    "6e2e359a2568f98041ba0728dd0d6981\
423e97e7aec1d4360c20a27afccfd9fae0b\
424f91b65c5524733ab8f593dabcd62b357\
4251639d624e65152ab8f530c359f0861d8\
42607ca0dbf500d6a6156a38e088a22b65e\
42752bc514d16ccf806818ce91ab7793736\
4285af90bbf74a35be6b40b8eedf2785e42\
429874d",
430                )
431                .unwrap(),
432            },
433            // https://www.rfc-editor.org/rfc/rfc8439#appendix-A.2 Test vector #1
434            TestDjb {
435                key: [0u8; 32],
436                nonce: [0u8; 8],
437                initial_counter: 0,
438                plaintext: [0u8; 64].to_vec(),
439                expected_ciphertext: hex::decode(
440                    "76b8e0ada0f13d90405d6ae55386bd28\
441bdd219b8a08ded1aa836efcc8b770dc7\
442da41597c5157488d7724e03fb8d84a37\
4436a43b8f41518a11cc387b669b2ee6586",
444                )
445                .unwrap(),
446            },
447            // https://www.rfc-editor.org/rfc/rfc8439#appendix-A.2 Test Vector #2
448            TestDjb {
449                key: hex::decode("0000000000000000000000000000000000000000000000000000000000000001")
450                    .unwrap()
451                    .try_into()
452                    .unwrap(),
453                nonce: hex::decode("0000000000000002").unwrap().try_into().unwrap(),
454                initial_counter: 1,
455                plaintext: hex::decode(
456                    "416e79207375626d697373696f6e2074\
4576f20746865204945544620696e74656e\
4586465642062792074686520436f6e7472\
459696275746f7220666f72207075626c69\
460636174696f6e20617320616c6c206f72\
4612070617274206f6620616e2049455446\
46220496e7465726e65742d447261667420\
4636f722052464320616e6420616e792073\
464746174656d656e74206d616465207769\
4657468696e2074686520636f6e74657874\
466206f6620616e20494554462061637469\
4677669747920697320636f6e7369646572\
468656420616e20224945544620436f6e74\
4697269627574696f6e222e205375636820\
47073746174656d656e747320696e636c75\
4716465206f72616c2073746174656d656e\
472747320696e2049455446207365737369\
4736f6e732c2061732077656c6c20617320\
4747772697474656e20616e6420656c6563\
47574726f6e696320636f6d6d756e696361\
47674696f6e73206d61646520617420616e\
477792074696d65206f7220706c6163652c\
47820776869636820617265206164647265\
4797373656420746f",
480                )
481                .unwrap(),
482                expected_ciphertext: hex::decode(
483                    "a3fbf07df3fa2fde4f376ca23e827370\
48441605d9f4f4f57bd8cff2c1d4b7955ec\
4852a97948bd3722915c8f3d337f7d37005\
4860e9e96d647b7c39f56e031ca5eb6250d\
4874042e02785ececfa4b4bb5e8ead0440e\
48820b6e8db09d881a7c6132f420e527950\
48942bdfa7773d8a9051447b3291ce1411c\
490680465552aa6c405b7764d5e87bea85a\
491d00f8449ed8f72d0d662ab052691ca66\
492424bc86d2df80ea41f43abf937d3259d\
493c4b2d0dfb48a6c9139ddd7f76966e928\
494e635553ba76c5c879d7b35d49eb2e62b\
4950871cdac638939e25e8a1e0ef9d5280f\
496a8ca328b351c3c765989cbcf3daa8b6c\
497cc3aaf9f3979c92b3720fc88dc95ed84\
498a1be059c6499b9fda236e7e818b04b0b\
499c39c1e876b193bfe5569753f88128cc0\
5008aaa9b63d1a16f80ef2554d7189c411f\
5015869ca52c5b83fa36ff216b9c1d30062\
502bebcfd2dc5bce0911934fda79a86f6e6\
50398ced759c3ff9b6477338f3da4f9cd85\
50414ea9982ccafb341b2384dd902f3d1ab\
5057ac61dd29c6f21ba5b862f3730e37cfd\
506c4fd806c22f221",
507                )
508                .unwrap(),
509            },
510            // https://www.rfc-editor.org/rfc/rfc8439#appendix-A.2 Test Vector #3
511            TestDjb {
512                key: hex::decode("1c9240a5eb55d38af333888604f6b5f0473917c1402b80099dca5cbc207075c0")
513                    .unwrap()
514                    .try_into()
515                    .unwrap(),
516                nonce: hex::decode("0000000000000002").unwrap().try_into().unwrap(),
517                initial_counter: 42,
518                plaintext: hex::decode(
519                    "2754776173206272696c6c69672c2061\
5206e642074686520736c6974687920746f\
5217665730a446964206779726520616e64\
5222067696d626c6520696e207468652077\
5236162653a0a416c6c206d696d73792077\
5246572652074686520626f726f676f7665\
525732c0a416e6420746865206d6f6d6520\
5267261746873206f757467726162652e",
527                )
528                .unwrap(),
529                expected_ciphertext: hex::decode(
530                    "62e6347f95ed87a45ffae7426f27a1df\
5315fb69110044c0d73118effa95b01e5cf\
532166d3df2d721caf9b21e5fb14c616871\
533fd84c54f9d65b283196c7fe4f60553eb\
534f39c6402c42234e32a356b3e764312a6\
5351a5532055716ead6962568f87d3f3f77\
53604c6a8d1bcd1bf4d50d6154b6da731b1\
53787b58dfd728afa36757a797ac188d1",
538                )
539                .unwrap(),
540            },
541        ];
542
543        for (i, test) in tests.into_iter().enumerate() {
544            let mut cipher = ChaCha20Djb::new(&test.key, &test.nonce);
545            cipher.set_counter(test.initial_counter);
546
547            let mut plaintext = test.plaintext.clone();
548            cipher.xor_keystream(&mut plaintext);
549
550            assert_eq!(
551                plaintext,
552                test.expected_ciphertext,
553                "test [{i}] failed
554Got ciphertext: {}
555Expected ciphertext: {}",
556                hex::encode(&plaintext),
557                hex::encode(&test.expected_ciphertext),
558            );
559
560            let mut cipher = ChaCha20Djb::new(&test.key, &test.nonce);
561            cipher.set_counter(test.initial_counter);
562            cipher.xor_keystream(&mut plaintext);
563
564            assert_eq!(
565                plaintext,
566                test.plaintext,
567                "test [{i}] failed. Initial plaintext != decrypt(encrypt(plaintext))
568Got: {}
569Expected: {}",
570                hex::encode(&plaintext),
571                hex::encode(&test.plaintext),
572            );
573
574            // ensure that the encryption is correct even for plaintexts that are not % 64 (block size)
575            // thus:
576            // cipher.xor_keystream(plaintext[0..10])
577            // cipher.xor_keystream(plaintext[10..30])
578            // cipher.xor_keystream(plaintext[30..5])
579            // should be equal to:
580            // cipher.xor_keystream(plaintext[0..35])
581
582            let mut cipher = ChaCha20Djb::new(&test.key, &test.nonce);
583            cipher.xor_keystream(&mut plaintext);
584            for n in 0..10 {
585                let mut partial_plaintext: Vec<u8> = test.plaintext.clone();
586
587                let mut cipher = ChaCha20Djb::new(&test.key, &test.nonce);
588                cipher.xor_keystream(&mut partial_plaintext[..n]);
589                cipher.xor_keystream(&mut partial_plaintext[n..]);
590
591                assert_eq!(
592                    plaintext,
593                    partial_plaintext,
594                    "test [{i}] failed. partial encryption is not valid for n = {n}
595            Got: {}
596            Expected: {}",
597                    hex::encode(&partial_plaintext),
598                    hex::encode(&plaintext),
599                )
600            }
601        }
602    }
603
604    #[test]
605    fn chacha20_keystream_leftover_multi_call() {
606        // Regression: verifies that leftovers are correctly consumed across 3+
607        // partial calls. The old length-based approach did not compact the array
608        // after partial consumption, causing stale keystream reuse on subsequent calls.
609        let key = hex::decode("000102030405060708090a0b0c0d0e0f101112131415161718191a1b1c1d1e1f")
610            .unwrap()
611            .try_into()
612            .unwrap();
613        let nonce = hex::decode("0000004a00000000").unwrap().try_into().unwrap();
614        let plaintext = hex::decode(
615            "4c616469657320616e642047656e746c\
616656d656e206f662074686520636c6173\
61773206f66202739393a20496620492063\
6186f756c64206f6666657220796f75206f\
6196e6c79206f6e652074697020666f7220\
620746865206675747572652c2073756e73\
621637265656e20776f756c642062652069\
622742e",
623        )
624        .unwrap();
625
626        let mut expected = plaintext.clone();
627        ChaCha20Djb::new(&key, &nonce).xor_keystream(&mut expected);
628
629        // call 1: partial block -> leaves leftover
630        // call 2: partially consumes leftover
631        // call 3: consumes remaining leftover + fresh blocks
632        {
633            let mut buf = plaintext.clone();
634            let mut cipher = ChaCha20Djb::new(&key, &nonce);
635            cipher.xor_keystream(&mut buf[..10]);
636            cipher.xor_keystream(&mut buf[10..15]);
637            cipher.xor_keystream(&mut buf[15..]);
638            assert_eq!(buf, expected, "partial leftover consumption");
639        }
640
641        // call 1: partial block -> leaves leftover
642        // call 2: exactly exhausts leftover
643        // call 3: fresh blocks
644        {
645            let mut buf = plaintext.clone();
646            let mut cipher = ChaCha20Djb::new(&key, &nonce);
647            cipher.xor_keystream(&mut buf[..10]);
648            cipher.xor_keystream(&mut buf[10..64]);
649            cipher.xor_keystream(&mut buf[64..]);
650            assert_eq!(buf, expected, "exact leftover exhaustion");
651        }
652
653        // call 1: partial block -> leaves leftover
654        // call 2 + call 3: two rounds of partial consumption
655        // call 4: consumes remaining leftover + fresh blocks
656        {
657            let mut buf = plaintext.clone();
658            let mut cipher = ChaCha20Djb::new(&key, &nonce);
659            cipher.xor_keystream(&mut buf[..8]);
660            cipher.xor_keystream(&mut buf[8..13]);
661            cipher.xor_keystream(&mut buf[13..33]);
662            cipher.xor_keystream(&mut buf[33..]);
663            assert_eq!(buf, expected, "multiple partial leftover consumptions");
664        }
665    }
666
667    #[test]
668    fn chacha20_ietf_test_vectors() {
669        // IETF ChaCha20 test vectors from RFC 8439 Appendix A.2.
670        // These use the IETF layout: 32-bit counter (word 12), 96-bit nonce (words 13–15).
671        // The nonce is the full 12-byte value (little-endian) that RFC 8439 places in words 13–15.
672
673        struct TestIetf {
674            key: [u8; 32],
675            nonce: [u8; 12],
676            initial_counter: u32,
677            plaintext: Vec<u8>,
678            expected_ciphertext: Vec<u8>,
679        }
680
681        let tests = vec![
682            // RFC 8439 section 2.4.2 — the "Ladies and Gentlemen" test vector
683            // IETF state: word 12 = counter, words 13-15 = 96-bit nonce LE
684            // So state[13]=0, state[14]=0x4a000000, state[15]=0
685            TestIetf {
686                key: hex::decode("000102030405060708090a0b0c0d0e0f101112131415161718191a1b1c1d1e1f")
687                    .unwrap()
688                    .try_into()
689                    .unwrap(),
690                nonce: hex::decode("000000000000004a00000000").unwrap().try_into().unwrap(),
691                initial_counter: 1,
692                plaintext: hex::decode(
693                    "4c616469657320616e642047656e746c\
694656d656e206f662074686520636c6173\
69573206f66202739393a20496620492063\
6966f756c64206f6666657220796f75206f\
6976e6c79206f6e652074697020666f7220\
698746865206675747572652c2073756e73\
699637265656e20776f756c642062652069\
700742e",
701                )
702                .unwrap(),
703                expected_ciphertext: hex::decode(
704                    "6e2e359a2568f98041ba0728dd0d6981\
705e97e7aec1d4360c20a27afccfd9fae0b\
706f91b65c5524733ab8f593dabcd62b357\
7071639d624e65152ab8f530c359f0861d8\
70807ca0dbf500d6a6156a38e088a22b65e\
70952bc514d16ccf806818ce91ab7793736\
7105af90bbf74a35be6b40b8eedf2785e42\
711874d",
712                )
713                .unwrap(),
714            },
715            // RFC 8439 Appendix A.2 Vector #1 — all zeros key, nonce, counter=0
716            TestIetf {
717                key: [0u8; 32],
718                nonce: [0u8; 12],
719                initial_counter: 0,
720                plaintext: [0u8; 64].to_vec(),
721                expected_ciphertext: hex::decode(
722                    "76b8e0ada0f13d90405d6ae55386bd28\
723bdd219b8a08ded1aa836efcc8b770dc7\
724da41597c5157488d7724e03fb8d84a37\
7256a43b8f41518a11cc387b669b2ee6586",
726                )
727                .unwrap(),
728            },
729            // RFC 8439 Appendix A.2 Vector #2 — counter=1
730            TestIetf {
731                key: hex::decode("0000000000000000000000000000000000000000000000000000000000000001")
732                    .unwrap()
733                    .try_into()
734                    .unwrap(),
735                nonce: hex::decode("000000000000000000000002").unwrap().try_into().unwrap(),
736                initial_counter: 1,
737                plaintext: hex::decode(
738                    "416e79207375626d697373696f6e2074\
7396f20746865204945544620696e74656e\
7406465642062792074686520436f6e7472\
741696275746f7220666f72207075626c69\
742636174696f6e20617320616c6c206f72\
7432070617274206f6620616e2049455446\
74420496e7465726e65742d447261667420\
7456f722052464320616e6420616e792073\
746746174656d656e74206d616465207769\
7477468696e2074686520636f6e74657874\
748206f6620616e20494554462061637469\
7497669747920697320636f6e7369646572\
750656420616e20224945544620436f6e74\
7517269627574696f6e222e205375636820\
75273746174656d656e747320696e636c75\
7536465206f72616c2073746174656d656e\
754747320696e2049455446207365737369\
7556f6e732c2061732077656c6c20617320\
7567772697474656e20616e6420656c6563\
75774726f6e696320636f6d6d756e696361\
75874696f6e73206d61646520617420616e\
759792074696d65206f7220706c6163652c\
76020776869636820617265206164647265\
7617373656420746f",
762                )
763                .unwrap(),
764                expected_ciphertext: hex::decode(
765                    "a3fbf07df3fa2fde4f376ca23e827370\
76641605d9f4f4f57bd8cff2c1d4b7955ec\
7672a97948bd3722915c8f3d337f7d37005\
7680e9e96d647b7c39f56e031ca5eb6250d\
7694042e02785ececfa4b4bb5e8ead0440e\
77020b6e8db09d881a7c6132f420e527950\
77142bdfa7773d8a9051447b3291ce1411c\
772680465552aa6c405b7764d5e87bea85a\
773d00f8449ed8f72d0d662ab052691ca66\
774424bc86d2df80ea41f43abf937d3259d\
775c4b2d0dfb48a6c9139ddd7f76966e928\
776e635553ba76c5c879d7b35d49eb2e62b\
7770871cdac638939e25e8a1e0ef9d5280f\
778a8ca328b351c3c765989cbcf3daa8b6c\
779cc3aaf9f3979c92b3720fc88dc95ed84\
780a1be059c6499b9fda236e7e818b04b0b\
781c39c1e876b193bfe5569753f88128cc0\
7828aaa9b63d1a16f80ef2554d7189c411f\
7835869ca52c5b83fa36ff216b9c1d30062\
784bebcfd2dc5bce0911934fda79a86f6e6\
78598ced759c3ff9b6477338f3da4f9cd85\
78614ea9982ccafb341b2384dd902f3d1ab\
7877ac61dd29c6f21ba5b862f3730e37cfd\
788c4fd806c22f221",
789                )
790                .unwrap(),
791            },
792            // RFC 8439 Appendix A.2 Vector #3 — counter=42
793            TestIetf {
794                key: hex::decode("1c9240a5eb55d38af333888604f6b5f0473917c1402b80099dca5cbc207075c0")
795                    .unwrap()
796                    .try_into()
797                    .unwrap(),
798                nonce: hex::decode("000000000000000000000002").unwrap().try_into().unwrap(),
799                initial_counter: 42,
800                plaintext: hex::decode(
801                    "2754776173206272696c6c69672c2061\
8026e642074686520736c6974687920746f\
8037665730a446964206779726520616e64\
8042067696d626c6520696e207468652077\
8056162653a0a416c6c206d696d73792077\
8066572652074686520626f726f676f7665\
807732c0a416e6420746865206d6f6d6520\
8087261746873206f757467726162652e",
809                )
810                .unwrap(),
811                expected_ciphertext: hex::decode(
812                    "62e6347f95ed87a45ffae7426f27a1df\
8135fb69110044c0d73118effa95b01e5cf\
814166d3df2d721caf9b21e5fb14c616871\
815fd84c54f9d65b283196c7fe4f60553eb\
816f39c6402c42234e32a356b3e764312a6\
8171a5532055716ead6962568f87d3f3f77\
81804c6a8d1bcd1bf4d50d6154b6da731b1\
81987b58dfd728afa36757a797ac188d1",
820                )
821                .unwrap(),
822            },
823        ];
824
825        for (i, test) in tests.into_iter().enumerate() {
826            let mut cipher = ChaCha20Ietf::new(&test.key, &test.nonce);
827            cipher.set_counter(test.initial_counter);
828
829            let mut plaintext = test.plaintext.clone();
830            cipher.xor_keystream(&mut plaintext);
831
832            assert_eq!(
833                plaintext,
834                test.expected_ciphertext,
835                "ietf test [{i}] failed
836Got ciphertext: {}
837Expected ciphertext: {}",
838                hex::encode(&plaintext),
839                hex::encode(&test.expected_ciphertext),
840            );
841
842            // decrypt
843            let mut cipher = ChaCha20Ietf::new(&test.key, &test.nonce);
844            cipher.set_counter(test.initial_counter);
845            cipher.xor_keystream(&mut plaintext);
846
847            assert_eq!(
848                plaintext,
849                test.plaintext,
850                "ietf test [{i}] failed. Initial plaintext != decrypt(encrypt(plaintext))
851Got: {}
852Expected: {}",
853                hex::encode(&plaintext),
854                hex::encode(&test.plaintext),
855            );
856
857            // partial encryption check
858            let mut cipher = ChaCha20Ietf::new(&test.key, &test.nonce);
859            cipher.set_counter(test.initial_counter);
860            cipher.xor_keystream(&mut plaintext);
861            for n in 0..10 {
862                let mut partial_plaintext: Vec<u8> = test.plaintext.clone();
863
864                let mut cipher = ChaCha20Ietf::new(&test.key, &test.nonce);
865                cipher.set_counter(test.initial_counter);
866                cipher.xor_keystream(&mut partial_plaintext[..n]);
867                cipher.xor_keystream(&mut partial_plaintext[n..]);
868
869                assert_eq!(
870                    plaintext,
871                    partial_plaintext,
872                    "ietf test [{i}] failed. partial encryption is not valid for n = {n}
873            Got: {}
874            Expected: {}",
875                    hex::encode(&partial_plaintext),
876                    hex::encode(&plaintext),
877                )
878            }
879        }
880    }
881
882    #[test]
883    fn chacha12_case_1() {
884        let nonce: &[u8; 8] = &[0xdb, 0x4b, 0x4a, 0x41, 0xd8, 0xdf, 0x18, 0xaa];
885        let key: &[u8; 32] = &[
886            0x27, 0xfc, 0x12, 0x0b, 0x01, 0x3b, 0x82, 0x9f, 0x1f, 0xae, 0xef, 0xd1, 0xab, 0x41, 0x7e, 0x86, 0x62, 0xf4,
887            0x3e, 0x0d, 0x73, 0xf9, 0x8d, 0xe8, 0x66, 0xe3, 0x46, 0x35, 0x31, 0x80, 0xfd, 0xb7,
888        ];
889
890        let mut buffer = [0u8; 100];
891        ChaCha12Djb::new(key, nonce).xor_keystream(&mut buffer);
892
893        assert_eq!(
894            buffer,
895            [
896                0x5f, 0x3c, 0x8c, 0x19, 0x0a, 0x78, 0xab, 0x7f, 0xe8, 0x08, 0xca, 0xe9, 0xcb, 0xcb, 0x0a, 0x98, 0x37,
897                0xc8, 0x93, 0x49, 0x2d, 0x96, 0x3a, 0x1c, 0x2e, 0xda, 0x6c, 0x15, 0x58, 0xb0, 0x2c, 0x83, 0xfc, 0x02,
898                0xa4, 0x4c, 0xbb, 0xb7, 0xe6, 0x20, 0x4d, 0x51, 0xd1, 0xc2, 0x43, 0x0e, 0x9c, 0x0b, 0x58, 0xf2, 0x93,
899                0x7b, 0xf5, 0x93, 0x84, 0x0c, 0x85, 0x0b, 0xda, 0x90, 0x51, 0xa1, 0xf0, 0x51, 0xdd, 0xf0, 0x9d, 0x2a,
900                0x03, 0xeb, 0xf0, 0x9f, 0x01, 0xbd, 0xba, 0x9d, 0xa0, 0xb6, 0xda, 0x79, 0x1b, 0x2e, 0x64, 0x56, 0x41,
901                0x04, 0x7d, 0x11, 0xeb, 0xf8, 0x50, 0x87, 0xd4, 0xde, 0x5c, 0x01, 0x5f, 0xdd, 0xd0, 0x44,
902            ]
903        );
904    }
905
906    #[test]
907    fn chacha8_case_1() {
908        let key = &[
909            0x64, 0x1a, 0xea, 0xeb, 0x08, 0x03, 0x6b, 0x61, 0x7a, 0x42, 0xcf, 0x14, 0xe8, 0xc5, 0xd2, 0xd1, 0x15, 0xf8,
910            0xd7, 0xcb, 0x6e, 0xa5, 0xe2, 0x8b, 0x9b, 0xfa, 0xf8, 0x3e, 0x03, 0x84, 0x26, 0xa7,
911        ];
912        let nonce = &[0xa1, 0x4a, 0x11, 0x68, 0x27, 0x1d, 0x45, 0x9b];
913
914        let mut buffer = [0u8; 100];
915        ChaCha8Djb::new(key, nonce).xor_keystream(&mut buffer);
916
917        assert_eq!(
918            buffer,
919            [
920                0x17, 0x21, 0xc0, 0x44, 0xa8, 0xa6, 0x45, 0x35, 0x22, 0xdd, 0xdb, 0x31, 0x43, 0xd0, 0xbe, 0x35, 0x12,
921                0x63, 0x3c, 0xa3, 0xc7, 0x9b, 0xf8, 0xcc, 0xc3, 0x59, 0x4c, 0xb2, 0xc2, 0xf3, 0x10, 0xf7, 0xbd, 0x54,
922                0x4f, 0x55, 0xce, 0x0d, 0xb3, 0x81, 0x23, 0x41, 0x2d, 0x6c, 0x45, 0x20, 0x7d, 0x5c, 0xf9, 0xaf, 0x0c,
923                0x6c, 0x68, 0x0c, 0xce, 0x1f, 0x7e, 0x43, 0x38, 0x8d, 0x1b, 0x03, 0x46, 0xb7, 0x13, 0x3c, 0x59, 0xfd,
924                0x6a, 0xf4, 0xa5, 0xa5, 0x68, 0xaa, 0x33, 0x4c, 0xcd, 0xc3, 0x8a, 0xf5, 0xac, 0xe2, 0x01, 0xdf, 0x84,
925                0xd0, 0xa3, 0xca, 0x22, 0x54, 0x94, 0xca, 0x62, 0x09, 0x34, 0x5f, 0xcf, 0x30, 0x13, 0x2e,
926            ]
927        );
928    }
929
930    #[test]
931    fn xchacha20_test_vector_counter_0() {
932        // draft-irtf-cfrg-xchacha-03, Appendix A.3.2.1
933        let key: [u8; 32] = hex::decode("808182838485868788898a8b8c8d8e8f909192939495969798999a9b9c9d9e9f")
934            .unwrap()
935            .try_into()
936            .unwrap();
937        let nonce: [u8; 24] = hex::decode("404142434445464748494a4b4c4d4e4f5051525354555658")
938            .unwrap()
939            .try_into()
940            .unwrap();
941        let plaintext = hex::decode(concat!(
942            "5468652064686f6c65202870726f6e6f756e6365642022646f6c652229206973",
943            "20616c736f206b6e6f776e2061732074686520417369617469632077696c6420",
944            "646f672c2072656420646f672c20616e642077686973746c696e6720646f672e",
945            "2049742069732061626f7574207468652073697a65206f662061204765726d61",
946            "6e20736865706865726420627574206c6f6f6b73206d6f7265206c696b652061",
947            "206c6f6e672d6c656767656420666f782e205468697320686967686c7920656c",
948            "757369766520616e6420736b696c6c6564206a756d70657220697320636c6173",
949            "736966696564207769746820776f6c7665732c20636f796f7465732c206a6163",
950            "6b616c732c20616e6420666f78657320696e20746865207461786f6e6f6d6963",
951            "2066616d696c792043616e696461652e",
952        ))
953        .unwrap();
954        let expected_ciphertext = hex::decode(concat!(
955            "4559abba4e48c16102e8bb2c05e6947f50a786de162f9b0b7e592a9b53d0d4e9",
956            "8d8d6410d540a1a6375b26d80dace4fab52384c731acbf16a5923c0c48d3575d",
957            "4d0d2c673b666faa731061277701093a6bf7a158a8864292a41c48e3a9b4c0da",
958            "ece0f8d98d0d7e05b37a307bbb66333164ec9e1b24ea0d6c3ffddcec4f68e744",
959            "3056193a03c810e11344ca06d8ed8a2bfb1e8d48cfa6bc0eb4e2464b74814240",
960            "7c9f431aee769960e15ba8b96890466ef2457599852385c661f752ce20f9da0c",
961            "09ab6b19df74e76a95967446f8d0fd415e7bee2a12a114c20eb5292ae7a349ae",
962            "577820d5520a1f3fb62a17ce6a7e68fa7c79111d8860920bc048ef43fe84486c",
963            "cb87c25f0ae045f0cce1e7989a9aa220a28bdd4827e751a24a6d5c62d790a663",
964            "93b93111c1a55dd7421a10184974c7c5",
965        ))
966        .unwrap();
967
968        let mut cipher = XChaCha20::new(&key, &nonce);
969        let mut buf = plaintext.clone();
970        cipher.xor_keystream(&mut buf);
971        assert_eq!(buf, expected_ciphertext);
972
973        // Decrypt
974        cipher.set_counter(0);
975        cipher.xor_keystream(&mut buf);
976        assert_eq!(buf, plaintext);
977    }
978
979    #[test]
980    fn xchacha20_set_counter() {
981        let key = [0x55u8; 32];
982        let nonce = [0xaau8; 24];
983        let plaintext = b"test message for xchacha20";
984
985        let mut cipher1 = XChaCha20::new(&key, &nonce);
986        let mut buf1 = plaintext.to_vec();
987        cipher1.xor_keystream(&mut buf1);
988
989        // Same cipher but with set_counter(0) should produce same output
990        let mut cipher2 = XChaCha20::new(&key, &nonce);
991        cipher2.set_counter(0);
992        let mut buf2 = plaintext.to_vec();
993        cipher2.xor_keystream(&mut buf2);
994
995        assert_eq!(buf1, buf2);
996
997        // Different counter should produce different output
998        let mut cipher3 = XChaCha20::new(&key, &nonce);
999        cipher3.set_counter(1);
1000        let mut buf3 = plaintext.to_vec();
1001        cipher3.xor_keystream(&mut buf3);
1002        assert_ne!(buf1, buf3);
1003    }
1004
1005    // -------------------------------------------------------------------------
1006    // Edge-case leftover tests
1007    // -------------------------------------------------------------------------
1008
1009    /// Helper: encrypt in one shot to get the expected reference.
1010    fn encrypt_one_shot_djb(key: &[u8; 32], nonce: &[u8; 8], plaintext: &[u8]) -> Vec<u8> {
1011        let mut buf = plaintext.to_vec();
1012        ChaCha20Djb::new(key, nonce).xor_keystream(&mut buf);
1013        buf
1014    }
1015
1016    fn encrypt_one_shot_ietf(key: &[u8; 32], nonce: &[u8; 12], plaintext: &[u8]) -> Vec<u8> {
1017        let mut buf = plaintext.to_vec();
1018        ChaCha20Ietf::new(key, nonce).xor_keystream(&mut buf);
1019        buf
1020    }
1021
1022    fn encrypt_one_shot_xchacha(key: &[u8; 32], nonce: &[u8; 24], plaintext: &[u8]) -> Vec<u8> {
1023        let mut buf = plaintext.to_vec();
1024        XChaCha20::new(key, nonce).xor_keystream(&mut buf);
1025        buf
1026    }
1027
1028    fn test_key_32() -> [u8; 32] {
1029        let mut k = [0u8; 32];
1030        for i in 0..32 {
1031            k[i] = i as u8;
1032        }
1033        k
1034    }
1035
1036    fn test_nonce_8() -> [u8; 8] {
1037        [0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07]
1038    }
1039
1040    fn test_nonce_12() -> [u8; 12] {
1041        [0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0a, 0x0b]
1042    }
1043
1044    fn test_nonce_24() -> [u8; 24] {
1045        let mut n = [0u8; 24];
1046        for i in 0..24 {
1047            n[i] = i as u8;
1048        }
1049        n
1050    }
1051
1052    fn test_plaintext(len: usize) -> Vec<u8> {
1053        (0..len).map(|i| (i % 251) as u8).collect()
1054    }
1055
1056    /// IETF variant: multi-call leftover consumption (analogous to the DJB test above).
1057    #[test]
1058    fn chacha20_ietf_keystream_leftover_multi_call() {
1059        let key = test_key_32();
1060        let nonce = test_nonce_12();
1061        let pt = test_plaintext(300);
1062
1063        let expected = encrypt_one_shot_ietf(&key, &nonce, &pt);
1064
1065        // partial block -> leaves leftover, partially consume, then finish
1066        {
1067            let mut buf = pt.clone();
1068            let mut c = ChaCha20Ietf::new(&key, &nonce);
1069            c.xor_keystream(&mut buf[..10]);
1070            c.xor_keystream(&mut buf[10..15]);
1071            c.xor_keystream(&mut buf[15..]);
1072            assert_eq!(buf, expected, "ietf partial leftover consumption");
1073        }
1074
1075        // partial block -> exactly exhaust leftover -> fresh blocks
1076        {
1077            let mut buf = pt.clone();
1078            let mut c = ChaCha20Ietf::new(&key, &nonce);
1079            c.xor_keystream(&mut buf[..3]);
1080            c.xor_keystream(&mut buf[3..64]);
1081            c.xor_keystream(&mut buf[64..]);
1082            assert_eq!(buf, expected, "ietf exact leftover exhaustion");
1083        }
1084
1085        // three rounds of partial leftover consumption
1086        {
1087            let mut buf = pt.clone();
1088            let mut c = ChaCha20Ietf::new(&key, &nonce);
1089            c.xor_keystream(&mut buf[..5]);
1090            c.xor_keystream(&mut buf[5..12]);
1091            c.xor_keystream(&mut buf[12..20]);
1092            c.xor_keystream(&mut buf[20..]);
1093            assert_eq!(buf, expected, "ietf multiple partial leftover consumptions");
1094        }
1095    }
1096
1097    /// XChaCha20: multi-call leftover consumption.
1098    #[test]
1099    fn xchacha20_keystream_leftover_multi_call() {
1100        let key = test_key_32();
1101        let nonce = test_nonce_24();
1102        let pt = test_plaintext(300);
1103
1104        let expected = encrypt_one_shot_xchacha(&key, &nonce, &pt);
1105
1106        // partial block -> leaves leftover, partially consume, then finish
1107        {
1108            let mut buf = pt.clone();
1109            let mut c = XChaCha20::new(&key, &nonce);
1110            c.xor_keystream(&mut buf[..8]);
1111            c.xor_keystream(&mut buf[8..20]);
1112            c.xor_keystream(&mut buf[20..]);
1113            assert_eq!(buf, expected, "xchacha partial leftover consumption");
1114        }
1115
1116        // three rounds of partial consumption
1117        {
1118            let mut buf = pt.clone();
1119            let mut c = XChaCha20::new(&key, &nonce);
1120            c.xor_keystream(&mut buf[..13]);
1121            c.xor_keystream(&mut buf[13..27]);
1122            c.xor_keystream(&mut buf[27..40]);
1123            c.xor_keystream(&mut buf[40..]);
1124            assert_eq!(buf, expected, "xchacha multiple partial leftover consumptions");
1125        }
1126    }
1127
1128    /// Very small chunks: 1-byte calls to stress the leftover offset state machine.
1129    #[test]
1130    fn chacha_keystream_leftover_tiny_chunks() {
1131        let key = test_key_32();
1132        let nonce = test_nonce_8();
1133        let pt = test_plaintext(200);
1134
1135        // 10 calls of 1 byte each, then the rest in one shot
1136        let mut cipher = ChaCha20Djb::new(&key, &nonce);
1137        let mut buf = pt.clone();
1138        for n in 0..10 {
1139            cipher.xor_keystream(&mut buf[n..n + 1]);
1140        }
1141        cipher.xor_keystream(&mut buf[10..]);
1142        let expected = encrypt_one_shot_djb(&key, &nonce, &pt);
1143        assert_eq!(buf, expected, "tiny-chunk DJB failed");
1144
1145        // IETF variant
1146        let nonce12 = test_nonce_12();
1147        let mut buf = pt.clone();
1148        let mut cipher = ChaCha20Ietf::new(&key, &nonce12);
1149        for n in 0..10 {
1150            cipher.xor_keystream(&mut buf[n..n + 1]);
1151        }
1152        cipher.xor_keystream(&mut buf[10..]);
1153        let expected = encrypt_one_shot_ietf(&key, &nonce12, &pt);
1154        assert_eq!(buf, expected, "tiny-chunk IETF failed");
1155
1156        // XChaCha variant
1157        let nonce24 = test_nonce_24();
1158        let mut buf = pt.clone();
1159        let mut cipher = XChaCha20::new(&key, &nonce24);
1160        for n in 0..10 {
1161            cipher.xor_keystream(&mut buf[n..n + 1]);
1162        }
1163        cipher.xor_keystream(&mut buf[10..]);
1164        let expected = encrypt_one_shot_xchacha(&key, &nonce24, &pt);
1165        assert_eq!(buf, expected, "tiny-chunk XChaCha failed");
1166    }
1167
1168    /// Boundary sizes: 63, 64, 65, 127, 128, 129 bytes.
1169    /// Tests the leftover offset formula at exact transition points.
1170    #[test]
1171    fn chacha_keystream_leftover_boundary_sizes() {
1172        let key = test_key_32();
1173        let nonce = test_nonce_8();
1174
1175        for &len in &[63usize, 64, 65, 127, 128, 129, 191, 192, 193, 255, 256, 257] {
1176            let pt = test_plaintext(len);
1177            let expected = encrypt_one_shot_djb(&key, &nonce, &pt);
1178
1179            // split into 3 chunks: a, b, c where a+b+c = len
1180            // Use varying split sizes to exercise different leftover states
1181            for a in [0usize, 1, 8, 31, 32, 33, 62, 63].iter().copied() {
1182                if a > len {
1183                    continue;
1184                }
1185                for b in [0usize, 1, 7, 32, 63, 64].iter().copied() {
1186                    if a + b > len {
1187                        continue;
1188                    }
1189                    let c = len - a - b;
1190                    let splits = [a, b, c];
1191
1192                    let mut cipher = ChaCha20Djb::new(&key, &nonce);
1193                    let mut buf = pt.clone();
1194                    let mut offset = 0;
1195                    for &size in &splits {
1196                        cipher.xor_keystream(&mut buf[offset..offset + size]);
1197                        offset += size;
1198                    }
1199                    assert_eq!(buf, expected, "Boundary DJB failed for len={len} splits={splits:?}",);
1200                }
1201            }
1202        }
1203    }
1204
1205    /// Boundary sizes for IETF variant.
1206    #[test]
1207    fn chacha_keystream_leftover_boundary_sizes_ietf() {
1208        let key = test_key_32();
1209        let nonce = test_nonce_12();
1210
1211        for &len in &[63usize, 64, 65, 127, 128, 129] {
1212            let pt = test_plaintext(len);
1213            let expected = encrypt_one_shot_ietf(&key, &nonce, &pt);
1214
1215            for &a in &[0usize, 1, 31, 32, 33, 63] {
1216                if a > len {
1217                    continue;
1218                }
1219                let mut cipher = ChaCha20Ietf::new(&key, &nonce);
1220                let mut buf = pt.clone();
1221                cipher.xor_keystream(&mut buf[..a]);
1222                cipher.xor_keystream(&mut buf[a..]);
1223                assert_eq!(buf, expected, "Boundary IETF failed for len={len} a={a}",);
1224            }
1225        }
1226    }
1227
1228    /// Zero-byte intermediate calls: should be no-ops that don't corrupt leftover state.
1229    #[test]
1230    fn chacha_keystream_leftover_zero_byte_intermediate() {
1231        let key = test_key_32();
1232        let nonce = test_nonce_8();
1233        let pt = test_plaintext(150);
1234        let expected = encrypt_one_shot_djb(&key, &nonce, &pt);
1235
1236        let mut buf = pt.clone();
1237        let mut cipher = ChaCha20Djb::new(&key, &nonce);
1238
1239        // first partial call
1240        cipher.xor_keystream(&mut buf[..10]);
1241        // zero-byte call in the middle
1242        cipher.xor_keystream(&mut []);
1243        // second partial call
1244        cipher.xor_keystream(&mut buf[10..20]);
1245        // another zero-byte call
1246        cipher.xor_keystream(&mut []);
1247        // final call
1248        cipher.xor_keystream(&mut buf[20..]);
1249
1250        assert_eq!(buf, expected, "zero-byte intermediate DJB failed");
1251    }
1252
1253    /// set_counter mid-stream, then partial encryption: verifies leftover is cleared
1254    /// and counter is correctly reset.
1255    #[test]
1256    fn chacha_keystream_leftover_set_counter_mid_stream() {
1257        let key = test_key_32();
1258        let nonce = test_nonce_8();
1259        let pt = test_plaintext(200);
1260
1261        // one-shot reference
1262        let full_encrypted = encrypt_one_shot_djb(&key, &nonce, &pt);
1263
1264        // encrypt all at once with a fresh cipher
1265        let mut buf = pt.clone();
1266        let mut cipher = ChaCha20Djb::new(&key, &nonce);
1267        cipher.xor_keystream(&mut buf);
1268        assert_eq!(buf, full_encrypted, "baseline DJB");
1269
1270        // encrypt 10 bytes, reset counter to 0, re-encrypt those 10 bytes
1271        // (this XORs again -> back to plaintext), then encrypt the rest.
1272        let mut buf = pt.clone();
1273        let mut cipher = ChaCha20Djb::new(&key, &nonce);
1274        cipher.xor_keystream(&mut buf[..10]);
1275        cipher.set_counter(0);
1276        cipher.xor_keystream(&mut buf[..10]); // re-XOR -> bytes 0-9 are plaintext again
1277        cipher.xor_keystream(&mut buf[10..]);
1278
1279        // first 10 bytes are back to plaintext
1280        assert_eq!(
1281            &buf[..10],
1282            &pt[..10],
1283            "set_counter mid-stream: first 10 bytes should be plaintext"
1284        );
1285        // remaining bytes match the one-shot encryption
1286        assert_eq!(&buf[10..], &full_encrypted[10..], "set_counter mid-stream DJB failed");
1287
1288        // IETF variant: same pattern
1289        let nonce12 = test_nonce_12();
1290        let full_encrypted = encrypt_one_shot_ietf(&key, &nonce12, &pt);
1291
1292        let mut buf = pt.clone();
1293        let mut cipher = ChaCha20Ietf::new(&key, &nonce12);
1294        cipher.xor_keystream(&mut buf[..10]);
1295        cipher.set_counter(0);
1296        cipher.xor_keystream(&mut buf[..10]);
1297        cipher.xor_keystream(&mut buf[10..]);
1298        assert_eq!(&buf[..10], &pt[..10], "set_counter mid-stream: IETF first 10 bytes");
1299        assert_eq!(&buf[10..], &full_encrypted[10..], "set_counter mid-stream IETF failed");
1300    }
1301
1302    /// Stress test: many sequential calls with random-sized splits at block boundaries.
1303    #[test]
1304    fn chacha_keystream_leftover_stress_random_splits() {
1305        let key = test_key_32();
1306        let nonce = test_nonce_8();
1307
1308        let sizes = [50usize, 127, 128, 129, 200, 256, 300, 400, 512];
1309        let split_patterns: &[&[usize]] = &[
1310            &[1, 2, 3, 4, 5],
1311            &[7, 13, 23, 31],
1312            &[32, 32, 32],
1313            &[63, 1],
1314            &[64, 64],
1315            &[65, 63],
1316            &[33, 33, 33, 33],
1317            &[10, 10, 10, 10, 10],
1318            &[50, 50, 50],
1319        ];
1320
1321        for &len in &sizes {
1322            let pt = test_plaintext(len);
1323
1324            // ChaCha20
1325            let expected20 = {
1326                let mut b = pt.clone();
1327                ChaCha20Djb::new(&key, &nonce).xor_keystream(&mut b);
1328                b
1329            };
1330            for &splits in split_patterns {
1331                let mut buf = pt.clone();
1332                let mut offset = 0;
1333                let mut c = ChaCha20Djb::new(&key, &nonce);
1334                for &size in splits {
1335                    let end = core::cmp::min(offset + size, buf.len());
1336                    c.xor_keystream(&mut buf[offset..end]);
1337                    offset = end;
1338                    if offset >= buf.len() {
1339                        break;
1340                    }
1341                }
1342                if offset < buf.len() {
1343                    c.xor_keystream(&mut buf[offset..]);
1344                }
1345                assert_eq!(buf, expected20, "Stress ChaCha20 failed for len={len} splits={splits:?}",);
1346            }
1347
1348            // ChaCha12
1349            let expected12 = {
1350                let mut b = pt.clone();
1351                ChaCha12Djb::new(&key, &nonce).xor_keystream(&mut b);
1352                b
1353            };
1354            for &splits in split_patterns {
1355                let mut buf = pt.clone();
1356                let mut offset = 0;
1357                let mut c = ChaCha12Djb::new(&key, &nonce);
1358                for &size in splits {
1359                    let end = core::cmp::min(offset + size, buf.len());
1360                    c.xor_keystream(&mut buf[offset..end]);
1361                    offset = end;
1362                    if offset >= buf.len() {
1363                        break;
1364                    }
1365                }
1366                if offset < buf.len() {
1367                    c.xor_keystream(&mut buf[offset..]);
1368                }
1369                assert_eq!(buf, expected12, "Stress ChaCha12 failed for len={len} splits={splits:?}",);
1370            }
1371
1372            // ChaCha8
1373            let expected8 = {
1374                let mut b = pt.clone();
1375                ChaCha8Djb::new(&key, &nonce).xor_keystream(&mut b);
1376                b
1377            };
1378            for &splits in split_patterns {
1379                let mut buf = pt.clone();
1380                let mut offset = 0;
1381                let mut c = ChaCha8Djb::new(&key, &nonce);
1382                for &size in splits {
1383                    let end = core::cmp::min(offset + size, buf.len());
1384                    c.xor_keystream(&mut buf[offset..end]);
1385                    offset = end;
1386                    if offset >= buf.len() {
1387                        break;
1388                    }
1389                }
1390                if offset < buf.len() {
1391                    c.xor_keystream(&mut buf[offset..]);
1392                }
1393                assert_eq!(buf, expected8, "Stress ChaCha8 failed for len={len} splits={splits:?}",);
1394            }
1395        }
1396    }
1397
1398    /// IETF-specific variant of the stress test.
1399    #[test]
1400    fn chacha_keystream_leftover_stress_random_splits_ietf() {
1401        let key = test_key_32();
1402        let nonce = test_nonce_12();
1403
1404        for &len in &[50usize, 127, 128, 129, 200, 256, 300] {
1405            let pt = test_plaintext(len);
1406            let expected = encrypt_one_shot_ietf(&key, &nonce, &pt);
1407
1408            let split_patterns: &[&[usize]] = &[
1409                &[1, 2, 3, 4, 5],
1410                &[7, 13, 23, 31],
1411                &[63, 1],
1412                &[64, 64],
1413                &[65, 63],
1414                &[33, 33, 33, 33],
1415                &[50, 50, 50],
1416            ];
1417
1418            for &splits in split_patterns {
1419                let mut buf = pt.clone();
1420                let mut offset = 0;
1421                let mut cipher = ChaCha20Ietf::new(&key, &nonce);
1422                for &size in splits {
1423                    let end = core::cmp::min(offset + size, buf.len());
1424                    cipher.xor_keystream(&mut buf[offset..end]);
1425                    offset = end;
1426                    if offset >= buf.len() {
1427                        break;
1428                    }
1429                }
1430                if offset < buf.len() {
1431                    cipher.xor_keystream(&mut buf[offset..]);
1432                }
1433                assert_eq!(buf, expected, "Stress IETF failed for len={len} splits={splits:?}",);
1434            }
1435        }
1436    }
1437
1438    /// XChaCha-specific variant of the stress test.
1439    #[test]
1440    fn keystream_leftover_stress_random_splits_xchacha() {
1441        let key = test_key_32();
1442        let nonce = test_nonce_24();
1443
1444        for &len in &[50usize, 127, 128, 129, 200] {
1445            let pt = test_plaintext(len);
1446            let expected = encrypt_one_shot_xchacha(&key, &nonce, &pt);
1447
1448            let split_patterns: &[&[usize]] = &[&[1, 2, 3, 4, 5], &[7, 13, 23], &[63, 1], &[64, 64], &[65, 63]];
1449
1450            for &splits in split_patterns {
1451                let mut buf = pt.clone();
1452                let mut offset = 0;
1453                let mut cipher = XChaCha20::new(&key, &nonce);
1454                for &size in splits {
1455                    let end = core::cmp::min(offset + size, buf.len());
1456                    cipher.xor_keystream(&mut buf[offset..end]);
1457                    offset = end;
1458                    if offset >= buf.len() {
1459                        break;
1460                    }
1461                }
1462                if offset < buf.len() {
1463                    cipher.xor_keystream(&mut buf[offset..]);
1464                }
1465                assert_eq!(buf, expected, "Stress XChaCha failed for len={len} splits={splits:?}",);
1466            }
1467        }
1468    }
1469
1470    // -------------------------------------------------------------------------
1471    // RFC 8439 counter-wrap / length-limit tests
1472    // -------------------------------------------------------------------------
1473
1474    #[test]
1475    fn chacha20_ietf_max_len_constant() {
1476        // (2^32 - 2) blocks of 64 bytes, i.e. one block of margin below the RFC 8439 counter space.
1477        assert_eq!(CHACHA20_IETF_MAX_LEN, (u32::MAX as u64 - 1) * 64);
1478        assert_eq!(CHACHA20_IETF_MAX_LEN, 274_877_906_816);
1479    }
1480
1481    /// The last usable 32-bit counter value still produces keystream.
1482    #[test]
1483    fn chacha20_ietf_last_counter_block_is_allowed() {
1484        let key = test_key_32();
1485        let nonce = test_nonce_12();
1486        let counter = u32::MAX - 1;
1487
1488        let mut a = [0u8; 64];
1489        let mut cipher = ChaCha20Ietf::new(&key, &nonce);
1490        cipher.set_counter(counter);
1491        cipher.xor_keystream(&mut a);
1492
1493        let mut b = [0u8; 64];
1494        let mut reference = ChaCha20Ietf::new(&key, &nonce);
1495        reference.set_counter(counter);
1496        reference.xor_keystream(&mut b);
1497
1498        assert_eq!(a, b);
1499        assert_ne!(a, [0u8; 64], "keystream must be non-trivial");
1500    }
1501
1502    /// A call that would advance the counter past `u32::MAX` must panic rather than wrap and
1503    /// silently reuse keystream.
1504    #[test]
1505    #[should_panic(expected = "input exceeds maximum allowed")]
1506    fn chacha20_ietf_counter_wrap_panics() {
1507        let key = test_key_32();
1508        let nonce = test_nonce_12();
1509
1510        let mut cipher = ChaCha20Ietf::new(&key, &nonce);
1511        cipher.set_counter(u32::MAX);
1512        let mut buf = [0u8; 64];
1513        cipher.xor_keystream(&mut buf);
1514    }
1515
1516    /// Streaming across the end of the counter space: partially consuming leftover must not be
1517    /// mistaken for needing a fresh (wrapping) block.
1518    #[test]
1519    fn chacha20_ietf_leftover_near_counter_end_does_not_panic() {
1520        let key = test_key_32();
1521        let nonce = test_nonce_12();
1522
1523        let mut cipher = ChaCha20Ietf::new(&key, &nonce);
1524        cipher.set_counter(u32::MAX - 1);
1525
1526        // First call consumes block u32::MAX - 1 and buffers 54 bytes of leftover.
1527        let mut first = [0u8; 10];
1528        cipher.xor_keystream(&mut first);
1529
1530        // Second call is fully served from the leftover: no new block, no panic.
1531        let mut second = [0u8; 5];
1532        cipher.xor_keystream(&mut second);
1533
1534        let mut expected_first = [0u8; 10];
1535        let mut expected_second = [0u8; 5];
1536        let mut reference = ChaCha20Ietf::new(&key, &nonce);
1537        reference.set_counter(u32::MAX - 1);
1538        reference.xor_keystream(&mut expected_first);
1539        reference.xor_keystream(&mut expected_second);
1540        assert_eq!(first, expected_first);
1541        assert_eq!(second, expected_second);
1542    }
1543
1544    /// Once the leftover is exhausted, the next call needs a block and must panic.
1545    #[test]
1546    #[should_panic(expected = "input exceeds maximum allowed")]
1547    fn chacha20_ietf_wrap_after_leftover_panics() {
1548        let key = test_key_32();
1549        let nonce = test_nonce_12();
1550
1551        let mut cipher = ChaCha20Ietf::new(&key, &nonce);
1552        cipher.set_counter(u32::MAX - 1);
1553
1554        let mut first = [0u8; 10];
1555        cipher.xor_keystream(&mut first);
1556        let mut second = [0u8; 5];
1557        cipher.xor_keystream(&mut second);
1558
1559        // 5 bytes of leftover remain; asking for more requires the wrapped block.
1560        let mut third = [0u8; 64];
1561        cipher.xor_keystream(&mut third);
1562    }
1563
1564    /// XChaCha20 wraps the IETF ChaCha20 stream and must enforce the same limit.
1565    #[test]
1566    #[should_panic(expected = "input exceeds maximum allowed")]
1567    fn xchacha20_counter_wrap_panics() {
1568        let key = test_key_32();
1569        let nonce = test_nonce_24();
1570
1571        let mut cipher = XChaCha20::new(&key, &nonce);
1572        cipher.set_counter(u32::MAX);
1573        let mut buf = [0u8; 64];
1574        cipher.xor_keystream(&mut buf);
1575    }
1576}