use block struct, make stable

This commit is contained in:
Neemek 2025-10-03 23:38:34 +02:00
parent 1942ba7d8f
commit 87a4a08f63
2 changed files with 126 additions and 107 deletions

View file

@ -1,21 +1,3 @@
#![feature(portable_simd)]
#![feature(slice_as_array)]
#![feature(test)]
extern crate test;
use std::simd::u32x16;
#[inline]
fn many_u8_to_few_u32(data: [u8; 64]) -> [u32; 16] {
let mut out = [0u32; 16];
for (i, c) in data.chunks_exact(4).enumerate() {
out[i] = u32::from_ne_bytes(*c.as_array().unwrap());
}
out
}
#[inline]
fn few_u32_to_many_u8(data: &[u32; 16]) -> [u8; 64] {
@ -28,52 +10,89 @@ fn few_u32_to_many_u8(data: &[u32; 16]) -> [u8; 64] {
out
}
#[inline]
fn make_chacha_block(key: &[u32; 8], counter: u64, nonce: u64) -> u32x16 {
let [k1, k2, k3, k4, k5, k6, k7, k8] = key;
#[derive(Debug)]
pub struct Block{
key: [u32; 8],
counter: u32,
nonce: [u32; 3],
rounds: usize,
u32x16::from_array([
u32::from_ne_bytes(*b"expa"),
u32::from_ne_bytes(*b"nd 3"),
u32::from_ne_bytes(*b"2-by"),
u32::from_ne_bytes(*b"te k"),
*k1,
*k2,
*k3,
*k4,
*k5,
*k6,
*k7,
*k8,
(counter >> 32) as u32,
(counter & 0xFFFFFFFF) as u32,
(nonce >> 32) as u32,
(nonce & 0xFFFFFFFF) as u32,
])
chunk: [u8; 64],
byte_index: usize,
}
pub fn do_chacha(input: &[u8], key_bytes: &[u8; 32], nonce: u64, initial_counter: u64, rounds: usize) -> Vec<u8> {
const CHACHA_CONSTANT: [u32; 4] = [u32::from_ne_bytes(*b"expa"), u32::from_ne_bytes(*b"nd 3"), u32::from_ne_bytes(*b"2-by"), u32::from_ne_bytes(*b"te k")];
impl Block {
pub fn new(key: [u32; 8], counter: u32, nonce: [u32; 3], rounds: usize) -> Block {
Block{
key,
counter,
nonce,
rounds,
chunk: few_u32_to_many_u8(&chacha_block(&key, counter, nonce, rounds)),
byte_index: 0,
}
}
pub fn next(&mut self) {
self.byte_index = 0;
self.counter += 1;
self.chunk = few_u32_to_many_u8(&chacha_block(&self.key, self.counter, self.nonce, self.rounds));
}
pub fn process(&mut self, data: &mut [u8]) {
for i in 0..data.len() {
data[i] = data[i] ^ self.chunk[self.byte_index];
self.byte_index += 1;
if self.byte_index >= 64 {
self.next()
}
}
}
}
#[inline]
fn make_chacha_block(key: &[u32; 8], counter: u32, nonce: [u32; 3]) -> [u32; 16] {
[
CHACHA_CONSTANT[0],
CHACHA_CONSTANT[1],
CHACHA_CONSTANT[2],
CHACHA_CONSTANT[3],
key[0],
key[1],
key[2],
key[3],
key[4],
key[5],
key[6],
key[7],
counter,
nonce[0],
nonce[1],
nonce[2],
]
}
pub fn do_chacha(input: &[u8], key_bytes: &[u8; 32], nonce: [u32; 3], initial_counter: u32, rounds: usize) -> Vec<u8> {
let mut key: [u32; 8] = [0u32; 8];
for i in 0..8 {
key[i] = u32::from_le_bytes(key_bytes[i * 4..(i + 1) * 4].try_into().unwrap());
}
let mut bytes: Vec<u8> = Vec::new();
for (i, c) in input.chunks(64).enumerate() {
let mut x = [0u8; 64];
x[..c.len()].copy_from_slice(c);
bytes.extend(input);
let cipher = chacha_block(&key, i as u64 + initial_counter, nonce, rounds);
let data = u32x16::from_array(many_u8_to_few_u32(x));
let new = cipher ^ data;
bytes.extend_from_slice(&few_u32_to_many_u8(new.as_array())[..c.len()]);
}
Block::new(key, initial_counter, nonce, rounds).process(&mut bytes);
bytes
}
fn chacha_rounds(input: u32x16, rounds: usize) -> u32x16 {
let mut x: u32x16 = input;
fn chacha_rounds(input: [u32; 16], rounds: usize) -> [u32; 16] {
let mut x: [u32; 16] = input;
for i in 0..rounds {
match (i & 1) == 0 {
@ -94,18 +113,22 @@ fn chacha_rounds(input: u32x16, rounds: usize) -> u32x16 {
}
}
x + input
for i in 0..x.len() {
x[i] = x[i].wrapping_add(input[i]);
}
x
}
#[inline]
fn chacha_block(key: &[u32; 8], counter: u64, nonce: u64, rounds: usize) -> u32x16 {
fn chacha_block(key: &[u32; 8], counter: u32, nonce: [u32; 3], rounds: usize) -> [u32; 16] {
let block = make_chacha_block(key, counter, nonce);
chacha_rounds(block, rounds)
}
#[inline]
fn chacha_quarter_round(x: &mut u32x16, a: usize, b: usize, c: usize, d: usize) {
fn chacha_quarter_round(x: &mut [u32; 16], a: usize, b: usize, c: usize, d: usize) {
x[a] = x[a].wrapping_add(x[b]);
x[d] ^= x[a];
x[d] = x[d].rotate_left(16);
@ -123,55 +146,51 @@ fn chacha_quarter_round(x: &mut u32x16, a: usize, b: usize, c: usize, d: usize)
x[b] = x[b].rotate_left(7);
}
#[cfg(test)]
mod tests {
use test::{Bencher, black_box};
use crate::{chacha_quarter_round, chacha_rounds, do_chacha, few_u32_to_many_u8, make_chacha_block};
use std::simd::u32x16;
/*
use test::{Bencher, black_box};
use crate::{chacha_quarter_round, chacha_rounds, do_chacha, few_u32_to_many_u8, make_chacha_block};
#[test]
fn test_chacha_quarter_round() {
let mut data = [0u32; 16];
data[0] = 0x11111111;
data[1] = 0x01020304;
data[2] = 0x9b8d6f43;
data[3] = 0x01234567;
#[test]
fn test_chacha_quarter_round() {
let mut data = [0u32; 16];
data[0] = 0x11111111;
data[1] = 0x01020304;
data[2] = 0x9b8d6f43;
data[3] = 0x01234567;
let mut arr = u32x16::from_array(data);
chacha_quarter_round(&mut data, 0, 1, 2, 3);
chacha_quarter_round(&mut arr, 0, 1, 2, 3);
assert_eq!(arr[0], 0xea2a92f4);
assert_eq!(arr[1], 0xcb1cf8ce);
assert_eq!(arr[2], 0x4581472e);
assert_eq!(arr[3], 0x5881c4bb);
}
#[bench]
fn bench_make_chacha_block(b: &mut Bencher) {
b.iter(|| {
black_box(make_chacha_block(&[1u32; 8], 0, 0))
})
}
#[bench]
fn bench_chacha_rounds(b: &mut Bencher) {
let data = u32x16::from_array([69u32; 16]);
b.iter(|| {
black_box(chacha_rounds(data, 20));
})
}
#[bench]
fn bench_chacha_operation(b: &mut Bencher) {
// somewhat random data
let data = few_u32_to_many_u8(chacha_rounds(u32x16::from_array([69u32; 16]), 20).as_array());
let key = &[0u8; 32];
b.iter(|| {
black_box(do_chacha(data.as_slice(), key, 0, 0, 20));
})
}
assert_eq!(data[0], 0xea2a92f4);
assert_eq!(data[1], 0xcb1cf8ce);
assert_eq!(data[2], 0x4581472e);
assert_eq!(data[3], 0x5881c4bb);
}
#[bench]
fn bench_make_chacha_block(b: &mut Bencher) {
b.iter(|| {
black_box(make_chacha_block(&[1u32; 8], 0, [0; 3]))
})
}
#[bench]
fn bench_chacha_rounds(b: &mut Bencher) {
let data = [69u32; 16];
b.iter(|| {
black_box(chacha_rounds(data, 20));
})
}
#[bench]
fn bench_chacha_operation(b: &mut Bencher) {
// somewhat random data
let data = few_u32_to_many_u8(&chacha_rounds([69u32; 16], 20));
let key = &[0u8; 32];
b.iter(|| {
black_box(do_chacha(data.as_slice(), key, [0; 3], 0, 20));
})
}
*/

View file

@ -17,7 +17,7 @@ struct Args {
/// The nonce in the initial key value
#[arg(short, long, default_value = "0", value_parser=maybe_hex::<u64>)]
nonce: u64,
nonce: u128,
}
fn main() {
@ -25,16 +25,16 @@ fn main() {
let stdin = stdin();
let mut f = File::open(args.key_file).expect("failed to open key file");
let key_bytes = &mut [0u8; 32];
f.read_exact(key_bytes).expect("couldn't read key file");
let mut input = Vec::new();
let mut inp = stdin.lock();
inp.read_to_end(&mut input)
.expect("failed to read from stdin");
let mut f = File::open(args.key_file).expect("failed to open key file");
let key_bytes = &mut [0u8; 32];
f.read_exact(key_bytes).expect("couldn't read key file");
let bytes = do_chacha(input.as_slice(), key_bytes, args.nonce, 0, args.rounds);
let bytes = do_chacha(input.as_slice(), key_bytes, [(args.nonce >> 64& 0xFFFFFFFF) as u32, (args.nonce >> 32 & 0xFFFFFFFF) as u32, args.nonce as u32], 0, args.rounds);
let stdout = stdout();
let mut out = stdout.lock();