From ed6c154a1fcc724c165c7709e8f0cfac93e465d9 Mon Sep 17 00:00:00 2001
From: Agnes Kiss <kisakishy@gmail.com>
Date: Mon, 24 Jun 2024 08:54:49 +0200
Subject: [PATCH] Modifications in preparation for including preOT

---
 src/otext/block.rs        | 16 +++----
 src/otext/crypto_utils.rs | 92 ++++++++++++++++++++++++++++++++++++++-
 2 files changed, 99 insertions(+), 9 deletions(-)

diff --git a/src/otext/block.rs b/src/otext/block.rs
index 8a5b152..2d30374 100644
--- a/src/otext/block.rs
+++ b/src/otext/block.rs
@@ -18,17 +18,21 @@ pub const ZERO_BLOCK: Block = 0;
 pub const ALL_ONE_BLOCK: Block = u128::MAX;
 
 /// XOR blocks
-pub fn xor_blocks_arr(res: &mut [Block], x: &[Block], y: &[Block], nblocks: usize) {
+pub fn xor_blocks_arr(x: &[Block], y: &[Block], nblocks: usize) -> Vec<Block> {
+    let mut res: Vec<Block> = vec![ZERO_BLOCK; nblocks];
     for i in 0..nblocks {
         res[i] = x[i] ^ y[i];
     }
+    res
 }
 
 /// XOR single block
-pub fn xor_blocks_arr_single(res: &mut [Block], x: &[Block], y: Block, nblocks: usize) {
+pub fn xor_blocks_arr_single(x: &[Block], y: Block, nblocks: usize) -> Vec<Block> {
+    let mut res: Vec<Block> = vec![ZERO_BLOCK; nblocks];
     for i in 0..nblocks {
         res[i] = x[i] ^ y;
     }
+    res
 }
 
 /// Compare blocks
@@ -57,9 +61,7 @@ mod tests {
     fn test_xor_blocks_arr() {
         let x = vec![make_block(1, 2); 3];
         let y = vec![make_block(3, 4); 3];
-        let mut res = vec![ZERO_BLOCK; 3];
-
-        xor_blocks_arr(&mut res, &x, &y, 3);
+        let res = xor_blocks_arr(&x, &y, 3);
 
         for block in &res {
             assert_eq!(*block, make_block(1 ^ 3, 2 ^ 4));
@@ -70,9 +72,7 @@ mod tests {
     fn test_xor_blocks_arr_single() {
         let x = vec![make_block(1, 2); 3];
         let y = make_block(5, 6);
-        let mut res = vec![ZERO_BLOCK; 3];
-
-        xor_blocks_arr_single(&mut res, &x, y, 3);
+        let res = xor_blocks_arr_single(&x, y, 3);
 
         for block in &res {
             assert_eq!(*block, make_block(1 ^ 5, 2 ^ 6));
diff --git a/src/otext/crypto_utils.rs b/src/otext/crypto_utils.rs
index 85083dc..c90b990 100644
--- a/src/otext/crypto_utils.rs
+++ b/src/otext/crypto_utils.rs
@@ -1,8 +1,12 @@
 //! Smaller crypto utilities
 use blake3::Hasher;
 use super::block::ZERO_BLOCK;
+use aes::Aes128;
+use aes::cipher::{BlockEncrypt, KeyInit};
 
-use super::block::Block;
+use super::block::{Block, xor_blocks_arr};
+
+const AES_BATCH_SIZE: usize = 8;
 
 /// Function to hash data using BLAKE3 and store the result in a Block[2] array
 pub fn hash_once(dgst: &mut [Block; 2], data: &Block) {
@@ -87,4 +91,90 @@ pub fn vector_inn_prdt_sum_red(a: &[Block], b: &[Block], sz: usize) -> Block {
         r ^= r1;
     }
     r
+}
+
+pub struct PRP {
+    aes_key: Aes128,
+}
+
+impl PRP {
+    // Constructor with no key
+    pub fn new() -> PRP {
+        PRP {
+            aes_key: Aes128::new(&ZERO_BLOCK.to_le_bytes().into()),
+        }
+    }
+
+    // Constructor with a key
+    pub fn with_key(key: Block) -> PRP {
+        PRP {
+            aes_key: Aes128::new(&key.to_le_bytes().into()),
+        }
+    }
+
+    // Permute blocks
+    pub fn permute_block(&self, data: &mut [Block], nblocks: usize) {
+        for chunk in data.chunks_exact_mut(AES_BATCH_SIZE) {
+            for block in chunk.iter_mut() {
+                let mut block_bytes = block.to_le_bytes();
+                self.aes_key.encrypt_block((&mut block_bytes).into());
+                *block = u128::from_le_bytes(block_bytes);
+            }
+        }
+
+        let remainder = nblocks % AES_BATCH_SIZE;
+        if remainder > 0 {
+            let start = nblocks - remainder;
+            for i in start..nblocks {
+                let mut block_bytes = data[i].to_le_bytes();
+                self.aes_key.encrypt_block((&mut block_bytes).into());
+                data[i] = u128::from_le_bytes(block_bytes);
+            }
+        }
+    }
+}
+
+pub struct CCRH {
+    prp: PRP,
+}
+
+impl CCRH {
+    pub fn new(key: Block) -> Self {
+        CCRH { prp: PRP::with_key(key) }
+    }
+
+    pub fn h(&self, input: Block) -> Block {
+        let t = sigma(input);
+        self.prp.permute_block(&mut [t], 1);
+        t ^ input
+    }
+
+    pub fn h_fixed<const N: usize>(&self, input: &[Block; N]) -> Vec<Block> {
+        let mut tmp = [0u128; N];
+        for i in 0..N {
+            tmp[i] = sigma(input[i]);
+        }
+        self.prp.permute_block(&mut tmp, N);
+        xor_blocks_arr(&tmp, input, N)
+    }
+
+    pub fn hn(&self, out: &mut Vec<Block>, input: &[Block], length: usize, scratch: Option<&mut [Block]>) -> Vec<Block> {
+        let mut local_scratch = vec![0u128; length];
+        let scratch = scratch.unwrap_or(local_scratch.as_mut());
+
+        for i in 0..length {
+            scratch[i] = sigma(input[i]);
+            out[i] = scratch[i];
+        }
+
+        self.prp.permute_block(scratch, length);
+        xor_blocks_arr(scratch, out, length)
+    }
+}
+
+fn sigma(block: Block) -> Block {
+    let xl = block >> 64;
+    let xr = block << 64;
+    let xlxl = xl ^ (xl << 64);
+    xlxl ^ xr
 }
\ No newline at end of file