use super::{overlapping, quic::Sample, Nonce};
use crate::cpu;
use cfg_if::cfg_if;
cfg_if! {
if #[cfg(any(
all(target_arch = "aarch64", target_endian = "little"),
all(target_arch = "arm", target_endian = "little"),
target_arch = "x86",
target_arch = "x86_64"
))] {
#[macro_use]
mod ffi;
#[cfg(any(target_arch = "x86", test))]
mod fallback;
} else {
mod fallback;
}
}
use crate::polyfill::ArraySplitMap;
pub type Overlapping<'o> = overlapping::Overlapping<'o, u8>;
#[derive(Clone)]
pub struct Key {
words: [u32; KEY_LEN / 4],
}
impl Key {
pub(super) fn new(value: [u8; KEY_LEN]) -> Self {
Self {
words: value.array_split_map(u32::from_le_bytes),
}
}
}
impl Key {
#[inline]
pub(super) fn encrypt_single_block_with_ctr_0<const N: usize>(
&self,
nonce: Nonce,
in_out: &mut [u8; N],
cpu: cpu::Features,
) -> Counter {
assert!(N <= BLOCK_LEN);
let (zero, one) = Counter::zero_one_less_safe(nonce);
self.encrypt(zero, in_out.as_mut().into(), cpu);
one
}
#[inline]
pub fn new_mask(&self, sample: Sample) -> [u8; 5] {
let cpu = cpu::features(); let (ctr, nonce) = sample.split_at(4);
let ctr = u32::from_le_bytes(ctr.try_into().unwrap());
let nonce = Nonce::assume_unique_for_key(nonce.try_into().unwrap());
let ctr = Counter::from_nonce_and_ctr(nonce, ctr);
let mut out: [u8; 5] = [0; 5];
self.encrypt(ctr, out.as_mut().into(), cpu);
out
}
#[inline(always)]
pub(super) fn encrypt(&self, counter: Counter, in_out: Overlapping<'_>, cpu: cpu::Features) {
cfg_if! {
if #[cfg(all(target_arch = "aarch64", target_endian = "little"))] {
use cpu::{GetFeature as _, arm::Neon};
const NEON_MIN_LEN: usize = 192 + 1;
if in_out.len() >= NEON_MIN_LEN {
if let Some(cpu) = cpu.get_feature() {
return chacha20_ctr32_ffi!(
unsafe { (NEON_MIN_LEN, Neon, Overlapping<'_>) => ChaCha20_ctr32_neon },
self, counter, in_out, cpu);
}
}
if in_out.len() >= 1 {
chacha20_ctr32_ffi!(
unsafe { (1, (), Overlapping<'_>) => ChaCha20_ctr32_nohw },
self, counter, in_out, ())
}
} else if #[cfg(all(target_arch = "arm", target_endian = "little"))] {
use cpu::{GetFeature as _, arm::Neon};
const NEON_MIN_LEN: usize = 192 + 1;
if in_out.len() >= NEON_MIN_LEN {
if let Some(cpu) = cpu.get_feature() {
return chacha20_ctr32_ffi!(
unsafe { (NEON_MIN_LEN, Neon, &mut [u8]) => ChaCha20_ctr32_neon },
self, counter, in_out.copy_within(), cpu);
}
}
if in_out.len() >= 1 {
chacha20_ctr32_ffi!(
unsafe { (1, (), &mut [u8]) => ChaCha20_ctr32_nohw },
self, counter, in_out.copy_within(), ())
}
} else if #[cfg(target_arch = "x86")] {
use cpu::{GetFeature as _, intel::Ssse3};
if in_out.len() >= 1 {
if let Some(cpu) = cpu.get_feature() {
chacha20_ctr32_ffi!(
unsafe { (1, Ssse3, &mut [u8]) => ChaCha20_ctr32_ssse3 },
self, counter, in_out.copy_within(), cpu)
} else {
let _: cpu::Features = cpu;
fallback::ChaCha20_ctr32(self, counter, in_out)
}
}
} else if #[cfg(target_arch = "x86_64")] {
use cpu::{GetFeature, intel::{Avx2, Ssse3}};
const SSE_MIN_LEN: usize = 128 + 1; if in_out.len() >= SSE_MIN_LEN {
let values = cpu.values();
if let Some(cpu) = values.get_feature() {
return chacha20_ctr32_ffi!(
unsafe { (SSE_MIN_LEN, Avx2, Overlapping<'_>) => ChaCha20_ctr32_avx2 },
self, counter, in_out, cpu);
}
if let Some(cpu) = values.get_feature() {
return chacha20_ctr32_ffi!(
unsafe { (SSE_MIN_LEN, Ssse3, Overlapping<'_>) =>
ChaCha20_ctr32_ssse3_4x },
self, counter, in_out, cpu);
}
}
if in_out.len() >= 1 {
chacha20_ctr32_ffi!(
unsafe { (1, (), Overlapping<'_>) => ChaCha20_ctr32_nohw },
self, counter, in_out, ())
}
} else {
let _: cpu::Features = cpu;
fallback::ChaCha20_ctr32(self, counter, in_out)
}
}
}
#[inline]
pub(super) fn words_less_safe(&self) -> &[u32; KEY_LEN / 4] {
&self.words
}
}
#[repr(transparent)]
pub struct Counter([u32; 4]);
impl Counter {
fn zero_one_less_safe(nonce: Nonce) -> (Self, Self) {
let ctr0 @ Self([_, n0, n1, n2]) = Self::from_nonce_and_ctr(nonce, 0);
let ctr1 = Self([1, n0, n1, n2]);
(ctr0, ctr1)
}
fn from_nonce_and_ctr(nonce: Nonce, ctr: u32) -> Self {
let [n0, n1, n2] = nonce.as_ref().array_split_map(u32::from_le_bytes);
Self([ctr, n0, n1, n2])
}
#[cfg(any(
test,
not(any(
all(target_arch = "aarch64", target_endian = "little"),
all(target_arch = "arm", target_endian = "little"),
target_arch = "x86_64"
))
))]
fn into_words_less_safe(self) -> [u32; 4] {
self.0
}
}
pub const KEY_LEN: usize = 32;
const BLOCK_LEN: usize = 64;
#[cfg(test)]
mod tests {
extern crate alloc;
use super::{super::overlapping::IndexError, *};
use crate::{error, test};
use alloc::vec;
const MAX_ALIGNMENT_AND_OFFSET: (usize, usize) = (15, 259);
const MAX_ALIGNMENT_AND_OFFSET_SUBSET: (usize, usize) =
if cfg!(any(not(debug_assertions), feature = "slow_tests")) {
MAX_ALIGNMENT_AND_OFFSET
} else {
(0, 0)
};
#[test]
fn chacha20_test_default() {
let max_offset = if cfg!(any(
all(target_arch = "aarch64", target_endian = "little"),
all(target_arch = "arm", target_endian = "little"),
target_arch = "x86",
target_arch = "x86_64"
)) {
MAX_ALIGNMENT_AND_OFFSET
} else {
MAX_ALIGNMENT_AND_OFFSET_SUBSET
};
chacha20_test(max_offset, Key::encrypt);
}
#[test]
fn chacha20_test_fallback() {
chacha20_test(MAX_ALIGNMENT_AND_OFFSET_SUBSET, |key, ctr, in_out, _cpu| {
fallback::ChaCha20_ctr32(key, ctr, in_out)
});
}
fn chacha20_test(
max_alignment_and_offset: (usize, usize),
f: impl for<'k, 'o> Fn(&'k Key, Counter, Overlapping<'o>, cpu::Features),
) {
let cpu = cpu::features();
let mut buf = vec![0u8; 1300];
test::run(test_file!("chacha_tests.txt"), move |section, test_case| {
assert_eq!(section, "");
let key = test_case.consume_bytes("Key");
let key: &[u8; KEY_LEN] = key.as_slice().try_into()?;
let key = Key::new(*key);
let ctr = test_case.consume_usize("Ctr");
let nonce = test_case.consume_bytes("Nonce");
let input = test_case.consume_bytes("Input");
let output = test_case.consume_bytes("Output");
for len in 0..=input.len() {
#[allow(clippy::cast_possible_truncation)]
chacha20_test_case_inner(
&key,
&nonce,
ctr as u32,
&input[..len],
&output[..len],
&mut buf,
max_alignment_and_offset,
cpu,
&f,
);
}
Ok(())
});
}
fn chacha20_test_case_inner(
key: &Key,
nonce: &[u8],
ctr: u32,
input: &[u8],
expected: &[u8],
buf: &mut [u8],
(max_alignment, max_offset): (usize, usize),
cpu: cpu::Features,
f: &impl for<'k, 'o> Fn(&'k Key, Counter, Overlapping<'o>, cpu::Features),
) {
const ARBITRARY: u8 = 123;
for alignment in 0..=max_alignment {
buf[..alignment].fill(ARBITRARY);
let buf = &mut buf[alignment..];
for offset in 0..=max_offset {
let buf = &mut buf[..(offset + input.len())];
buf[..offset].fill(ARBITRARY);
let src = offset..;
buf[src.clone()].copy_from_slice(input);
let ctr = Counter::from_nonce_and_ctr(
Nonce::try_assume_unique_for_key(nonce).unwrap(),
ctr,
);
let in_out = Overlapping::new(buf, src)
.map_err(error::erase::<IndexError>)
.unwrap();
f(key, ctr, in_out, cpu);
assert_eq!(&buf[..input.len()], expected)
}
}
}
}