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/// The "sigma" constant which is the value of the first row of ChaCha's state.
13pub(crate) const CONSTANT: [u32; 4] = [
14    0x61707865, // "expa"
15    0x3320646e, // "nd 3"
16    0x79622d32, // "2-by"
17    0x6b206574, // "te k"
18];
19
20pub type ChaCha8Djb = ChaCha<8, false>;
21pub type ChaCha12Djb = ChaCha<12, false>;
22pub type ChaCha20Djb = ChaCha<20, false>;
23pub type ChaCha20Ietf = ChaCha<20, true>;
24pub type XChaCha20 = XChaCha<20>;
25
26/// ChaCha stream cipher.
27///
28/// `IS_IETF` selects the nonce/counter layout:
29/// - `false` (DJB original): 64-bit counter at words 12–13, 64-bit nonce at words 14–15.
30/// - `true` (IETF / RFC 8439): 32-bit counter at word 12, 96-bit nonce at words 13–15.
31#[cfg_attr(feature = "zeroize", derive(Zeroize, ZeroizeOnDrop))]
32pub struct ChaCha<const ROUNDS: usize, const IS_IETF: bool> {
33    state: [u32; STATE_WORDS],
34    /// ChaCha is a stream cipher that works with 64-byte blocks.
35    /// It means that consumers of this packages should be able to call `xor_keystream` multiple
36    /// times even if there input is not aligned with ChaCha blocks.
37    /// Thus calling multiple times `xor_keystream`:
38    /// xor_keystream(plaintext[0..3]), xor_keystream(plaintext[3..50]), xor_keystream(plaintext[50..150]);
39    /// Should be equal to calling it only once:
40    /// xor_keystream(plaintext[0..150]);
41    /// For that, we keep the last computed keystream block, as well as an offset indicating where
42    /// the unconsumed tail starts.
43    /// The full leftover block is stored in `keystream_leftover` minus 1 byte because if there is leftover
44    /// it means that the leftover is <= (BLOCK_SIZE - 1).
45    /// When `keystream_leftover_offset == (BLOCK_SIZE - 1)`, there is no leftover (empty slice).
46    /// NOTE: the `keystream_leftover` buffer is valid only if the previous call to `xor_keystream` had
47    /// an `input.len() % 64 != 0`, Otherwise there is no need to preserve the last keystream block.
48    keystream_leftover: [u8; BLOCK_SIZE - 1],
49    keystream_leftover_offset: u8,
50}
51
52impl<const ROUNDS: usize, const IS_IETF: bool> ChaCha<ROUNDS, IS_IETF> {
53    #[inline(always)]
54    fn extract_counter(state: &[u32; STATE_WORDS]) -> u64 {
55        if IS_IETF {
56            state[12] as u64
57        } else {
58            ((state[13] as u64) << 32) | (state[12] as u64)
59        }
60    }
61
62    #[inline(always)]
63    fn inject_counter(state: &mut [u32; STATE_WORDS], counter: u64) {
64        state[12] = counter as u32;
65        if !IS_IETF {
66            state[13] = (counter >> 32) as u32;
67        }
68    }
69}
70
71impl<const ROUNDS: usize> ChaCha<ROUNDS, false> {
72    /// Create a new ChaCha instance with the DJB nonce layout (8-byte nonce, 64-bit counter).
73    pub fn new(key: &[u8; 32], nonce: &[u8; 8]) -> ChaCha<ROUNDS, false> {
74        let mut state = [0u32; STATE_WORDS];
75
76        state[..4].copy_from_slice(&CONSTANT);
77
78        for (state_word, key_chunk) in state[4..12].iter_mut().zip(key.chunks_exact(4)) {
79            *state_word = u32::from_le_bytes(key_chunk.try_into().unwrap());
80        }
81
82        state[14] = u32::from_le_bytes(nonce[0..4].try_into().unwrap());
83        state[15] = u32::from_le_bytes(nonce[4..8].try_into().unwrap());
84
85        return ChaCha {
86            state,
87            keystream_leftover: [0u8; BLOCK_SIZE - 1],
88            keystream_leftover_offset: (BLOCK_SIZE - 1) as u8,
89        };
90    }
91
92    /// Set the ChaCha counter (words 12 and 13). It can be used to move forward and backward in the
93    /// keystream.
94    #[inline(always)]
95    pub fn set_counter(&mut self, counter: u64) {
96        Self::inject_counter(&mut self.state, counter);
97        self.keystream_leftover_offset = (BLOCK_SIZE - 1) as u8;
98    }
99}
100
101impl<const ROUNDS: usize> ChaCha<ROUNDS, true> {
102    /// Create a new ChaCha instance with the IETF (RFC 8439) nonce layout (12-byte nonce, 32-bit counter).
103    pub fn new(key: &[u8; 32], nonce: &[u8; 12]) -> ChaCha<ROUNDS, true> {
104        let mut state = [0u32; STATE_WORDS];
105
106        state[..4].copy_from_slice(&CONSTANT);
107
108        for (state_word, key_chunk) in state[4..12].iter_mut().zip(key.chunks_exact(4)) {
109            *state_word = u32::from_le_bytes(key_chunk.try_into().unwrap());
110        }
111
112        state[13] = u32::from_le_bytes(nonce[0..4].try_into().unwrap());
113        state[14] = u32::from_le_bytes(nonce[4..8].try_into().unwrap());
114        state[15] = u32::from_le_bytes(nonce[8..12].try_into().unwrap());
115
116        return ChaCha {
117            state,
118            keystream_leftover: [0u8; BLOCK_SIZE - 1],
119            keystream_leftover_offset: (BLOCK_SIZE - 1) as u8,
120        };
121    }
122
123    /// Set the ChaCha counter (word 12). The counter is a u32. It can be used to move forward
124    /// and backward in the keystream.
125    #[inline(always)]
126    pub fn set_counter(&mut self, counter: u32) {
127        Self::inject_counter(&mut self.state, counter as u64);
128        self.keystream_leftover_offset = (BLOCK_SIZE - 1) as u8;
129    }
130}
131
132impl<const ROUNDS: usize, const IS_IETF: bool> StreamCipher for ChaCha<ROUNDS, IS_IETF> {
133    /// XOR `plaintext` with the ChaCha keystream.
134    fn xor_keystream(&mut self, mut in_out: &mut [u8]) {
135        if in_out.len() == 0 {
136            return;
137        }
138
139        // first, consume the keystream leftover, if any
140        if self.keystream_leftover_offset < (BLOCK_SIZE - 1) as u8 {
141            let keystream_leftover = &self.keystream_leftover[(self.keystream_leftover_offset as usize)..];
142
143            in_out
144                .iter_mut()
145                .zip(keystream_leftover)
146                .for_each(|(plaintext, keystream)| *plaintext ^= *keystream);
147
148            if in_out.len() > keystream_leftover.len() {
149                in_out = &mut in_out[keystream_leftover.len()..];
150            } else if in_out.len() < keystream_leftover.len() {
151                self.keystream_leftover_offset += in_out.len() as u8;
152                return;
153            } else {
154                // in_out.len() == keystream_leftover.len() -> in_out has consumed exactly all the
155                // leftover keystream
156                self.keystream_leftover_offset = (BLOCK_SIZE - 1) as u8;
157                return;
158            }
159        }
160        // at this point, we already know how many bytes of leftover there will be
161        self.keystream_leftover_offset = ((in_out.len() + BLOCK_SIZE - 1) % BLOCK_SIZE) as u8;
162
163        // runtime detection of CPU features for x86 and x86_64 when the "std" feature is enabled
164        #[cfg(all(feature = "std", target_arch = "x86_64"))]
165        if in_out.len() >= 128 && is_x86_feature_detected!("avx512f") {
166            use super::chacha_avx512::chacha_avx512;
167            unsafe { chacha_avx512::<ROUNDS, IS_IETF>(&mut self.state, in_out, &mut self.keystream_leftover) };
168            return;
169        }
170
171        // compile-time dispatch
172        #[cfg(all(not(feature = "std"), target_arch = "x86_64", target_feature = "avx512f"))]
173        if in_out.len() >= 128 {
174            use super::chacha_avx512::chacha_avx512;
175            unsafe { chacha_avx512::<ROUNDS, IS_IETF>(&mut self.state, in_out, &mut self.keystream_leftover) };
176            return;
177        }
178
179        #[cfg(all(target_arch = "x86_64", target_feature = "avx2"))]
180        if in_out.len() >= 128 {
181            use super::chacha_avx2::chacha_avx2;
182            unsafe { chacha_avx2::<ROUNDS, IS_IETF>(&mut self.state, in_out, &mut self.keystream_leftover) };
183            return;
184        }
185
186        // aarch64 assumes that NEON is always available so compile-time dispatch is enough
187        #[cfg(all(target_arch = "aarch64", target_feature = "neon"))]
188        if in_out.len() >= 128 {
189            use super::chacha_neon::chacha_neon;
190            unsafe { chacha_neon::<ROUNDS, IS_IETF>(&mut self.state, in_out, &mut self.keystream_leftover) };
191            return;
192        }
193
194        // wasm32 only supports compile-time features detection
195        #[cfg(all(target_arch = "wasm32", target_feature = "simd128"))]
196        if in_out.len() >= 128 {
197            use super::chacha_wasm_simd128::chacha_wasm_simd128;
198            chacha_wasm_simd128::<ROUNDS, IS_IETF>(&mut self.state, in_out, &mut self.keystream_leftover);
199            return;
200        }
201
202        // fallback for when SIMD acceleration is not available
203        chacha_generic::<ROUNDS, IS_IETF>(&mut self.state, &mut self.keystream_leftover, in_out);
204    }
205}
206
207#[inline]
208fn chacha_generic<const ROUNDS: usize, const IS_IETF: bool>(
209    mut state: &mut [u32; STATE_WORDS],
210    keystream_leftover: &mut [u8; BLOCK_SIZE - 1],
211    plaintext: &mut [u8],
212) {
213    let mut keystream = [0u8; BLOCK_SIZE];
214    let keystream_ptr = keystream.as_mut_ptr();
215    let mut counter = ChaCha::<ROUNDS, IS_IETF>::extract_counter(state);
216
217    // process the input by blocks of 64 bytes
218    for plaintext_block in plaintext.chunks_mut(BLOCK_SIZE) {
219        ChaCha::<ROUNDS, IS_IETF>::inject_counter(&mut state, counter);
220
221        // prepare temporary (working) state
222        let mut tmp_state = *state;
223
224        // perform the ROUNDS / 2 double rounds e.g. 10 double rounds for ChaCha20
225        for _ in 0..(ROUNDS / 2) {
226            // column rounds
227            quarter_round(&mut tmp_state, 0, 4, 8, 12);
228            quarter_round(&mut tmp_state, 1, 5, 9, 13);
229            quarter_round(&mut tmp_state, 2, 6, 10, 14);
230            quarter_round(&mut tmp_state, 3, 7, 11, 15);
231
232            // diagonal rounds
233            quarter_round(&mut tmp_state, 0, 5, 10, 15);
234            quarter_round(&mut tmp_state, 1, 6, 11, 12);
235            quarter_round(&mut tmp_state, 2, 7, 8, 13);
236            quarter_round(&mut tmp_state, 3, 4, 9, 14);
237        }
238
239        // add initial state to tmp_state to generate the keystream and "serialize" it to little endian
240        // for (tmp_word, state_word) in tmp_state.iter_mut().zip(state.iter()) {
241        //     *tmp_word = tmp_word.wrapping_add(*state_word).to_le();
242        // }
243        for word_index in 0..STATE_WORDS {
244            // first we add the initial state to the working state to get the keystream
245            tmp_state[word_index] = tmp_state[word_index].wrapping_add(state[word_index]);
246
247            // then we serialize the keystream
248            // SAFETY: this is safe because `tmp_state` and `keystream` both have fixed, known size.
249            // We are just merely converting [u32; STATE_WORDS] to [u8; BLOCK_SIZE] with the correct
250            // endianness.
251            unsafe {
252                core::ptr::copy_nonoverlapping(
253                    tmp_state[word_index].to_le_bytes().as_ptr(),
254                    keystream_ptr.add(word_index * 4),
255                    4,
256                );
257            }
258        }
259
260        // XOR plaintext with keystream
261        plaintext_block
262            .iter_mut()
263            .zip(keystream)
264            .for_each(|(plaintext, keystream)| *plaintext ^= keystream);
265
266        counter = counter.wrapping_add(1);
267    }
268
269    ChaCha::<ROUNDS, IS_IETF>::inject_counter(state, counter);
270
271    if plaintext.len() % BLOCK_SIZE != 0 {
272        // copy the last 63 bytes of the leftover keystream block
273        keystream_leftover.copy_from_slice(&keystream[1..]);
274    }
275}
276
277#[inline(always)]
278pub(crate) const fn quarter_round(state: &mut [u32; 16], a: usize, b: usize, c: usize, d: usize) {
279    // a += b; d ^= a; d <<<= 16
280    state[a] = state[a].wrapping_add(state[b]);
281    state[d] ^= state[a];
282    state[d] = state[d].rotate_left(16);
283
284    // c += d; b ^= c; b <<<= 12
285    state[c] = state[c].wrapping_add(state[d]);
286    state[b] ^= state[c];
287    state[b] = state[b].rotate_left(12);
288
289    // a += b; d ^= a; d <<<= 8
290    state[a] = state[a].wrapping_add(state[b]);
291    state[d] ^= state[a];
292    state[d] = state[d].rotate_left(8);
293
294    // c += d; b ^= c; b <<<= 7
295    state[c] = state[c].wrapping_add(state[d]);
296    state[b] ^= state[c];
297    state[b] = state[b].rotate_left(7);
298}
299
300/// XChaCha20 stream cipher with 24-byte (192-bit) nonce (draft-irtf-cfrg-xchacha-03).
301///
302/// Internally uses HChaCha20 to derive a subkey from the first 16 bytes of the nonce,
303/// then uses the IETF ChaCha20 variant with the remaining 8 nonce bytes.
304#[cfg_attr(feature = "zeroize", derive(Zeroize, ZeroizeOnDrop))]
305pub struct XChaCha<const ROUNDS: usize> {
306    inner: ChaCha<ROUNDS, true>,
307}
308
309impl<const ROUNDS: usize> XChaCha<ROUNDS> {
310    /// Creates a new XChaCha stream cipher from a 32-byte key and a 24-byte nonce.
311    ///
312    /// The first 16 bytes of the nonce are used with HChaCha20 to derive a subkey.
313    /// The last 8 bytes of the nonce become the ChaCha20 IETF nonce (prefixed with 4 zero bytes).
314    pub fn new(key: &[u8; 32], nonce: &[u8; 24]) -> XChaCha<ROUNDS> {
315        let subkey = super::hchacha20(key, nonce[..16].try_into().unwrap());
316        let mut ietf_nonce = [0u8; 12];
317        ietf_nonce[4..12].copy_from_slice(&nonce[16..24]);
318        return XChaCha {
319            inner: ChaCha::<ROUNDS, true>::new(&subkey, &ietf_nonce),
320        };
321    }
322
323    /// Sets the ChaCha20 block counter. Counter is a u32.
324    pub fn set_counter(&mut self, counter: u32) {
325        self.inner.set_counter(counter);
326    }
327}
328
329impl<const ROUNDS: usize> StreamCipher for XChaCha<ROUNDS> {
330    fn xor_keystream(&mut self, in_out: &mut [u8]) {
331        self.inner.xor_keystream(in_out);
332    }
333}
334
335#[cfg(test)]
336mod test {
337    use super::{ChaCha8Djb, ChaCha12Djb, ChaCha20Djb, ChaCha20Ietf, XChaCha20};
338    use crate::StreamCipher;
339
340    struct TestDjb {
341        key: [u8; 32],
342        nonce: [u8; 8],
343        initial_counter: u64,
344        plaintext: Vec<u8>,
345        expected_ciphertext: Vec<u8>,
346    }
347
348    #[test]
349    fn chacha20_test_vectors_djb() {
350        let tests = vec![
351            // https://www.rfc-editor.org/rfc/rfc8439#section-2.4.2
352            TestDjb {
353                key: hex::decode("000102030405060708090a0b0c0d0e0f101112131415161718191a1b1c1d1e1f")
354                    .unwrap()
355                    .try_into()
356                    .unwrap(),
357                nonce: hex::decode("0000004a00000000").unwrap().try_into().unwrap(),
358                initial_counter: 1,
359                plaintext: hex::decode(
360                    "4c616469657320616e642047656e746c\
361656d656e206f662074686520636c6173\
36273206f66202739393a20496620492063\
3636f756c64206f6666657220796f75206f\
3646e6c79206f6e652074697020666f7220\
365746865206675747572652c2073756e73\
366637265656e20776f756c642062652069\
367742e",
368                )
369                .unwrap(),
370                expected_ciphertext: hex::decode(
371                    "6e2e359a2568f98041ba0728dd0d6981\
372e97e7aec1d4360c20a27afccfd9fae0b\
373f91b65c5524733ab8f593dabcd62b357\
3741639d624e65152ab8f530c359f0861d8\
37507ca0dbf500d6a6156a38e088a22b65e\
37652bc514d16ccf806818ce91ab7793736\
3775af90bbf74a35be6b40b8eedf2785e42\
378874d",
379                )
380                .unwrap(),
381            },
382            // https://www.rfc-editor.org/rfc/rfc8439#appendix-A.2 Test vector #1
383            TestDjb {
384                key: [0u8; 32],
385                nonce: [0u8; 8],
386                initial_counter: 0,
387                plaintext: [0u8; 64].to_vec(),
388                expected_ciphertext: hex::decode(
389                    "76b8e0ada0f13d90405d6ae55386bd28\
390bdd219b8a08ded1aa836efcc8b770dc7\
391da41597c5157488d7724e03fb8d84a37\
3926a43b8f41518a11cc387b669b2ee6586",
393                )
394                .unwrap(),
395            },
396            // https://www.rfc-editor.org/rfc/rfc8439#appendix-A.2 Test Vector #2
397            TestDjb {
398                key: hex::decode("0000000000000000000000000000000000000000000000000000000000000001")
399                    .unwrap()
400                    .try_into()
401                    .unwrap(),
402                nonce: hex::decode("0000000000000002").unwrap().try_into().unwrap(),
403                initial_counter: 1,
404                plaintext: hex::decode(
405                    "416e79207375626d697373696f6e2074\
4066f20746865204945544620696e74656e\
4076465642062792074686520436f6e7472\
408696275746f7220666f72207075626c69\
409636174696f6e20617320616c6c206f72\
4102070617274206f6620616e2049455446\
41120496e7465726e65742d447261667420\
4126f722052464320616e6420616e792073\
413746174656d656e74206d616465207769\
4147468696e2074686520636f6e74657874\
415206f6620616e20494554462061637469\
4167669747920697320636f6e7369646572\
417656420616e20224945544620436f6e74\
4187269627574696f6e222e205375636820\
41973746174656d656e747320696e636c75\
4206465206f72616c2073746174656d656e\
421747320696e2049455446207365737369\
4226f6e732c2061732077656c6c20617320\
4237772697474656e20616e6420656c6563\
42474726f6e696320636f6d6d756e696361\
42574696f6e73206d61646520617420616e\
426792074696d65206f7220706c6163652c\
42720776869636820617265206164647265\
4287373656420746f",
429                )
430                .unwrap(),
431                expected_ciphertext: hex::decode(
432                    "a3fbf07df3fa2fde4f376ca23e827370\
43341605d9f4f4f57bd8cff2c1d4b7955ec\
4342a97948bd3722915c8f3d337f7d37005\
4350e9e96d647b7c39f56e031ca5eb6250d\
4364042e02785ececfa4b4bb5e8ead0440e\
43720b6e8db09d881a7c6132f420e527950\
43842bdfa7773d8a9051447b3291ce1411c\
439680465552aa6c405b7764d5e87bea85a\
440d00f8449ed8f72d0d662ab052691ca66\
441424bc86d2df80ea41f43abf937d3259d\
442c4b2d0dfb48a6c9139ddd7f76966e928\
443e635553ba76c5c879d7b35d49eb2e62b\
4440871cdac638939e25e8a1e0ef9d5280f\
445a8ca328b351c3c765989cbcf3daa8b6c\
446cc3aaf9f3979c92b3720fc88dc95ed84\
447a1be059c6499b9fda236e7e818b04b0b\
448c39c1e876b193bfe5569753f88128cc0\
4498aaa9b63d1a16f80ef2554d7189c411f\
4505869ca52c5b83fa36ff216b9c1d30062\
451bebcfd2dc5bce0911934fda79a86f6e6\
45298ced759c3ff9b6477338f3da4f9cd85\
45314ea9982ccafb341b2384dd902f3d1ab\
4547ac61dd29c6f21ba5b862f3730e37cfd\
455c4fd806c22f221",
456                )
457                .unwrap(),
458            },
459            // https://www.rfc-editor.org/rfc/rfc8439#appendix-A.2 Test Vector #3
460            TestDjb {
461                key: hex::decode("1c9240a5eb55d38af333888604f6b5f0473917c1402b80099dca5cbc207075c0")
462                    .unwrap()
463                    .try_into()
464                    .unwrap(),
465                nonce: hex::decode("0000000000000002").unwrap().try_into().unwrap(),
466                initial_counter: 42,
467                plaintext: hex::decode(
468                    "2754776173206272696c6c69672c2061\
4696e642074686520736c6974687920746f\
4707665730a446964206779726520616e64\
4712067696d626c6520696e207468652077\
4726162653a0a416c6c206d696d73792077\
4736572652074686520626f726f676f7665\
474732c0a416e6420746865206d6f6d6520\
4757261746873206f757467726162652e",
476                )
477                .unwrap(),
478                expected_ciphertext: hex::decode(
479                    "62e6347f95ed87a45ffae7426f27a1df\
4805fb69110044c0d73118effa95b01e5cf\
481166d3df2d721caf9b21e5fb14c616871\
482fd84c54f9d65b283196c7fe4f60553eb\
483f39c6402c42234e32a356b3e764312a6\
4841a5532055716ead6962568f87d3f3f77\
48504c6a8d1bcd1bf4d50d6154b6da731b1\
48687b58dfd728afa36757a797ac188d1",
487                )
488                .unwrap(),
489            },
490        ];
491
492        for (i, test) in tests.into_iter().enumerate() {
493            let mut cipher = ChaCha20Djb::new(&test.key, &test.nonce);
494            cipher.set_counter(test.initial_counter);
495
496            let mut plaintext = test.plaintext.clone();
497            cipher.xor_keystream(&mut plaintext);
498
499            assert_eq!(
500                plaintext,
501                test.expected_ciphertext,
502                "test [{i}] failed
503Got ciphertext: {}
504Expected ciphertext: {}",
505                hex::encode(&plaintext),
506                hex::encode(&test.expected_ciphertext),
507            );
508
509            let mut cipher = ChaCha20Djb::new(&test.key, &test.nonce);
510            cipher.set_counter(test.initial_counter);
511            cipher.xor_keystream(&mut plaintext);
512
513            assert_eq!(
514                plaintext,
515                test.plaintext,
516                "test [{i}] failed. Initial plaintext != decrypt(encrypt(plaintext))
517Got: {}
518Expected: {}",
519                hex::encode(&plaintext),
520                hex::encode(&test.plaintext),
521            );
522
523            // ensure that the encryption is correct even for plaintexts that are not % 64 (block size)
524            // thus:
525            // cipher.xor_keystream(plaintext[0..10])
526            // cipher.xor_keystream(plaintext[10..30])
527            // cipher.xor_keystream(plaintext[30..5])
528            // should be equal to:
529            // cipher.xor_keystream(plaintext[0..35])
530
531            let mut cipher = ChaCha20Djb::new(&test.key, &test.nonce);
532            cipher.xor_keystream(&mut plaintext);
533            for n in 0..10 {
534                let mut partial_plaintext: Vec<u8> = test.plaintext.clone();
535
536                let mut cipher = ChaCha20Djb::new(&test.key, &test.nonce);
537                cipher.xor_keystream(&mut partial_plaintext[..n]);
538                cipher.xor_keystream(&mut partial_plaintext[n..]);
539
540                assert_eq!(
541                    plaintext,
542                    partial_plaintext,
543                    "test [{i}] failed. partial encryption is not valid for n = {n}
544            Got: {}
545            Expected: {}",
546                    hex::encode(&partial_plaintext),
547                    hex::encode(&plaintext),
548                )
549            }
550        }
551    }
552
553    #[test]
554    fn chacha20_keystream_leftover_multi_call() {
555        // Regression: verifies that leftovers are correctly consumed across 3+
556        // partial calls. The old length-based approach did not compact the array
557        // after partial consumption, causing stale keystream reuse on subsequent calls.
558        let key = hex::decode("000102030405060708090a0b0c0d0e0f101112131415161718191a1b1c1d1e1f")
559            .unwrap()
560            .try_into()
561            .unwrap();
562        let nonce = hex::decode("0000004a00000000").unwrap().try_into().unwrap();
563        let plaintext = hex::decode(
564            "4c616469657320616e642047656e746c\
565656d656e206f662074686520636c6173\
56673206f66202739393a20496620492063\
5676f756c64206f6666657220796f75206f\
5686e6c79206f6e652074697020666f7220\
569746865206675747572652c2073756e73\
570637265656e20776f756c642062652069\
571742e",
572        )
573        .unwrap();
574
575        let mut expected = plaintext.clone();
576        ChaCha20Djb::new(&key, &nonce).xor_keystream(&mut expected);
577
578        // call 1: partial block -> leaves leftover
579        // call 2: partially consumes leftover
580        // call 3: consumes remaining leftover + fresh blocks
581        {
582            let mut buf = plaintext.clone();
583            let mut cipher = ChaCha20Djb::new(&key, &nonce);
584            cipher.xor_keystream(&mut buf[..10]);
585            cipher.xor_keystream(&mut buf[10..15]);
586            cipher.xor_keystream(&mut buf[15..]);
587            assert_eq!(buf, expected, "partial leftover consumption");
588        }
589
590        // call 1: partial block -> leaves leftover
591        // call 2: exactly exhausts leftover
592        // call 3: fresh blocks
593        {
594            let mut buf = plaintext.clone();
595            let mut cipher = ChaCha20Djb::new(&key, &nonce);
596            cipher.xor_keystream(&mut buf[..10]);
597            cipher.xor_keystream(&mut buf[10..64]);
598            cipher.xor_keystream(&mut buf[64..]);
599            assert_eq!(buf, expected, "exact leftover exhaustion");
600        }
601
602        // call 1: partial block -> leaves leftover
603        // call 2 + call 3: two rounds of partial consumption
604        // call 4: consumes remaining leftover + fresh blocks
605        {
606            let mut buf = plaintext.clone();
607            let mut cipher = ChaCha20Djb::new(&key, &nonce);
608            cipher.xor_keystream(&mut buf[..8]);
609            cipher.xor_keystream(&mut buf[8..13]);
610            cipher.xor_keystream(&mut buf[13..33]);
611            cipher.xor_keystream(&mut buf[33..]);
612            assert_eq!(buf, expected, "multiple partial leftover consumptions");
613        }
614    }
615
616    #[test]
617    fn chacha20_ietf_test_vectors() {
618        // IETF ChaCha20 test vectors from RFC 8439 Appendix A.2.
619        // These use the IETF layout: 32-bit counter (word 12), 96-bit nonce (words 13–15).
620        // The nonce is the full 12-byte value (little-endian) that RFC 8439 places in words 13–15.
621
622        struct TestIetf {
623            key: [u8; 32],
624            nonce: [u8; 12],
625            initial_counter: u32,
626            plaintext: Vec<u8>,
627            expected_ciphertext: Vec<u8>,
628        }
629
630        let tests = vec![
631            // RFC 8439 section 2.4.2 — the "Ladies and Gentlemen" test vector
632            // IETF state: word 12 = counter, words 13-15 = 96-bit nonce LE
633            // So state[13]=0, state[14]=0x4a000000, state[15]=0
634            TestIetf {
635                key: hex::decode("000102030405060708090a0b0c0d0e0f101112131415161718191a1b1c1d1e1f")
636                    .unwrap()
637                    .try_into()
638                    .unwrap(),
639                nonce: hex::decode("000000000000004a00000000").unwrap().try_into().unwrap(),
640                initial_counter: 1,
641                plaintext: hex::decode(
642                    "4c616469657320616e642047656e746c\
643656d656e206f662074686520636c6173\
64473206f66202739393a20496620492063\
6456f756c64206f6666657220796f75206f\
6466e6c79206f6e652074697020666f7220\
647746865206675747572652c2073756e73\
648637265656e20776f756c642062652069\
649742e",
650                )
651                .unwrap(),
652                expected_ciphertext: hex::decode(
653                    "6e2e359a2568f98041ba0728dd0d6981\
654e97e7aec1d4360c20a27afccfd9fae0b\
655f91b65c5524733ab8f593dabcd62b357\
6561639d624e65152ab8f530c359f0861d8\
65707ca0dbf500d6a6156a38e088a22b65e\
65852bc514d16ccf806818ce91ab7793736\
6595af90bbf74a35be6b40b8eedf2785e42\
660874d",
661                )
662                .unwrap(),
663            },
664            // RFC 8439 Appendix A.2 Vector #1 — all zeros key, nonce, counter=0
665            TestIetf {
666                key: [0u8; 32],
667                nonce: [0u8; 12],
668                initial_counter: 0,
669                plaintext: [0u8; 64].to_vec(),
670                expected_ciphertext: hex::decode(
671                    "76b8e0ada0f13d90405d6ae55386bd28\
672bdd219b8a08ded1aa836efcc8b770dc7\
673da41597c5157488d7724e03fb8d84a37\
6746a43b8f41518a11cc387b669b2ee6586",
675                )
676                .unwrap(),
677            },
678            // RFC 8439 Appendix A.2 Vector #2 — counter=1
679            TestIetf {
680                key: hex::decode("0000000000000000000000000000000000000000000000000000000000000001")
681                    .unwrap()
682                    .try_into()
683                    .unwrap(),
684                nonce: hex::decode("000000000000000000000002").unwrap().try_into().unwrap(),
685                initial_counter: 1,
686                plaintext: hex::decode(
687                    "416e79207375626d697373696f6e2074\
6886f20746865204945544620696e74656e\
6896465642062792074686520436f6e7472\
690696275746f7220666f72207075626c69\
691636174696f6e20617320616c6c206f72\
6922070617274206f6620616e2049455446\
69320496e7465726e65742d447261667420\
6946f722052464320616e6420616e792073\
695746174656d656e74206d616465207769\
6967468696e2074686520636f6e74657874\
697206f6620616e20494554462061637469\
6987669747920697320636f6e7369646572\
699656420616e20224945544620436f6e74\
7007269627574696f6e222e205375636820\
70173746174656d656e747320696e636c75\
7026465206f72616c2073746174656d656e\
703747320696e2049455446207365737369\
7046f6e732c2061732077656c6c20617320\
7057772697474656e20616e6420656c6563\
70674726f6e696320636f6d6d756e696361\
70774696f6e73206d61646520617420616e\
708792074696d65206f7220706c6163652c\
70920776869636820617265206164647265\
7107373656420746f",
711                )
712                .unwrap(),
713                expected_ciphertext: hex::decode(
714                    "a3fbf07df3fa2fde4f376ca23e827370\
71541605d9f4f4f57bd8cff2c1d4b7955ec\
7162a97948bd3722915c8f3d337f7d37005\
7170e9e96d647b7c39f56e031ca5eb6250d\
7184042e02785ececfa4b4bb5e8ead0440e\
71920b6e8db09d881a7c6132f420e527950\
72042bdfa7773d8a9051447b3291ce1411c\
721680465552aa6c405b7764d5e87bea85a\
722d00f8449ed8f72d0d662ab052691ca66\
723424bc86d2df80ea41f43abf937d3259d\
724c4b2d0dfb48a6c9139ddd7f76966e928\
725e635553ba76c5c879d7b35d49eb2e62b\
7260871cdac638939e25e8a1e0ef9d5280f\
727a8ca328b351c3c765989cbcf3daa8b6c\
728cc3aaf9f3979c92b3720fc88dc95ed84\
729a1be059c6499b9fda236e7e818b04b0b\
730c39c1e876b193bfe5569753f88128cc0\
7318aaa9b63d1a16f80ef2554d7189c411f\
7325869ca52c5b83fa36ff216b9c1d30062\
733bebcfd2dc5bce0911934fda79a86f6e6\
73498ced759c3ff9b6477338f3da4f9cd85\
73514ea9982ccafb341b2384dd902f3d1ab\
7367ac61dd29c6f21ba5b862f3730e37cfd\
737c4fd806c22f221",
738                )
739                .unwrap(),
740            },
741            // RFC 8439 Appendix A.2 Vector #3 — counter=42
742            TestIetf {
743                key: hex::decode("1c9240a5eb55d38af333888604f6b5f0473917c1402b80099dca5cbc207075c0")
744                    .unwrap()
745                    .try_into()
746                    .unwrap(),
747                nonce: hex::decode("000000000000000000000002").unwrap().try_into().unwrap(),
748                initial_counter: 42,
749                plaintext: hex::decode(
750                    "2754776173206272696c6c69672c2061\
7516e642074686520736c6974687920746f\
7527665730a446964206779726520616e64\
7532067696d626c6520696e207468652077\
7546162653a0a416c6c206d696d73792077\
7556572652074686520626f726f676f7665\
756732c0a416e6420746865206d6f6d6520\
7577261746873206f757467726162652e",
758                )
759                .unwrap(),
760                expected_ciphertext: hex::decode(
761                    "62e6347f95ed87a45ffae7426f27a1df\
7625fb69110044c0d73118effa95b01e5cf\
763166d3df2d721caf9b21e5fb14c616871\
764fd84c54f9d65b283196c7fe4f60553eb\
765f39c6402c42234e32a356b3e764312a6\
7661a5532055716ead6962568f87d3f3f77\
76704c6a8d1bcd1bf4d50d6154b6da731b1\
76887b58dfd728afa36757a797ac188d1",
769                )
770                .unwrap(),
771            },
772        ];
773
774        for (i, test) in tests.into_iter().enumerate() {
775            let mut cipher = ChaCha20Ietf::new(&test.key, &test.nonce);
776            cipher.set_counter(test.initial_counter);
777
778            let mut plaintext = test.plaintext.clone();
779            cipher.xor_keystream(&mut plaintext);
780
781            assert_eq!(
782                plaintext,
783                test.expected_ciphertext,
784                "ietf test [{i}] failed
785Got ciphertext: {}
786Expected ciphertext: {}",
787                hex::encode(&plaintext),
788                hex::encode(&test.expected_ciphertext),
789            );
790
791            // decrypt
792            let mut cipher = ChaCha20Ietf::new(&test.key, &test.nonce);
793            cipher.set_counter(test.initial_counter);
794            cipher.xor_keystream(&mut plaintext);
795
796            assert_eq!(
797                plaintext,
798                test.plaintext,
799                "ietf test [{i}] failed. Initial plaintext != decrypt(encrypt(plaintext))
800Got: {}
801Expected: {}",
802                hex::encode(&plaintext),
803                hex::encode(&test.plaintext),
804            );
805
806            // partial encryption check
807            let mut cipher = ChaCha20Ietf::new(&test.key, &test.nonce);
808            cipher.set_counter(test.initial_counter);
809            cipher.xor_keystream(&mut plaintext);
810            for n in 0..10 {
811                let mut partial_plaintext: Vec<u8> = test.plaintext.clone();
812
813                let mut cipher = ChaCha20Ietf::new(&test.key, &test.nonce);
814                cipher.set_counter(test.initial_counter);
815                cipher.xor_keystream(&mut partial_plaintext[..n]);
816                cipher.xor_keystream(&mut partial_plaintext[n..]);
817
818                assert_eq!(
819                    plaintext,
820                    partial_plaintext,
821                    "ietf test [{i}] failed. partial encryption is not valid for n = {n}
822            Got: {}
823            Expected: {}",
824                    hex::encode(&partial_plaintext),
825                    hex::encode(&plaintext),
826                )
827            }
828        }
829    }
830
831    #[test]
832    fn chacha12_case_1() {
833        let nonce: &[u8; 8] = &[0xdb, 0x4b, 0x4a, 0x41, 0xd8, 0xdf, 0x18, 0xaa];
834        let key: &[u8; 32] = &[
835            0x27, 0xfc, 0x12, 0x0b, 0x01, 0x3b, 0x82, 0x9f, 0x1f, 0xae, 0xef, 0xd1, 0xab, 0x41, 0x7e, 0x86, 0x62, 0xf4,
836            0x3e, 0x0d, 0x73, 0xf9, 0x8d, 0xe8, 0x66, 0xe3, 0x46, 0x35, 0x31, 0x80, 0xfd, 0xb7,
837        ];
838
839        let mut buffer = [0u8; 100];
840        ChaCha12Djb::new(key, nonce).xor_keystream(&mut buffer);
841
842        assert_eq!(
843            buffer,
844            [
845                0x5f, 0x3c, 0x8c, 0x19, 0x0a, 0x78, 0xab, 0x7f, 0xe8, 0x08, 0xca, 0xe9, 0xcb, 0xcb, 0x0a, 0x98, 0x37,
846                0xc8, 0x93, 0x49, 0x2d, 0x96, 0x3a, 0x1c, 0x2e, 0xda, 0x6c, 0x15, 0x58, 0xb0, 0x2c, 0x83, 0xfc, 0x02,
847                0xa4, 0x4c, 0xbb, 0xb7, 0xe6, 0x20, 0x4d, 0x51, 0xd1, 0xc2, 0x43, 0x0e, 0x9c, 0x0b, 0x58, 0xf2, 0x93,
848                0x7b, 0xf5, 0x93, 0x84, 0x0c, 0x85, 0x0b, 0xda, 0x90, 0x51, 0xa1, 0xf0, 0x51, 0xdd, 0xf0, 0x9d, 0x2a,
849                0x03, 0xeb, 0xf0, 0x9f, 0x01, 0xbd, 0xba, 0x9d, 0xa0, 0xb6, 0xda, 0x79, 0x1b, 0x2e, 0x64, 0x56, 0x41,
850                0x04, 0x7d, 0x11, 0xeb, 0xf8, 0x50, 0x87, 0xd4, 0xde, 0x5c, 0x01, 0x5f, 0xdd, 0xd0, 0x44,
851            ]
852        );
853    }
854
855    #[test]
856    fn chacha8_case_1() {
857        let key = &[
858            0x64, 0x1a, 0xea, 0xeb, 0x08, 0x03, 0x6b, 0x61, 0x7a, 0x42, 0xcf, 0x14, 0xe8, 0xc5, 0xd2, 0xd1, 0x15, 0xf8,
859            0xd7, 0xcb, 0x6e, 0xa5, 0xe2, 0x8b, 0x9b, 0xfa, 0xf8, 0x3e, 0x03, 0x84, 0x26, 0xa7,
860        ];
861        let nonce = &[0xa1, 0x4a, 0x11, 0x68, 0x27, 0x1d, 0x45, 0x9b];
862
863        let mut buffer = [0u8; 100];
864        ChaCha8Djb::new(key, nonce).xor_keystream(&mut buffer);
865
866        assert_eq!(
867            buffer,
868            [
869                0x17, 0x21, 0xc0, 0x44, 0xa8, 0xa6, 0x45, 0x35, 0x22, 0xdd, 0xdb, 0x31, 0x43, 0xd0, 0xbe, 0x35, 0x12,
870                0x63, 0x3c, 0xa3, 0xc7, 0x9b, 0xf8, 0xcc, 0xc3, 0x59, 0x4c, 0xb2, 0xc2, 0xf3, 0x10, 0xf7, 0xbd, 0x54,
871                0x4f, 0x55, 0xce, 0x0d, 0xb3, 0x81, 0x23, 0x41, 0x2d, 0x6c, 0x45, 0x20, 0x7d, 0x5c, 0xf9, 0xaf, 0x0c,
872                0x6c, 0x68, 0x0c, 0xce, 0x1f, 0x7e, 0x43, 0x38, 0x8d, 0x1b, 0x03, 0x46, 0xb7, 0x13, 0x3c, 0x59, 0xfd,
873                0x6a, 0xf4, 0xa5, 0xa5, 0x68, 0xaa, 0x33, 0x4c, 0xcd, 0xc3, 0x8a, 0xf5, 0xac, 0xe2, 0x01, 0xdf, 0x84,
874                0xd0, 0xa3, 0xca, 0x22, 0x54, 0x94, 0xca, 0x62, 0x09, 0x34, 0x5f, 0xcf, 0x30, 0x13, 0x2e,
875            ]
876        );
877    }
878
879    #[test]
880    fn xchacha20_test_vector_counter_0() {
881        // draft-irtf-cfrg-xchacha-03, Appendix A.3.2.1
882        let key: [u8; 32] = hex::decode("808182838485868788898a8b8c8d8e8f909192939495969798999a9b9c9d9e9f")
883            .unwrap()
884            .try_into()
885            .unwrap();
886        let nonce: [u8; 24] = hex::decode("404142434445464748494a4b4c4d4e4f5051525354555658")
887            .unwrap()
888            .try_into()
889            .unwrap();
890        let plaintext = hex::decode(concat!(
891            "5468652064686f6c65202870726f6e6f756e6365642022646f6c652229206973",
892            "20616c736f206b6e6f776e2061732074686520417369617469632077696c6420",
893            "646f672c2072656420646f672c20616e642077686973746c696e6720646f672e",
894            "2049742069732061626f7574207468652073697a65206f662061204765726d61",
895            "6e20736865706865726420627574206c6f6f6b73206d6f7265206c696b652061",
896            "206c6f6e672d6c656767656420666f782e205468697320686967686c7920656c",
897            "757369766520616e6420736b696c6c6564206a756d70657220697320636c6173",
898            "736966696564207769746820776f6c7665732c20636f796f7465732c206a6163",
899            "6b616c732c20616e6420666f78657320696e20746865207461786f6e6f6d6963",
900            "2066616d696c792043616e696461652e",
901        ))
902        .unwrap();
903        let expected_ciphertext = hex::decode(concat!(
904            "4559abba4e48c16102e8bb2c05e6947f50a786de162f9b0b7e592a9b53d0d4e9",
905            "8d8d6410d540a1a6375b26d80dace4fab52384c731acbf16a5923c0c48d3575d",
906            "4d0d2c673b666faa731061277701093a6bf7a158a8864292a41c48e3a9b4c0da",
907            "ece0f8d98d0d7e05b37a307bbb66333164ec9e1b24ea0d6c3ffddcec4f68e744",
908            "3056193a03c810e11344ca06d8ed8a2bfb1e8d48cfa6bc0eb4e2464b74814240",
909            "7c9f431aee769960e15ba8b96890466ef2457599852385c661f752ce20f9da0c",
910            "09ab6b19df74e76a95967446f8d0fd415e7bee2a12a114c20eb5292ae7a349ae",
911            "577820d5520a1f3fb62a17ce6a7e68fa7c79111d8860920bc048ef43fe84486c",
912            "cb87c25f0ae045f0cce1e7989a9aa220a28bdd4827e751a24a6d5c62d790a663",
913            "93b93111c1a55dd7421a10184974c7c5",
914        ))
915        .unwrap();
916
917        let mut cipher = XChaCha20::new(&key, &nonce);
918        let mut buf = plaintext.clone();
919        cipher.xor_keystream(&mut buf);
920        assert_eq!(buf, expected_ciphertext);
921
922        // Decrypt
923        cipher.set_counter(0);
924        cipher.xor_keystream(&mut buf);
925        assert_eq!(buf, plaintext);
926    }
927
928    #[test]
929    fn xchacha20_set_counter() {
930        let key = [0x55u8; 32];
931        let nonce = [0xaau8; 24];
932        let plaintext = b"test message for xchacha20";
933
934        let mut cipher1 = XChaCha20::new(&key, &nonce);
935        let mut buf1 = plaintext.to_vec();
936        cipher1.xor_keystream(&mut buf1);
937
938        // Same cipher but with set_counter(0) should produce same output
939        let mut cipher2 = XChaCha20::new(&key, &nonce);
940        cipher2.set_counter(0);
941        let mut buf2 = plaintext.to_vec();
942        cipher2.xor_keystream(&mut buf2);
943
944        assert_eq!(buf1, buf2);
945
946        // Different counter should produce different output
947        let mut cipher3 = XChaCha20::new(&key, &nonce);
948        cipher3.set_counter(1);
949        let mut buf3 = plaintext.to_vec();
950        cipher3.xor_keystream(&mut buf3);
951        assert_ne!(buf1, buf3);
952    }
953
954    // -------------------------------------------------------------------------
955    // Edge-case leftover tests
956    // -------------------------------------------------------------------------
957
958    /// Helper: encrypt in one shot to get the expected reference.
959    fn encrypt_one_shot_djb(key: &[u8; 32], nonce: &[u8; 8], plaintext: &[u8]) -> Vec<u8> {
960        let mut buf = plaintext.to_vec();
961        ChaCha20Djb::new(key, nonce).xor_keystream(&mut buf);
962        buf
963    }
964
965    fn encrypt_one_shot_ietf(key: &[u8; 32], nonce: &[u8; 12], plaintext: &[u8]) -> Vec<u8> {
966        let mut buf = plaintext.to_vec();
967        ChaCha20Ietf::new(key, nonce).xor_keystream(&mut buf);
968        buf
969    }
970
971    fn encrypt_one_shot_xchacha(key: &[u8; 32], nonce: &[u8; 24], plaintext: &[u8]) -> Vec<u8> {
972        let mut buf = plaintext.to_vec();
973        XChaCha20::new(key, nonce).xor_keystream(&mut buf);
974        buf
975    }
976
977    fn test_key_32() -> [u8; 32] {
978        let mut k = [0u8; 32];
979        for i in 0..32 {
980            k[i] = i as u8;
981        }
982        k
983    }
984
985    fn test_nonce_8() -> [u8; 8] {
986        [0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07]
987    }
988
989    fn test_nonce_12() -> [u8; 12] {
990        [0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0a, 0x0b]
991    }
992
993    fn test_nonce_24() -> [u8; 24] {
994        let mut n = [0u8; 24];
995        for i in 0..24 {
996            n[i] = i as u8;
997        }
998        n
999    }
1000
1001    fn test_plaintext(len: usize) -> Vec<u8> {
1002        (0..len).map(|i| (i % 251) as u8).collect()
1003    }
1004
1005    /// IETF variant: multi-call leftover consumption (analogous to the DJB test above).
1006    #[test]
1007    fn chacha20_ietf_keystream_leftover_multi_call() {
1008        let key = test_key_32();
1009        let nonce = test_nonce_12();
1010        let pt = test_plaintext(300);
1011
1012        let expected = encrypt_one_shot_ietf(&key, &nonce, &pt);
1013
1014        // partial block -> leaves leftover, partially consume, then finish
1015        {
1016            let mut buf = pt.clone();
1017            let mut c = ChaCha20Ietf::new(&key, &nonce);
1018            c.xor_keystream(&mut buf[..10]);
1019            c.xor_keystream(&mut buf[10..15]);
1020            c.xor_keystream(&mut buf[15..]);
1021            assert_eq!(buf, expected, "ietf partial leftover consumption");
1022        }
1023
1024        // partial block -> exactly exhaust leftover -> fresh blocks
1025        {
1026            let mut buf = pt.clone();
1027            let mut c = ChaCha20Ietf::new(&key, &nonce);
1028            c.xor_keystream(&mut buf[..3]);
1029            c.xor_keystream(&mut buf[3..64]);
1030            c.xor_keystream(&mut buf[64..]);
1031            assert_eq!(buf, expected, "ietf exact leftover exhaustion");
1032        }
1033
1034        // three rounds of partial leftover consumption
1035        {
1036            let mut buf = pt.clone();
1037            let mut c = ChaCha20Ietf::new(&key, &nonce);
1038            c.xor_keystream(&mut buf[..5]);
1039            c.xor_keystream(&mut buf[5..12]);
1040            c.xor_keystream(&mut buf[12..20]);
1041            c.xor_keystream(&mut buf[20..]);
1042            assert_eq!(buf, expected, "ietf multiple partial leftover consumptions");
1043        }
1044    }
1045
1046    /// XChaCha20: multi-call leftover consumption.
1047    #[test]
1048    fn xchacha20_keystream_leftover_multi_call() {
1049        let key = test_key_32();
1050        let nonce = test_nonce_24();
1051        let pt = test_plaintext(300);
1052
1053        let expected = encrypt_one_shot_xchacha(&key, &nonce, &pt);
1054
1055        // partial block -> leaves leftover, partially consume, then finish
1056        {
1057            let mut buf = pt.clone();
1058            let mut c = XChaCha20::new(&key, &nonce);
1059            c.xor_keystream(&mut buf[..8]);
1060            c.xor_keystream(&mut buf[8..20]);
1061            c.xor_keystream(&mut buf[20..]);
1062            assert_eq!(buf, expected, "xchacha partial leftover consumption");
1063        }
1064
1065        // three rounds of partial consumption
1066        {
1067            let mut buf = pt.clone();
1068            let mut c = XChaCha20::new(&key, &nonce);
1069            c.xor_keystream(&mut buf[..13]);
1070            c.xor_keystream(&mut buf[13..27]);
1071            c.xor_keystream(&mut buf[27..40]);
1072            c.xor_keystream(&mut buf[40..]);
1073            assert_eq!(buf, expected, "xchacha multiple partial leftover consumptions");
1074        }
1075    }
1076
1077    /// Very small chunks: 1-byte calls to stress the leftover offset state machine.
1078    #[test]
1079    fn chacha_keystream_leftover_tiny_chunks() {
1080        let key = test_key_32();
1081        let nonce = test_nonce_8();
1082        let pt = test_plaintext(200);
1083
1084        // 10 calls of 1 byte each, then the rest in one shot
1085        let mut cipher = ChaCha20Djb::new(&key, &nonce);
1086        let mut buf = pt.clone();
1087        for n in 0..10 {
1088            cipher.xor_keystream(&mut buf[n..n + 1]);
1089        }
1090        cipher.xor_keystream(&mut buf[10..]);
1091        let expected = encrypt_one_shot_djb(&key, &nonce, &pt);
1092        assert_eq!(buf, expected, "tiny-chunk DJB failed");
1093
1094        // IETF variant
1095        let nonce12 = test_nonce_12();
1096        let mut buf = pt.clone();
1097        let mut cipher = ChaCha20Ietf::new(&key, &nonce12);
1098        for n in 0..10 {
1099            cipher.xor_keystream(&mut buf[n..n + 1]);
1100        }
1101        cipher.xor_keystream(&mut buf[10..]);
1102        let expected = encrypt_one_shot_ietf(&key, &nonce12, &pt);
1103        assert_eq!(buf, expected, "tiny-chunk IETF failed");
1104
1105        // XChaCha variant
1106        let nonce24 = test_nonce_24();
1107        let mut buf = pt.clone();
1108        let mut cipher = XChaCha20::new(&key, &nonce24);
1109        for n in 0..10 {
1110            cipher.xor_keystream(&mut buf[n..n + 1]);
1111        }
1112        cipher.xor_keystream(&mut buf[10..]);
1113        let expected = encrypt_one_shot_xchacha(&key, &nonce24, &pt);
1114        assert_eq!(buf, expected, "tiny-chunk XChaCha failed");
1115    }
1116
1117    /// Boundary sizes: 63, 64, 65, 127, 128, 129 bytes.
1118    /// Tests the leftover offset formula at exact transition points.
1119    #[test]
1120    fn chacha_keystream_leftover_boundary_sizes() {
1121        let key = test_key_32();
1122        let nonce = test_nonce_8();
1123
1124        for &len in &[63usize, 64, 65, 127, 128, 129, 191, 192, 193, 255, 256, 257] {
1125            let pt = test_plaintext(len);
1126            let expected = encrypt_one_shot_djb(&key, &nonce, &pt);
1127
1128            // split into 3 chunks: a, b, c where a+b+c = len
1129            // Use varying split sizes to exercise different leftover states
1130            for a in [0usize, 1, 8, 31, 32, 33, 62, 63].iter().copied() {
1131                if a > len {
1132                    continue;
1133                }
1134                for b in [0usize, 1, 7, 32, 63, 64].iter().copied() {
1135                    if a + b > len {
1136                        continue;
1137                    }
1138                    let c = len - a - b;
1139                    let splits = [a, b, c];
1140
1141                    let mut cipher = ChaCha20Djb::new(&key, &nonce);
1142                    let mut buf = pt.clone();
1143                    let mut offset = 0;
1144                    for &size in &splits {
1145                        cipher.xor_keystream(&mut buf[offset..offset + size]);
1146                        offset += size;
1147                    }
1148                    assert_eq!(buf, expected, "Boundary DJB failed for len={len} splits={splits:?}",);
1149                }
1150            }
1151        }
1152    }
1153
1154    /// Boundary sizes for IETF variant.
1155    #[test]
1156    fn chacha_keystream_leftover_boundary_sizes_ietf() {
1157        let key = test_key_32();
1158        let nonce = test_nonce_12();
1159
1160        for &len in &[63usize, 64, 65, 127, 128, 129] {
1161            let pt = test_plaintext(len);
1162            let expected = encrypt_one_shot_ietf(&key, &nonce, &pt);
1163
1164            for &a in &[0usize, 1, 31, 32, 33, 63] {
1165                if a > len {
1166                    continue;
1167                }
1168                let mut cipher = ChaCha20Ietf::new(&key, &nonce);
1169                let mut buf = pt.clone();
1170                cipher.xor_keystream(&mut buf[..a]);
1171                cipher.xor_keystream(&mut buf[a..]);
1172                assert_eq!(buf, expected, "Boundary IETF failed for len={len} a={a}",);
1173            }
1174        }
1175    }
1176
1177    /// Zero-byte intermediate calls: should be no-ops that don't corrupt leftover state.
1178    #[test]
1179    fn chacha_keystream_leftover_zero_byte_intermediate() {
1180        let key = test_key_32();
1181        let nonce = test_nonce_8();
1182        let pt = test_plaintext(150);
1183        let expected = encrypt_one_shot_djb(&key, &nonce, &pt);
1184
1185        let mut buf = pt.clone();
1186        let mut cipher = ChaCha20Djb::new(&key, &nonce);
1187
1188        // first partial call
1189        cipher.xor_keystream(&mut buf[..10]);
1190        // zero-byte call in the middle
1191        cipher.xor_keystream(&mut []);
1192        // second partial call
1193        cipher.xor_keystream(&mut buf[10..20]);
1194        // another zero-byte call
1195        cipher.xor_keystream(&mut []);
1196        // final call
1197        cipher.xor_keystream(&mut buf[20..]);
1198
1199        assert_eq!(buf, expected, "zero-byte intermediate DJB failed");
1200    }
1201
1202    /// set_counter mid-stream, then partial encryption: verifies leftover is cleared
1203    /// and counter is correctly reset.
1204    #[test]
1205    fn chacha_keystream_leftover_set_counter_mid_stream() {
1206        let key = test_key_32();
1207        let nonce = test_nonce_8();
1208        let pt = test_plaintext(200);
1209
1210        // one-shot reference
1211        let full_encrypted = encrypt_one_shot_djb(&key, &nonce, &pt);
1212
1213        // encrypt all at once with a fresh cipher
1214        let mut buf = pt.clone();
1215        let mut cipher = ChaCha20Djb::new(&key, &nonce);
1216        cipher.xor_keystream(&mut buf);
1217        assert_eq!(buf, full_encrypted, "baseline DJB");
1218
1219        // encrypt 10 bytes, reset counter to 0, re-encrypt those 10 bytes
1220        // (this XORs again -> back to plaintext), then encrypt the rest.
1221        let mut buf = pt.clone();
1222        let mut cipher = ChaCha20Djb::new(&key, &nonce);
1223        cipher.xor_keystream(&mut buf[..10]);
1224        cipher.set_counter(0);
1225        cipher.xor_keystream(&mut buf[..10]); // re-XOR -> bytes 0-9 are plaintext again
1226        cipher.xor_keystream(&mut buf[10..]);
1227
1228        // first 10 bytes are back to plaintext
1229        assert_eq!(
1230            &buf[..10],
1231            &pt[..10],
1232            "set_counter mid-stream: first 10 bytes should be plaintext"
1233        );
1234        // remaining bytes match the one-shot encryption
1235        assert_eq!(&buf[10..], &full_encrypted[10..], "set_counter mid-stream DJB failed");
1236
1237        // IETF variant: same pattern
1238        let nonce12 = test_nonce_12();
1239        let full_encrypted = encrypt_one_shot_ietf(&key, &nonce12, &pt);
1240
1241        let mut buf = pt.clone();
1242        let mut cipher = ChaCha20Ietf::new(&key, &nonce12);
1243        cipher.xor_keystream(&mut buf[..10]);
1244        cipher.set_counter(0);
1245        cipher.xor_keystream(&mut buf[..10]);
1246        cipher.xor_keystream(&mut buf[10..]);
1247        assert_eq!(&buf[..10], &pt[..10], "set_counter mid-stream: IETF first 10 bytes");
1248        assert_eq!(&buf[10..], &full_encrypted[10..], "set_counter mid-stream IETF failed");
1249    }
1250
1251    /// Stress test: many sequential calls with random-sized splits at block boundaries.
1252    #[test]
1253    fn chacha_keystream_leftover_stress_random_splits() {
1254        let key = test_key_32();
1255        let nonce = test_nonce_8();
1256
1257        let sizes = [50usize, 127, 128, 129, 200, 256, 300, 400, 512];
1258        let split_patterns: &[&[usize]] = &[
1259            &[1, 2, 3, 4, 5],
1260            &[7, 13, 23, 31],
1261            &[32, 32, 32],
1262            &[63, 1],
1263            &[64, 64],
1264            &[65, 63],
1265            &[33, 33, 33, 33],
1266            &[10, 10, 10, 10, 10],
1267            &[50, 50, 50],
1268        ];
1269
1270        for &len in &sizes {
1271            let pt = test_plaintext(len);
1272
1273            // ChaCha20
1274            let expected20 = {
1275                let mut b = pt.clone();
1276                ChaCha20Djb::new(&key, &nonce).xor_keystream(&mut b);
1277                b
1278            };
1279            for &splits in split_patterns {
1280                let mut buf = pt.clone();
1281                let mut offset = 0;
1282                let mut c = ChaCha20Djb::new(&key, &nonce);
1283                for &size in splits {
1284                    let end = core::cmp::min(offset + size, buf.len());
1285                    c.xor_keystream(&mut buf[offset..end]);
1286                    offset = end;
1287                    if offset >= buf.len() {
1288                        break;
1289                    }
1290                }
1291                if offset < buf.len() {
1292                    c.xor_keystream(&mut buf[offset..]);
1293                }
1294                assert_eq!(buf, expected20, "Stress ChaCha20 failed for len={len} splits={splits:?}",);
1295            }
1296
1297            // ChaCha12
1298            let expected12 = {
1299                let mut b = pt.clone();
1300                ChaCha12Djb::new(&key, &nonce).xor_keystream(&mut b);
1301                b
1302            };
1303            for &splits in split_patterns {
1304                let mut buf = pt.clone();
1305                let mut offset = 0;
1306                let mut c = ChaCha12Djb::new(&key, &nonce);
1307                for &size in splits {
1308                    let end = core::cmp::min(offset + size, buf.len());
1309                    c.xor_keystream(&mut buf[offset..end]);
1310                    offset = end;
1311                    if offset >= buf.len() {
1312                        break;
1313                    }
1314                }
1315                if offset < buf.len() {
1316                    c.xor_keystream(&mut buf[offset..]);
1317                }
1318                assert_eq!(buf, expected12, "Stress ChaCha12 failed for len={len} splits={splits:?}",);
1319            }
1320
1321            // ChaCha8
1322            let expected8 = {
1323                let mut b = pt.clone();
1324                ChaCha8Djb::new(&key, &nonce).xor_keystream(&mut b);
1325                b
1326            };
1327            for &splits in split_patterns {
1328                let mut buf = pt.clone();
1329                let mut offset = 0;
1330                let mut c = ChaCha8Djb::new(&key, &nonce);
1331                for &size in splits {
1332                    let end = core::cmp::min(offset + size, buf.len());
1333                    c.xor_keystream(&mut buf[offset..end]);
1334                    offset = end;
1335                    if offset >= buf.len() {
1336                        break;
1337                    }
1338                }
1339                if offset < buf.len() {
1340                    c.xor_keystream(&mut buf[offset..]);
1341                }
1342                assert_eq!(buf, expected8, "Stress ChaCha8 failed for len={len} splits={splits:?}",);
1343            }
1344        }
1345    }
1346
1347    /// IETF-specific variant of the stress test.
1348    #[test]
1349    fn chacha_keystream_leftover_stress_random_splits_ietf() {
1350        let key = test_key_32();
1351        let nonce = test_nonce_12();
1352
1353        for &len in &[50usize, 127, 128, 129, 200, 256, 300] {
1354            let pt = test_plaintext(len);
1355            let expected = encrypt_one_shot_ietf(&key, &nonce, &pt);
1356
1357            let split_patterns: &[&[usize]] = &[
1358                &[1, 2, 3, 4, 5],
1359                &[7, 13, 23, 31],
1360                &[63, 1],
1361                &[64, 64],
1362                &[65, 63],
1363                &[33, 33, 33, 33],
1364                &[50, 50, 50],
1365            ];
1366
1367            for &splits in split_patterns {
1368                let mut buf = pt.clone();
1369                let mut offset = 0;
1370                let mut cipher = ChaCha20Ietf::new(&key, &nonce);
1371                for &size in splits {
1372                    let end = core::cmp::min(offset + size, buf.len());
1373                    cipher.xor_keystream(&mut buf[offset..end]);
1374                    offset = end;
1375                    if offset >= buf.len() {
1376                        break;
1377                    }
1378                }
1379                if offset < buf.len() {
1380                    cipher.xor_keystream(&mut buf[offset..]);
1381                }
1382                assert_eq!(buf, expected, "Stress IETF failed for len={len} splits={splits:?}",);
1383            }
1384        }
1385    }
1386
1387    /// XChaCha-specific variant of the stress test.
1388    #[test]
1389    fn keystream_leftover_stress_random_splits_xchacha() {
1390        let key = test_key_32();
1391        let nonce = test_nonce_24();
1392
1393        for &len in &[50usize, 127, 128, 129, 200] {
1394            let pt = test_plaintext(len);
1395            let expected = encrypt_one_shot_xchacha(&key, &nonce, &pt);
1396
1397            let split_patterns: &[&[usize]] = &[&[1, 2, 3, 4, 5], &[7, 13, 23], &[63, 1], &[64, 64], &[65, 63]];
1398
1399            for &splits in split_patterns {
1400                let mut buf = pt.clone();
1401                let mut offset = 0;
1402                let mut cipher = XChaCha20::new(&key, &nonce);
1403                for &size in splits {
1404                    let end = core::cmp::min(offset + size, buf.len());
1405                    cipher.xor_keystream(&mut buf[offset..end]);
1406                    offset = end;
1407                    if offset >= buf.len() {
1408                        break;
1409                    }
1410                }
1411                if offset < buf.len() {
1412                    cipher.xor_keystream(&mut buf[offset..]);
1413                }
1414                assert_eq!(buf, expected, "Stress XChaCha failed for len={len} splits={splits:?}",);
1415            }
1416        }
1417    }
1418}
crypto/chacha/chacha.rs

crypto/chacha/
chacha.rs
