Skip to main content

vector_core/crypto/chachapoly/
mod.rs

1//! ChaCha20-Poly1305 (RFC 8439), tuned for many short messages.
2//!
3//! Byte-for-byte the same construction as the `chacha20poly1305` crate; what
4//! differs is the cost shape. The Poly1305 key block is produced in the same SIMD
5//! pass as the first data blocks, and Poly1305 is scalar, so a chat message costs
6//! one keystream pass and no per-message vector setup.
7
8mod chacha;
9mod poly1305;
10
11use chacha::{Backend, State, BUF};
12use poly1305::Poly1305;
13use zeroize::Zeroize;
14
15pub const KEY_LEN: usize = 32;
16pub const NONCE_LEN: usize = 12;
17pub const TAG_LEN: usize = 16;
18
19/// RFC 8439's 32-bit block counter starts at 1 for data.
20const MAX_LEN: u64 = (u32::MAX as u64 - 1) * 64;
21
22#[derive(Debug, Clone, Copy, PartialEq, Eq)]
23pub struct Error;
24
25impl std::fmt::Display for Error {
26    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
27        f.write_str("aead error")
28    }
29}
30
31impl std::error::Error for Error {}
32
33/// A keyed ChaCha20-Poly1305 instance. The expanded key is zeroized on drop.
34pub struct ChaCha20Poly1305 {
35    key: [u32; 8],
36    backend: Backend,
37    /// Poly1305 input length from which the SIMD lanes take over.
38    poly_min: usize,
39}
40
41impl Drop for ChaCha20Poly1305 {
42    fn drop(&mut self) {
43        self.key.zeroize();
44    }
45}
46
47impl ChaCha20Poly1305 {
48    pub fn new(key: &[u8; KEY_LEN]) -> Self {
49        let mut k = [0u32; 8];
50        for (w, c) in k.iter_mut().zip(key.chunks_exact(4)) {
51            *w = u32::from_le_bytes(c.try_into().unwrap());
52        }
53        Self { key: k, backend: Backend::detect(), poly_min: poly1305::simd_min() }
54    }
55
56    #[cfg(test)]
57    fn with_backend(key: &[u8; KEY_LEN], backend: Backend) -> Self {
58        let mut s = Self::new(key);
59        s.backend = backend;
60        s
61    }
62
63    fn state(&self, nonce: &[u8; NONCE_LEN]) -> State {
64        let mut n = [0u32; 3];
65        for (w, c) in n.iter_mut().zip(nonce.chunks_exact(4)) {
66            *w = u32::from_le_bytes(c.try_into().unwrap());
67        }
68        State { key: self.key, nonce: n }
69    }
70
71    /// Encrypt `buf` in place and return the tag.
72    pub fn seal_in_place(&self, nonce: &[u8; NONCE_LEN], aad: &[u8], buf: &mut [u8]) -> Result<[u8; TAG_LEN], Error> {
73        if buf.len() as u64 > MAX_LEN {
74            return Err(Error);
75        }
76        let st = self.state(nonce);
77        let mut ks = [0u8; BUF];
78        let (mut mac, first) = self.first_pass(&st, buf.len(), &mut ks);
79        mac.padded(aad, self.poly_min);
80        let used = self.xor_stream(&st, buf, &mut ks, first);
81        mac.padded(buf, self.poly_min);
82        let tag = finish(mac, aad.len(), buf.len());
83        wipe(&mut ks[..used]);
84        Ok(tag)
85    }
86
87    /// Verify and decrypt `buf` in place. On a bad tag `buf` is left as ciphertext.
88    pub fn open_in_place(
89        &self,
90        nonce: &[u8; NONCE_LEN],
91        aad: &[u8],
92        buf: &mut [u8],
93        tag: &[u8; TAG_LEN],
94    ) -> Result<(), Error> {
95        if buf.len() as u64 > MAX_LEN {
96            return Err(Error);
97        }
98        let st = self.state(nonce);
99        let mut ks = [0u8; BUF];
100        let (mut mac, first) = self.first_pass(&st, buf.len(), &mut ks);
101        mac.padded(aad, self.poly_min);
102        mac.padded(buf, self.poly_min);
103        let expected = finish(mac, aad.len(), buf.len());
104        if !ct_eq(&expected, tag) {
105            wipe(&mut ks[..64 * first]);
106            return Err(Error);
107        }
108        let used = self.xor_stream(&st, buf, &mut ks, first);
109        wipe(&mut ks[..used]);
110        Ok(())
111    }
112
113    /// Opens `nonce(12) || ciphertext || tag(16)` in place and returns the plaintext
114    /// range within `data`. The storage format every at-rest blob uses.
115    pub fn open_framed_in_place<'a>(&self, data: &'a mut [u8]) -> Result<&'a mut [u8], Error> {
116        if data.len() < NONCE_LEN + TAG_LEN {
117            return Err(Error);
118        }
119        let (nonce, rest) = data.split_at_mut(NONCE_LEN);
120        let nonce: [u8; NONCE_LEN] = (&*nonce).try_into().unwrap();
121        let (body, tag) = rest.split_at_mut(rest.len() - TAG_LEN);
122        let tag: [u8; TAG_LEN] = (&*tag).try_into().unwrap();
123        self.open_in_place(&nonce, &[], body, &tag)?;
124        Ok(body)
125    }
126
127    /// Generates the first keystream batch (counter 0 onward) into `ks` and keys
128    /// Poly1305 from block 0. Returns the MAC and the batch's block count; blocks
129    /// 1.. of that batch are data keystream that `xor_stream` consumes first.
130    #[inline]
131    fn first_pass(&self, st: &State, len: usize, ks: &mut [u8; BUF]) -> (Poly1305, usize) {
132        let want = 1 + len.div_ceil(64);
133        let n = self.backend.keystream(st, 0, want, ks);
134        let pk: &[u8; 32] = ks[..32].try_into().unwrap();
135        (Poly1305::new(pk), n)
136    }
137
138    /// Returns how many leading bytes of `ks` hold keystream and need wiping.
139    #[inline]
140    fn xor_stream(&self, st: &State, buf: &mut [u8], ks: &mut [u8; BUF], first: usize) -> usize {
141        let mut used = 64 * first;
142        let avail = 64 * (first - 1);
143        let head = avail.min(buf.len());
144        xor(&mut buf[..head], &ks[64..64 + head]);
145        let mut counter = first as u32;
146        let mut rest = &mut buf[head..];
147        while !rest.is_empty() {
148            let want = rest.len().div_ceil(64);
149            let n = self.backend.keystream(st, counter, want, ks);
150            used = used.max(64 * n);
151            let take = (64 * n).min(rest.len());
152            let (now, later) = rest.split_at_mut(take);
153            xor(now, &ks[..take]);
154            rest = later;
155            counter = counter.wrapping_add(n as u32);
156        }
157        used
158    }
159}
160
161#[inline]
162fn finish(mut mac: Poly1305, aad_len: usize, ct_len: usize) -> [u8; TAG_LEN] {
163    let mut lens = [0u8; 16];
164    lens[..8].copy_from_slice(&(aad_len as u64).to_le_bytes());
165    lens[8..].copy_from_slice(&(ct_len as u64).to_le_bytes());
166    mac.blocks(&lens);
167    mac.finish()
168}
169
170#[inline(always)]
171fn xor(dst: &mut [u8], ks: &[u8]) {
172    for (d, k) in dst.iter_mut().zip(ks) {
173        *d ^= *k;
174    }
175}
176
177/// Zeroize with word-wide volatile stores. `zeroize` on a byte slice issues one
178/// volatile store per byte, which costs more than the cipher on short messages.
179#[inline]
180pub fn wipe(buf: &mut [u8]) {
181    // SAFETY: the slice is `len` writable bytes.
182    unsafe { wipe_raw(buf.as_mut_ptr(), buf.len()) }
183}
184
185/// [`wipe`] over a vector's whole allocation, spare capacity included, then empty it.
186#[inline]
187pub fn wipe_vec(v: &mut Vec<u8>) {
188    // SAFETY: the allocation is `capacity` writable bytes; zeros are a valid u8 anywhere.
189    unsafe {
190        wipe_raw(v.as_mut_ptr(), v.capacity());
191        v.set_len(0);
192    }
193}
194
195/// SAFETY: `p` must be valid for `len` byte writes.
196#[inline]
197unsafe fn wipe_raw(p: *mut u8, len: usize) {
198    use core::ptr::write_volatile;
199    let head = p.align_offset(8).min(len);
200    let words = (len - head) / 8;
201    for i in 0..head {
202        write_volatile(p.add(i), 0);
203    }
204    let w = p.add(head) as *mut u64;
205    for i in 0..words {
206        write_volatile(w.add(i), 0);
207    }
208    for i in head + 8 * words..len {
209        write_volatile(p.add(i), 0);
210    }
211    core::sync::atomic::compiler_fence(core::sync::atomic::Ordering::SeqCst);
212}
213
214#[inline]
215fn ct_eq(a: &[u8; 16], b: &[u8; 16]) -> bool {
216    let d = u128::from_ne_bytes(*a) ^ u128::from_ne_bytes(*b);
217    let folded = (d as u64) | ((d >> 64) as u64);
218    core::hint::black_box(folded) == 0
219}
220
221/// The AEAD and Poly1305 with their paths forced, for measuring crossovers on a device.
222#[doc(hidden)]
223#[cfg(target_pointer_width = "64")]
224pub mod bench {
225    use super::chacha::Backend;
226    use super::poly1305::r64::{self, Poly1305};
227    use super::ChaCha20Poly1305;
228
229    /// ChaCha20 backends this CPU can run, by name.
230    pub fn chacha_backends() -> Vec<&'static str> {
231        Backend::available().into_iter().map(|(n, _)| n).collect()
232    }
233
234    /// The ChaCha20 backend production picks on this CPU.
235    pub fn default_chacha_backend() -> &'static str {
236        let chosen = Backend::detect();
237        Backend::available().into_iter().find(|(_, b)| *b == chosen).map_or("portable", |(n, _)| n)
238    }
239
240    /// Whether this CPU has the Poly1305 lanes.
241    pub fn poly_simd_available() -> bool {
242        r64::simd_available()
243    }
244
245    /// The production Poly1305 lanes threshold on this CPU (`usize::MAX`: lanes off).
246    pub fn poly_simd_min() -> usize {
247        super::poly1305::simd_min()
248    }
249
250    /// An AEAD on the named ChaCha20 backend with the given Poly1305 lanes threshold, or
251    /// None if this CPU can't run that backend.
252    pub fn cipher(key: &[u8; 32], chacha: &str, poly_simd_min: usize) -> Option<ChaCha20Poly1305> {
253        let (_, backend) = Backend::available().into_iter().find(|(n, _)| *n == chacha)?;
254        let mut c = ChaCha20Poly1305::new(key);
255        c.backend = backend;
256        c.poly_min = poly_simd_min;
257        Some(c)
258    }
259
260    /// Poly1305 over `data` (a multiple of 64 bytes) with the scalar loop.
261    pub fn poly1305_scalar(key: &[u8; 32], data: &[u8]) -> [u8; 16] {
262        let mut p = Poly1305::new(key);
263        p.blocks(data);
264        p.finish()
265    }
266
267    /// Poly1305 over `data` (a nonzero multiple of 64 bytes) with the lanes, or None
268    /// where this CPU has none.
269    pub fn poly1305_simd(key: &[u8; 32], data: &[u8]) -> Option<[u8; 16]> {
270        #[cfg(target_arch = "x86_64")]
271        if r64::simd_available() && !data.is_empty() && data.len().is_multiple_of(64) {
272            let mut p = Poly1305::new(key);
273            // SAFETY: lanes available; the length is a nonzero multiple of 64.
274            unsafe { p.blocks_simd(data) };
275            return Some(p.finish());
276        }
277        let _ = (key, data);
278        None
279    }
280}
281
282#[cfg(test)]
283mod tests;
284