feat: v0.2.0-alpha-1, performance improved from 36MB/s to ~450MB/s

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "aes-gcm-stream"
-version = "0.1.1"
+version = "0.1.2"
 edition = "2021"
 authors = ["Hatter Jiang"]
 repository = "https://git.hatter.ink/hatter/aes-gcm-stream"
@@ -13,9 +13,12 @@ categories = ["cryptography"]
 
 [dependencies]
 aes = { version = "0.8.3", features = ["zeroize"] }
+ghash = "0.5.0"
 zeroize = { version = "1.6.0", features = ["zeroize_derive"] }
 
-[dev-dependencies]
 hex = "0.4.3"
+
+[dev-dependencies]
+#hex = "0.4.3"
 aes-gcm = { version = "0.10.2", features = ["zeroize"] }
 benchmark-simple = "0.1.8"

examples/bench.rs

@@ -2,7 +2,7 @@ use aes_gcm::{Aes256Gcm, KeyInit};
 use aes_gcm::aead::{Aead, Nonce};
 use benchmark_simple::{Bench, Options};
 
-use aes_gcm_stream::{Aes128GcmStreamEncryptor, Aes192GcmStreamEncryptor, Aes256GcmStreamEncryptor};
+use aes_gcm_stream::{Aes128GcmStreamEncryptor, Aes192GcmStreamEncryptor, Aes256GcmStreamEncryptor, Aes256GcmStreamEncryptor2};
 
 fn test_aes128_encrypt(m: &mut [u8]) {
     let key = [0u8; 16];
@@ -25,7 +25,7 @@ fn test_aes192_encrypt(m: &mut [u8]) {
 fn test_aes256_encrypt(m: &mut [u8]) {
     let key = [0u8; 32];
     let nonce = [0u8; 12];
-    let mut encryptor = Aes256GcmStreamEncryptor::new(key, &nonce);
+    let mut encryptor = Aes256GcmStreamEncryptor2::new(key, &nonce);
 
     encryptor.update(m);
     encryptor.finalize();
@@ -55,11 +55,11 @@ fn main() {
         ..Default::default()
     };
 
-    let res = bench.run(options, || test_aes128_encrypt(&mut m));
+    // let res = bench.run(options, || test_aes128_encrypt(&mut m));
-    println!("AES128 encrypt         : {}", res.throughput(m.len() as _));
+    // println!("AES128 encrypt         : {}", res.throughput(m.len() as _));
-
+    //
-    let res = bench.run(options, || test_aes192_encrypt(&mut m));
+    // let res = bench.run(options, || test_aes192_encrypt(&mut m));
-    println!("AES192 encrypt         : {}", res.throughput(m.len() as _));
+    // println!("AES192 encrypt         : {}", res.throughput(m.len() as _));
 
     let res = bench.run(options, || test_aes256_encrypt(&mut m));
     println!("AES256 encrypt         : {}", res.throughput(m.len() as _));

src/decryptor.rs

@@ -42,6 +42,7 @@ impl $module {
             message_len: 0,
         };
         let (ghash_key, normalized_nonce) = s.normalize_nonce(nonce);
+        println!("<<< KEY: {}", hex::encode(ghash_key.to_be_bytes()));
         s.ghash_key = ghash_key;
         s.init_nonce = normalized_nonce;
         s.encryption_nonce = normalized_nonce;
@@ -126,7 +127,9 @@ impl $module {
         let mut bs = self.init_nonce.to_be_bytes().clone();
         let block = Block::<$aesn>::from_mut_slice(&mut bs);
         self.crypto.encrypt_block(block);
+        println!("<<< final enc block: {}", hex::encode(&block.as_slice()));
         let tag_trunk = self.ghash_val.to_be_bytes();
+        println!("<<< final block: {}", hex::encode(&tag_trunk));
         let y = u8to128(&tag_trunk) ^ u8to128(&block.as_slice());
         y.to_be_bytes().to_vec()
     }
@@ -138,6 +141,7 @@ impl $module {
             let blocks_count = integrality_buffer_slice_len / 16;
             for i in 0..blocks_count {
                 let buf = &integrality_buffer_slice[i * 16..(i + 1) * 16];
+                println!("<<< block: {}", hex::encode(buf));
                 self.ghash_val = gmul_128(self.ghash_val ^ u8to128(buf), self.ghash_key)
             }
             self.integrality_buffer = integrality_buffer_slice[blocks_count * 16..].to_vec();

src/encryptor.rs

@@ -66,11 +66,21 @@ impl $module {
         }
         let blocks_count = message_buffer_len / 16;
         let mut encrypted_message = Vec::with_capacity(blocks_count * 16);
-        for i in 0..blocks_count {
+        let mut blocks = Vec::with_capacity(blocks_count);
+        for _ in 0..blocks_count {
             self.encryption_nonce = inc_32(self.encryption_nonce);
-            let mut ctr = self.encryption_nonce.to_be_bytes();
+            let ctr = self.encryption_nonce.to_be_bytes();
-            let block = Block::<$aesn>::from_mut_slice(&mut ctr);
+            let block = Block::<$aesn>::clone_from_slice(&ctr);
-            self.crypto.encrypt_block(block);
+            blocks.push(block);
+        }
+        // println!("block site: {}", blocks.len());
+            self.crypto.encrypt_blocks(&mut blocks);
+        for i in 0..blocks_count {
+            // self.encryption_nonce = inc_32(self.encryption_nonce);
+            // let mut ctr = self.encryption_nonce.to_be_bytes();
+            // let block = Block::<$aesn>::from_mut_slice(&mut ctr);
+            // self.crypto.encrypt_block(block);
+            let block = &blocks[i];
             let chunk = &message_buffer_slice[i * 16..(i + 1) * 16];
             let y = u8to128(chunk) ^ u8to128(&block.as_slice());
             encrypted_message.extend_from_slice(&y.to_be_bytes());
@@ -125,6 +135,7 @@ impl $module {
     }
 
     fn update_integrality_buffer(&mut self) {
+        // self.integrality_buffer.clear();
         let integrality_buffer_slice = self.integrality_buffer.as_slice();
         let integrality_buffer_slice_len = integrality_buffer_slice.len();
         if integrality_buffer_slice_len >= 16 {

src/encryptor2.rs

@@ -0,0 +1,145 @@
+use aes::Aes256;
+use aes::cipher::{Block, BlockEncrypt, BlockSizeUser, KeyInit};
+use aes::cipher::consts::U16;
+use aes::cipher::generic_array::GenericArray;
+use ghash::GHash;
+use ghash::universal_hash::UniversalHash;
+use zeroize::Zeroize;
+
+use crate::util::{inc_32, msb_s, normalize_nonce, u8to128};
+
+struct AesBlock {}
+
+impl BlockSizeUser for AesBlock {
+    type BlockSize = U16;
+}
+
+const BLOCK_SIZE: usize = 16;
+
+// #[derive(ZeroizeOnDrop)]
+pub struct Aes256GcmStreamEncryptor2 {
+    cipher: Aes256,
+    message_buffer: Vec<u8>,
+    ghash: GHash,
+    ghash_key: u128,
+    init_nonce: u128,
+    encryption_nonce: u128,
+    adata_len: usize,
+    message_len: usize,
+}
+
+impl Aes256GcmStreamEncryptor2 {
+    pub fn new(key: [u8; 32], nonce: &[u8]) -> Self {
+        let key = GenericArray::from(key);
+        let aes = Aes256::new(&key);
+
+        let mut ghash_key = ghash::Key::default();
+        aes.encrypt_block(&mut ghash_key);
+        let ghash = GHash::new(&ghash_key);
+        ghash_key.zeroize();
+
+        let mut s = Self {
+            cipher: aes,
+            message_buffer: vec![],
+            ghash,
+            ghash_key: 0,
+            init_nonce: 0,
+            encryption_nonce: 0,
+            adata_len: 0,
+            message_len: 0,
+        };
+        let (ghash_key, normalized_nonce) = s.normalize_nonce(nonce);
+        s.ghash_key = ghash_key;
+        s.init_nonce = normalized_nonce;
+        s.encryption_nonce = normalized_nonce;
+        s
+    }
+
+    pub fn init_adata(&mut self, adata: &[u8]) {
+        if adata.len() > 0 {
+            self.adata_len += adata.len();
+            self.ghash.update_padded(adata);
+        }
+    }
+
+    pub fn update(&mut self, message: &[u8]) -> Vec<u8> {
+        self.message_buffer.extend_from_slice(message);
+        let message_buffer_slice = self.message_buffer.as_slice();
+        if message_buffer_slice.len() < BLOCK_SIZE {
+            return Vec::with_capacity(0);
+        }
+        let blocks_count = message_buffer_slice.len() / BLOCK_SIZE;
+        let mut blocks = Vec::with_capacity(blocks_count);
+        for _ in 0..blocks_count {
+            self.encryption_nonce = inc_32(self.encryption_nonce);
+            let ctr = self.encryption_nonce.to_be_bytes();
+            blocks.push(Block::<AesBlock>::clone_from_slice(&ctr));
+        }
+        self.cipher.encrypt_blocks(&mut blocks);
+
+        let mut encrypted_message = message_buffer_slice[0..blocks_count * BLOCK_SIZE].to_vec();
+        for i in 0..blocks_count {
+            let chunk = &mut encrypted_message[i * BLOCK_SIZE..(i + 1) * BLOCK_SIZE];
+            let block = blocks[i].as_slice();
+            for k in 0..BLOCK_SIZE {
+                chunk[k] ^= block[k];
+            }
+        }
+        self.ghash.update_padded(&encrypted_message);
+        self.message_buffer = message_buffer_slice[blocks_count * BLOCK_SIZE..].to_vec();
+        self.message_len += encrypted_message.len();
+
+        encrypted_message
+    }
+
+    pub fn finalize(&mut self) -> (Vec<u8>, Vec<u8>) {
+        let mut final_encrypted_message = Vec::with_capacity(BLOCK_SIZE);
+        if !self.message_buffer.is_empty() {
+            // last block and this block len may less than 128 bits (16 bytes)
+            self.encryption_nonce = inc_32(self.encryption_nonce);
+            let mut ctr = self.encryption_nonce.to_be_bytes();
+            let block = Block::<AesBlock>::from_mut_slice(&mut ctr);
+            self.cipher.encrypt_block(block);
+
+            let chunk = self.message_buffer.as_slice();
+            let msb = msb_s(chunk.len() * 8, block.as_slice());
+            let y = u8to128(chunk) ^ u8to128(&msb);
+            final_encrypted_message.extend_from_slice(&y.to_be_bytes()[16 - chunk.len()..16]);
+            self.ghash.update_padded(&final_encrypted_message);
+            self.message_len += final_encrypted_message.len();
+        }
+        let adata_bit_len = self.adata_len * 8;
+        let message_bit_len = self.message_len * 8;
+        let mut adata_and_message_len = Vec::with_capacity(BLOCK_SIZE);
+        adata_and_message_len.extend_from_slice(&(adata_bit_len as u64).to_be_bytes());
+        adata_and_message_len.extend_from_slice(&(message_bit_len as u64).to_be_bytes());
+        self.ghash.update_padded(&adata_and_message_len);
+
+
+        let tag = self.compute_tag();
+
+        (final_encrypted_message, tag)
+    }
+
+    fn compute_tag(&mut self) -> Vec<u8> {
+        let mut bs = self.init_nonce.to_be_bytes().clone();
+        let block = Block::<AesBlock>::from_mut_slice(&mut bs);
+        self.cipher.encrypt_block(block);
+        let ghash = self.ghash.clone().finalize();
+        let tag_trunk = ghash.as_slice();
+        let y = u8to128(&tag_trunk) ^ u8to128(&block.as_slice());
+        y.to_be_bytes().to_vec()
+    }
+
+    fn ghash_key(&mut self) -> u128 {
+        let mut block = [0u8; BLOCK_SIZE];
+        let block = Block::<AesBlock>::from_mut_slice(&mut block);
+        self.cipher.encrypt_block(block);
+        u8to128(&block.as_slice())
+    }
+
+    fn normalize_nonce(&mut self, nonce_bytes: &[u8]) -> (u128, u128) {
+        let ghash_key = self.ghash_key();
+        normalize_nonce(ghash_key, nonce_bytes)
+    }
+}

src/lib.rs

@@ -17,11 +17,14 @@ pub use decryptor::Aes256GcmStreamDecryptor;
 pub use encryptor::Aes128GcmStreamEncryptor;
 pub use encryptor::Aes192GcmStreamEncryptor;
 pub use encryptor::Aes256GcmStreamEncryptor;
+pub use encryptor2::Aes256GcmStreamEncryptor2;
 
 mod util;
 mod encryptor;
 mod decryptor;
 
+mod encryptor2;
+
 #[test]
 fn test128() {
     use aes_gcm::{aead::{Aead, Nonce, Payload}, Aes128Gcm, KeyInit};
@@ -198,6 +201,7 @@ fn test256() {
     ];
 
     for (key, nonce, aad, plaintext) in knp {
+        println!("======= {} {} {} {}", hex::encode(key), hex::encode(nonce), hex::encode(aad), hex::encode(plaintext));
         // encrypt
         let mut ciphertext = vec![];
         let mut encryptor = Aes256GcmStreamEncryptor::new(key.clone(), &nonce);