Merge code from dshulyak/master (#3)

Co-authored-by: Dmitry <yashulyak@gmail.com>

Cargo.toml

@@ -3,4 +3,22 @@ name = "post-rs"
 version = "0.1.0"
 edition = "2021"
 
+[lib]
+name = "post"
+
 [dependencies]
+rand = "0.8.5"
+aes = "0.8.2"
+cipher = "0.4.2"
+
+[dev-dependencies]
+criterion = "0.4"
+
+[features]
+default = ["block-padding"]
+block-padding = ["cipher/block-padding"]
+
+[[bench]]
+name = "aes"
+harness = false
+opt-level = 3

benches/aes.rs

@@ -0,0 +1,32 @@
+use std::hint::black_box;
+
+use criterion::{criterion_group, criterion_main, Criterion};
+use rand::{thread_rng, RngCore};
+
+fn prover_bench<const N: usize>(c: &mut Criterion) {
+    let mut group = c.benchmark_group("aes");
+
+    let challenge = b"dsadassdada12311";
+    let mut data: Vec<u8> = vec![0; 512 << 20];
+    thread_rng().fill_bytes(&mut data);
+    let mut prover: post::Prover<N> = post::Prover::new(challenge, 0);
+
+    let name = (N * 2).to_string();
+    group.throughput(criterion::Throughput::Bytes(data.len() as u64));
+    group.bench_function(name, |b| {
+        b.iter(|| {
+            let f = black_box(|_, _| Ok(()));
+            prover.prove(&data, f);
+        });
+    });
+}
+
+criterion_group!(
+    benches,
+    prover_bench::<1>,
+    prover_bench::<2>,
+    prover_bench::<3>,
+    prover_bench::<5>,
+    prover_bench::<10>,
+);
+criterion_main!(benches);

src/lib.rs

@@ -1,14 +1,4 @@
-pub fn add(left: usize, right: usize) -> usize {
-    left + right
-}
-
-#[cfg(test)]
-mod tests {
-    use super::*;
-
-    #[test]
-    fn it_works() {
-        let result = add(2, 2);
-        assert_eq!(result, 4);
-    }
-}
+mod prove;
+mod verify;
+pub use crate::prove::*;
+pub use crate::verify::*;

src/main.rs

@@ -0,0 +1,38 @@
+use std::{fs, io::Read, process::exit, sync::mpsc, time::SystemTime};
+
+fn main() {
+    let mut f = fs::File::open("/dev/random").expect("can't read from /dev/random");
+    let gb = 1;
+    let mut stream = vec![0u8; gb << 30];
+    let challenge = b"3213123122dsadsa";
+    let mut prover = post::Prover::<1>::new(challenge, u64::MAX >> 28);
+    let (tx, rx) = mpsc::channel();
+
+    for _ in 0..10 {
+        let start = SystemTime::now();
+        match f.read_exact(&mut stream) {
+            Err(err) => {
+                println!("reader failed {err}");
+                exit(1);
+            }
+            Ok(()) => {
+                println!("got data in {:?}", start.elapsed().unwrap());
+                let start = SystemTime::now();
+                prover.prove(&stream, |nonce, index| match tx.send((nonce, index)) {
+                    Ok(()) => Ok(()),
+                    Err(_) => Err(()),
+                });
+                println!(
+                    "proving finished in {:?} for file of size {:?}. estimated throughput {:?} GB/s",
+                    start.elapsed().unwrap(),
+                    stream.len(),
+                    gb as f64 * (1000 as f64 / start.elapsed().unwrap().as_millis() as f64),
+                );
+            }
+        };
+    }
+    drop(tx);
+    while let Ok((nonce, index)) = rx.recv() {
+        println!("nonce={nonce} index={index}");
+    }
+}

src/prove.rs

@@ -0,0 +1,108 @@
+use aes::{
+    cipher::{BlockEncrypt, KeyInit},
+    Aes128,
+};
+use cipher::block_padding::NoPadding;
+use std::sync::mpsc;
+
+const BLOCK_SIZE: usize = 16; // size of the aes block
+const BATCH: usize = 8; // will use encrypt8 asm method
+const CHUNK_SIZE: usize = BLOCK_SIZE * BATCH;
+
+pub struct Prover<const N: usize = 1> {
+    ciphers: [Aes128; N],
+    output: [u8; CHUNK_SIZE],
+    d: u64,
+}
+
+impl<const N: usize> Prover<N> {
+    pub fn new(challenge: &[u8; 16], d: u64) -> Self {
+        let ciphers: [Aes128; N] = (0..N as u32)
+            .map(|i: u32| {
+                let mut key = [0u8; BLOCK_SIZE];
+                key[..12].copy_from_slice(&challenge[..12]);
+                key[12..].copy_from_slice(&i.to_le_bytes());
+                Aes128::new(&key.into())
+            })
+            .collect::<Vec<_>>()
+            .try_into()
+            .unwrap();
+        let output = [0u8; CHUNK_SIZE];
+        Prover { ciphers, output, d }
+    }
+
+    pub fn prove<F>(&mut self, stream: &[u8], mut consume: F)
+    where
+        F: FnMut(u64, u64) -> Result<(), ()>,
+    {
+        for i in 0..stream.len() / (CHUNK_SIZE) {
+            let chunk = &stream[i * CHUNK_SIZE..(i + 1) * CHUNK_SIZE];
+            for (j, cipher) in self.ciphers.iter().enumerate() {
+                cipher
+                    .encrypt_padded_b2b::<NoPadding>(chunk, &mut self.output)
+                    .unwrap();
+                let ints = unsafe {
+                    let (_, ints, _) = self.output.align_to::<u64>();
+                    ints
+                };
+                for (out_i, out) in ints.into_iter().enumerate() {
+                    if out.to_le() <= self.d {
+                        let j = j * 2;
+                        let i = i * 8 + out_i;
+                        match consume((j + i % 2) as u64, (i / 2) as u64) {
+                            Ok(()) => {}
+                            Err(()) => return,
+                        }
+                    }
+                }
+            }
+        }
+    }
+}
+
+pub fn prove(stream: &[u8], challenge: &[u8; 16], d: u64, tx: &mpsc::Sender<(u64, u64)>) {
+    let mut prover = Prover::<1>::new(challenge, d);
+    prover.prove(stream, |nonce, index| -> Result<(), ()> {
+        match tx.send((nonce, index)) {
+            Ok(()) => Ok(()),
+            Err(_) => Err(()),
+        }
+    });
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn sanity() {
+        let (tx, rx) = mpsc::channel();
+        let challenge = b"dsadsadasdsaaaaa";
+        let stream = vec![0u8; 16 * 8];
+        prove(&stream, &challenge, u64::MAX, &tx);
+        drop(tx);
+        let rst: Vec<(u64, u64)> = rx.into_iter().collect();
+        assert_eq!(rst.len(), 16);
+        assert_eq!(
+            rst,
+            vec![
+                (0, 0),
+                (1, 0),
+                (0, 1),
+                (1, 1),
+                (0, 2),
+                (1, 2),
+                (0, 3),
+                (1, 3),
+                (0, 4),
+                (1, 4),
+                (0, 5),
+                (1, 5),
+                (0, 6),
+                (1, 6),
+                (0, 7),
+                (1, 7)
+            ],
+        );
+    }
+}

src/verify.rs

@@ -0,0 +1,35 @@
+use aes::{
+    cipher::{BlockEncrypt, KeyInit},
+    Aes128,
+};
+
+pub struct Proof {
+    nonce: u64,
+    indices: [u64; 725],
+}
+
+fn work_oracle(_i: u64) -> u8 {
+    0
+}
+
+pub fn verify(challenge: &[u8; 16], d: u64, proof: &Proof) -> bool {
+    let mut key = [0u8; 16];
+    key[..12].copy_from_slice(&challenge[..12]);
+    key[12..].copy_from_slice(&proof.nonce.to_le_bytes());
+    let cipher = Aes128::new(&key.into());
+
+    let mut output = [0u8; 16];
+    let output_index = proof.nonce as usize % 2;
+
+    for i in proof.indices.into_iter() {
+        let labels = (i..i + 16).map(|i| work_oracle(i)).collect();
+        cipher.encrypt_block_b2b(&labels, (&mut output).into());
+        unsafe {
+            let (_, ints, _) = output.align_to::<u64>();
+            if ints[output_index].to_le() > d {
+                return false;
+            }
+        }
+    }
+    true
+}