Skip to main content

dcrypt_algorithms/stream/chacha/chacha20/
mod.rs

1//! ChaCha20 stream cipher implementation
2//!
3//! This module implements the ChaCha20 stream cipher as defined in RFC 8439.
4
5use crate::error::{Error, Result};
6use crate::types::nonce::ChaCha20Compatible;
7use crate::types::Nonce;
8use dcrypt_common::security::{EphemeralSecret, SecretBuffer};
9use dcrypt_internal::zeroing::{Zeroize, ZeroizeOnDrop, Zeroizing};
10
11/// Size of ChaCha20 key in bytes
12pub const CHACHA20_KEY_SIZE: usize = 32;
13/// Size of ChaCha20 nonce in bytes
14pub const CHACHA20_NONCE_SIZE: usize = 12;
15/// Size of ChaCha20 block in bytes
16pub const CHACHA20_BLOCK_SIZE: usize = 64;
17
18/// ChaCha20 stream cipher
19#[derive(Clone)]
20pub struct ChaCha20 {
21    /// The key schedule
22    state: [u32; 16],
23    /// Keystream buffer
24    buffer: [u8; CHACHA20_BLOCK_SIZE],
25    /// Current position in the buffer
26    position: usize,
27    /// Current block counter
28    counter: u32,
29    /// Set after counter `u32::MAX` has been consumed.
30    exhausted: bool,
31}
32
33impl Zeroize for ChaCha20 {
34    fn zeroize(&mut self) {
35        self.state.zeroize();
36        self.buffer.zeroize();
37        self.position.zeroize();
38        self.counter.zeroize();
39        self.exhausted.zeroize();
40    }
41}
42
43impl Drop for ChaCha20 {
44    fn drop(&mut self) {
45        self.zeroize();
46    }
47}
48
49impl ZeroizeOnDrop for ChaCha20 {}
50
51impl ChaCha20 {
52    /// Creates a new ChaCha20 instance with the specified key and nonce
53    pub fn new<const N: usize>(key: &[u8; CHACHA20_KEY_SIZE], nonce: &Nonce<N>) -> Self
54    where
55        Nonce<N>: ChaCha20Compatible,
56    {
57        // Wrap key in SecretBuffer for secure handling
58        let key_buf = SecretBuffer::new(*key);
59        Self::with_counter_secure(&key_buf, nonce, 0)
60    }
61
62    /// Creates a new ChaCha20 instance with the specified key, nonce, and counter
63    pub fn with_counter<const N: usize>(
64        key: &[u8; CHACHA20_KEY_SIZE],
65        nonce: &Nonce<N>,
66        counter: u32,
67    ) -> Self
68    where
69        Nonce<N>: ChaCha20Compatible,
70    {
71        // Wrap key in SecretBuffer for secure handling
72        let key_buf = SecretBuffer::new(*key);
73        Self::with_counter_secure(&key_buf, nonce, counter)
74    }
75
76    /// Internal method that works with SecretBuffer for secure key handling
77    fn with_counter_secure<const N: usize>(
78        key: &SecretBuffer<CHACHA20_KEY_SIZE>,
79        nonce: &Nonce<N>,
80        counter: u32,
81    ) -> Self
82    where
83        Nonce<N>: ChaCha20Compatible,
84    {
85        // Initialize state with constants and key
86        let mut state = Zeroizing::new([0u32; 16]);
87
88        // "expand 32-byte k" in little-endian
89        state[0] = 0x61707865;
90        state[1] = 0x3320646e;
91        state[2] = 0x79622d32;
92        state[3] = 0x6b206574;
93
94        // Key (8 words) - use secure key access
95        let key_bytes = key.as_ref();
96        for i in 0..8 {
97            state[4 + i] =
98                u32::from_le_bytes(key_bytes[i * 4..i * 4 + 4].try_into().expect("four bytes"));
99        }
100
101        // Counter (1 word)
102        state[12] = counter;
103
104        // Nonce (3 words)
105        let nonce_bytes = nonce.as_ref();
106        state[13] = u32::from_le_bytes(nonce_bytes[0..4].try_into().expect("four bytes"));
107        state[14] = u32::from_le_bytes(nonce_bytes[4..8].try_into().expect("four bytes"));
108        state[15] = u32::from_le_bytes(nonce_bytes[8..12].try_into().expect("four bytes"));
109
110        let instance = Self {
111            state: state.into_inner(),
112            buffer: [0; CHACHA20_BLOCK_SIZE],
113            position: CHACHA20_BLOCK_SIZE, // Force initial keystream generation
114            counter,
115            exhausted: false,
116        };
117        instance
118    }
119
120    /// The ChaCha20 quarter round function
121    #[inline]
122    fn quarter_round(state: &mut [u32], a: usize, b: usize, c: usize, d: usize) {
123        state[a] = state[a].wrapping_add(state[b]);
124        state[d] ^= state[a];
125        state[d] = state[d].rotate_left(16);
126
127        state[c] = state[c].wrapping_add(state[d]);
128        state[b] ^= state[c];
129        state[b] = state[b].rotate_left(12);
130
131        state[a] = state[a].wrapping_add(state[b]);
132        state[d] ^= state[a];
133        state[d] = state[d].rotate_left(8);
134
135        state[c] = state[c].wrapping_add(state[d]);
136        state[b] ^= state[c];
137        state[b] = state[b].rotate_left(7);
138    }
139
140    /// Generate a block of keystream
141    fn generate_keystream(&mut self) -> Result<()> {
142        if self.exhausted {
143            return Err(Error::Processing {
144                operation: "ChaCha20",
145                details: "block counter exhausted",
146            });
147        }
148
149        // Create a working copy of the state
150        let mut working_state = self.state;
151
152        // Ensure the current counter is set in the working state
153        working_state[12] = self.counter;
154
155        // 20 rounds of ChaCha20: 10 column rounds, 10 diagonal rounds
156        for _ in 0..10 {
157            // Column rounds
158            Self::quarter_round(&mut working_state, 0, 4, 8, 12);
159            Self::quarter_round(&mut working_state, 1, 5, 9, 13);
160            Self::quarter_round(&mut working_state, 2, 6, 10, 14);
161            Self::quarter_round(&mut working_state, 3, 7, 11, 15);
162
163            // Diagonal rounds
164            Self::quarter_round(&mut working_state, 0, 5, 10, 15);
165            Self::quarter_round(&mut working_state, 1, 6, 11, 12);
166            Self::quarter_round(&mut working_state, 2, 7, 8, 13);
167            Self::quarter_round(&mut working_state, 3, 4, 9, 14);
168        }
169
170        // Create output by adding the working state to the original state
171        // Use EphemeralSecret to ensure intermediate values are zeroized
172        let mut output_state = EphemeralSecret::new([0u32; 16]);
173        for i in 0..16 {
174            let original_val = if i == 12 { self.counter } else { self.state[i] };
175            output_state[i] = working_state[i].wrapping_add(original_val);
176        }
177
178        // Convert to bytes (little-endian)
179        for i in 0..16 {
180            self.buffer[i * 4..i * 4 + 4].copy_from_slice(&output_state[i].to_le_bytes());
181        }
182        working_state.zeroize();
183
184        // Reset position and increment counter for next block
185        self.position = 0;
186        if self.counter == u32::MAX {
187            self.exhausted = true;
188        } else {
189            self.counter += 1;
190        }
191        Ok(())
192    }
193
194    fn ensure_capacity(&self, data_len: usize) -> Result<()> {
195        let buffered = if self.position < CHACHA20_BLOCK_SIZE {
196            CHACHA20_BLOCK_SIZE - self.position
197        } else {
198            0
199        };
200        let bytes_requiring_blocks = data_len.saturating_sub(buffered);
201        let blocks_required = bytes_requiring_blocks.div_ceil(CHACHA20_BLOCK_SIZE) as u64;
202        let blocks_available = if self.exhausted {
203            0
204        } else {
205            u64::from(u32::MAX - self.counter) + 1
206        };
207
208        if blocks_required > blocks_available {
209            return Err(Error::Processing {
210                operation: "ChaCha20",
211                details: "message would wrap the block counter",
212            });
213        }
214        Ok(())
215    }
216
217    /// Encrypt or decrypt data in place using the ChaCha20 stream cipher
218    pub fn process(&mut self, data: &mut [u8]) -> Result<()> {
219        self.ensure_capacity(data.len())?;
220        for byte in data.iter_mut() {
221            // Generate new keystream block if needed
222            if self.position >= CHACHA20_BLOCK_SIZE {
223                self.generate_keystream()?;
224            }
225
226            // XOR data with keystream
227            *byte ^= self.buffer[self.position];
228            self.position += 1;
229        }
230        Ok(())
231    }
232
233    /// Encrypt data in place
234    pub fn encrypt(&mut self, data: &mut [u8]) -> Result<()> {
235        self.process(data)
236    }
237
238    /// Decrypt data in place
239    pub fn decrypt(&mut self, data: &mut [u8]) -> Result<()> {
240        self.process(data)
241    }
242
243    /// Generate keystream directly into an output buffer
244    pub fn keystream(&mut self, output: &mut [u8]) -> Result<()> {
245        // Zero the output buffer
246        for byte in output.iter_mut() {
247            *byte = 0;
248        }
249
250        // Force generation from a block boundary (ignore any leftover position)
251        self.position = CHACHA20_BLOCK_SIZE;
252
253        // Then run the encryption pass to copy the keystream
254        self.process(output)
255    }
256
257    /// Seek to a specific block position
258    ///
259    /// `block_offset` is the number of full blocks that have been consumed;
260    /// after seeking, the next generated block will be at `block_offset + 1`.
261    pub fn seek(&mut self, block_offset: u32) -> Result<()> {
262        // Set counter so that generate_keystream() yields the next block
263        self.counter = block_offset.checked_add(1).ok_or(Error::Processing {
264            operation: "ChaCha20 seek",
265            details: "block offset would wrap the counter",
266        })?;
267        self.exhausted = false;
268
269        // Force regeneration on next use
270        self.position = CHACHA20_BLOCK_SIZE;
271
272        // Clear any old keystream
273        self.buffer.zeroize();
274        Ok(())
275    }
276
277    /// Reset to initial state with the same key
278    pub fn reset(&mut self) {
279        self.counter = self.state[12]; // Restore original counter
280        self.exhausted = false;
281        self.position = CHACHA20_BLOCK_SIZE; // Force keystream regeneration
282        self.buffer.zeroize(); // Clear keystream buffer
283    }
284}
285
286/// HChaCha20 core used to derive the XChaCha20 subkey.
287///
288/// Unlike a ChaCha20 block, HChaCha20 does not add the initial state back to
289/// the post-round state. It returns words 0..=3 and 12..=15.
290pub(crate) fn hchacha20(
291    key: &[u8; CHACHA20_KEY_SIZE],
292    nonce: &[u8; 16],
293) -> Zeroizing<[u8; CHACHA20_KEY_SIZE]> {
294    let mut state = Zeroizing::new([0u32; 16]);
295    state[0] = 0x6170_7865;
296    state[1] = 0x3320_646e;
297    state[2] = 0x7962_2d32;
298    state[3] = 0x6b20_6574;
299    for i in 0..8 {
300        state[4 + i] = u32::from_le_bytes(key[i * 4..i * 4 + 4].try_into().expect("four bytes"));
301    }
302    for i in 0..4 {
303        state[12 + i] = u32::from_le_bytes(nonce[i * 4..i * 4 + 4].try_into().expect("four bytes"));
304    }
305
306    for _ in 0..10 {
307        ChaCha20::quarter_round(&mut state[..], 0, 4, 8, 12);
308        ChaCha20::quarter_round(&mut state[..], 1, 5, 9, 13);
309        ChaCha20::quarter_round(&mut state[..], 2, 6, 10, 14);
310        ChaCha20::quarter_round(&mut state[..], 3, 7, 11, 15);
311        ChaCha20::quarter_round(&mut state[..], 0, 5, 10, 15);
312        ChaCha20::quarter_round(&mut state[..], 1, 6, 11, 12);
313        ChaCha20::quarter_round(&mut state[..], 2, 7, 8, 13);
314        ChaCha20::quarter_round(&mut state[..], 3, 4, 9, 14);
315    }
316
317    let words = Zeroizing::new([
318        state[0], state[1], state[2], state[3], state[12], state[13], state[14], state[15],
319    ]);
320    let mut out = Zeroizing::new([0u8; CHACHA20_KEY_SIZE]);
321    for (chunk, word) in out.chunks_exact_mut(4).zip(words.iter().copied()) {
322        chunk[0] = word as u8;
323        chunk[1] = (word >> 8) as u8;
324        chunk[2] = (word >> 16) as u8;
325        chunk[3] = (word >> 24) as u8;
326    }
327    out
328}
329
330#[cfg(test)]
331mod tests;