feat: Aes algorithm hardware acceleration

nervosnetwork · Oct 11, 2019 · 01a436e · 01a436e
1 parent f6cf320
commit 01a436e
Show file tree

Hide file tree

Showing 3 changed files with 35 additions and 16 deletions.
diff --git a/secio/Cargo.toml b/secio/Cargo.toml
@@ -32,13 +32,14 @@ rand = "0.6"
 ring = "0.16.5"
 twofish = "0.2.0"
 aes-ctr = "0.3.0"
+aesni = { version = "0.6", default-features = false, features = ["nocheck"] }
 ctr = "0.3.0"
 unsigned-varint = "0.2.2"
 bs58 = "0.2.0"
 
 [dev-dependencies]
 env_logger = "0.6"
-criterion = "0.2"
+criterion = "0.3"
 
 [features]
 default = []

diff --git a/secio/src/stream_cipher.rs b/secio/src/stream_cipher.rs
@@ -5,6 +5,7 @@ use crate::codec::StreamCipher;
 use aes_ctr::stream_cipher::generic_array::GenericArray;
 use aes_ctr::stream_cipher::NewStreamCipher;
 use aes_ctr::{Aes128Ctr, Aes256Ctr};
+use aesni::{Aes128Ctr as NIAes128Ctr, Aes256Ctr as NIAes256Ctr};
 use ctr::Ctr128;
 use twofish::Twofish;
 
@@ -39,18 +40,35 @@ impl Cipher {
 /// Returns your stream cipher depending on `Cipher`.
 #[inline]
 pub fn ctr_init(key_size: Cipher, key: &[u8], iv: &[u8]) -> StreamCipher {
-    match key_size {
-        Cipher::Aes128 => Box::new(Aes128Ctr::new(
-            GenericArray::from_slice(key),
-            GenericArray::from_slice(iv),
-        )),
-        Cipher::Aes256 => Box::new(Aes256Ctr::new(
-            GenericArray::from_slice(key),
-            GenericArray::from_slice(iv),
-        )),
-        Cipher::TwofishCtr => Box::new(Ctr128::<Twofish>::new(
-            GenericArray::from_slice(key),
-            GenericArray::from_slice(iv),
-        )),
+    if is_x86_feature_detected!("aes") && is_x86_feature_detected!("sse3") {
+        match key_size {
+            Cipher::Aes128 => Box::new(NIAes128Ctr::new(
+                GenericArray::from_slice(key),
+                GenericArray::from_slice(iv),
+            )),
+            Cipher::Aes256 => Box::new(NIAes256Ctr::new(
+                GenericArray::from_slice(key),
+                GenericArray::from_slice(iv),
+            )),
+            Cipher::TwofishCtr => Box::new(Ctr128::<Twofish>::new(
+                GenericArray::from_slice(key),
+                GenericArray::from_slice(iv),
+            )),
+        }
+    } else {
+        match key_size {
+            Cipher::Aes128 => Box::new(Aes128Ctr::new(
+                GenericArray::from_slice(key),
+                GenericArray::from_slice(iv),
+            )),
+            Cipher::Aes256 => Box::new(Aes256Ctr::new(
+                GenericArray::from_slice(key),
+                GenericArray::from_slice(iv),
+            )),
+            Cipher::TwofishCtr => Box::new(Ctr128::<Twofish>::new(
+                GenericArray::from_slice(key),
+                GenericArray::from_slice(iv),
+            )),
+        }
     }
 }
diff --git a/tests/test_kill.rs b/tests/test_kill.rs
@@ -17,10 +17,10 @@ use tentacle::{
 };
 
 /// Get current used memory(bytes)
-fn current_used_memory() -> Option<f32> {
+fn current_used_memory() -> Option<f64> {
     let sys = System::new();
     match sys.memory() {
-        Ok(mem) => Some((mem.total.as_usize() - mem.free.as_usize()) as f32),
+        Ok(mem) => Some((mem.total.as_u64() - mem.free.as_u64()) as f64),
         Err(_) => None,
     }
 }