Initial commit

src/main.rs

@@ -0,0 +1,157 @@
+#![feature(portable_simd)]
+#![feature(slice_as_array)]
+
+use std::fs::File;
+use std::io::{stdin, stdout, Read, Write};
+use std::simd::u32x16;
+use clap::Parser;
+use clap_num::maybe_hex;
+
+#[derive(Parser, Debug)]
+#[command(version, about, long_about = None)]
+struct Args {
+    /// Where the key is stored
+    #[arg(short, long)]
+    key_file: String,
+
+    /// The amount of chacha rounds to do
+    #[arg(short, long, default_value = "20")]
+    rounds: usize,
+
+    /// The nonce in the initial key value
+    #[arg(short, long, default_value = "0", value_parser=maybe_hex::<u64>)]
+    nonce: u64,
+}
+
+fn main() {
+    let args = Args::parse();
+
+    let stdin = stdin();
+
+    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 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 cipher = make_chacha_block(&key, 0, args.nonce);
+
+    let mut output = Vec::new();
+    let mut data_chunks = input.chunks_exact(64);
+    for c in &mut data_chunks {
+        cipher = chacha_rounds(cipher, args.rounds);
+        let data = u32x16::from(many_u8_to_few_u32(c.try_into().unwrap()));
+        cipher ^= data;
+        output.push(cipher);
+    }
+
+    let mut bytes: Vec<u8> = Vec::new();
+    for d in output {
+        bytes.extend_from_slice(few_u32_to_many_u8(d.as_array()).as_slice());
+    }
+
+    let remainder = data_chunks.remainder();
+    let remaining = remainder.len();
+    if remaining != 0 {
+        let mut padded = [0u8; 64];
+        padded[0..remainder.len()].copy_from_slice(remainder);
+
+        cipher = chacha_rounds(cipher, 20);
+        let data = u32x16::from(many_u8_to_few_u32(padded));
+        cipher ^= data;
+        bytes.extend_from_slice(&few_u32_to_many_u8(cipher.as_array())[..remaining]);
+    }
+
+    let stdout = stdout();
+    let mut out = stdout.lock();
+    out.write_all(bytes.as_slice()).unwrap();
+}
+
+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.try_into().unwrap());
+    }
+
+    out
+}
+
+fn few_u32_to_many_u8(data: &[u32; 16]) -> [u8; 64] {
+    let mut out = [0u8; 64];
+
+    for (i, c) in data.iter().enumerate() {
+        out[i*4..][..4].copy_from_slice(&u32::to_ne_bytes(*c));
+    }
+
+    out
+}
+
+#[inline]
+fn make_chacha_block(key: &[u32; 8], counter: u64, nonce: u64) -> u32x16 {
+    let [k1, k2, k3, k4, k5, k6, k7, k8] = key;
+
+    u32x16::from([
+        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
+    ])
+}
+
+fn chacha_rounds(input: u32x16, rounds: usize) -> u32x16 {
+    let x: u32x16 = input;
+
+    for i in 0..rounds {
+        match i & 1 {
+            // odd rounds
+            0 => {
+                chacha_quarter_round(x, 0, 4,  8, 12);
+                chacha_quarter_round(x, 1, 5,  9, 13);
+                chacha_quarter_round(x, 2, 6, 10, 14);
+                chacha_quarter_round(x, 3, 7, 11, 15);
+            }
+            // even rounds
+            1 => {
+                chacha_quarter_round(x, 0, 5, 10, 15);
+                chacha_quarter_round(x, 1, 6, 11, 12);
+                chacha_quarter_round(x, 2, 7,  8, 13);
+                chacha_quarter_round(x, 3, 4,  9, 14);
+            }
+
+            _ => ()
+        }
+
+    }
+
+    x + input
+}
+
+#[inline]
+fn chacha_quarter_round(mut x: u32x16, a: usize, b: usize, c: usize, d: usize) {
+    x[a] += x[b];
+    x[d] ^= x[a];
+    x[d] = x[d].rotate_left(16);
+
+    x[c] += x[d];
+    x[b] ^= x[c];
+    x[b] = x[b].rotate_left(12);
+
+    x[a] += x[b];
+    x[d] ^= x[a];
+    x[d] = x[d].rotate_left(8);
+
+    x[c] += x[d];
+    x[b] ^= x[c];
+    x[b] = x[b].rotate_left(7);
+}