use block struct, make stable

src/lib.rs

@@ -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]
+#[derive(Debug)]
-fn make_chacha_block(key: &[u32; 8], counter: u64, nonce: u64) -> u32x16 {
+pub struct Block{
-    let [k1, k2, k3, k4, k5, k6, k7, k8] = key;
+    key: [u32; 8],
+    counter: u32,
+    nonce: [u32; 3],
+    rounds: usize,
 
-    u32x16::from_array([
+    chunk: [u8; 64],
-        u32::from_ne_bytes(*b"expa"),
+    byte_index: usize,
-        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,
-    ])
 }
 
-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() {
+    bytes.extend(input);
-        let mut x = [0u8; 64];
-        x[..c.len()].copy_from_slice(c);
 
-        let cipher = chacha_block(&key, i as u64 + initial_counter, nonce, rounds);
+    Block::new(key, initial_counter, nonce, rounds).process(&mut bytes);
-        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()]);
-    }
 
     bytes
 }
 
-fn chacha_rounds(input: u32x16, rounds: usize) -> u32x16 {
+fn chacha_rounds(input: [u32; 16], rounds: usize) -> [u32; 16] {
-    let mut x: u32x16 = input;
+    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 test::{Bencher, black_box};
+use crate::{chacha_quarter_round, chacha_rounds, do_chacha, few_u32_to_many_u8, make_chacha_block};
-    use crate::{chacha_quarter_round, chacha_rounds, do_chacha, few_u32_to_many_u8, make_chacha_block};
-    use std::simd::u32x16;
 
-    #[test]
+#[test]
-    fn test_chacha_quarter_round() {
+fn test_chacha_quarter_round() {
-        let mut data = [0u32; 16];
+    let mut data = [0u32; 16];
-        data[0] = 0x11111111;
+    data[0] = 0x11111111;
-        data[1] = 0x01020304;
+    data[1] = 0x01020304;
-        data[2] = 0x9b8d6f43;
+    data[2] = 0x9b8d6f43;
-        data[3] = 0x01234567;
+    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!(data[0], 0xea2a92f4);
-
+    assert_eq!(data[1], 0xcb1cf8ce);
-        assert_eq!(arr[0], 0xea2a92f4);
+    assert_eq!(data[2], 0x4581472e);
-        assert_eq!(arr[1], 0xcb1cf8ce);
+    assert_eq!(data[3], 0x5881c4bb);
-        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));
-        })
-    }
 }
+
+#[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));
+    })
+}
+ */

src/main.rs

@@ -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 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 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 stdout = stdout();
     let mut out = stdout.lock();