Implement much faster sha256 and sha512.

sha2/Cargo.toml

@@ -10,5 +10,15 @@ keywords = ["crypto", "sha2", "asm"]
 categories = ["cryptography", "no-std"]
 edition = "2018"
 
+[dependencies]
+cfg-if = "1.0.0"
+
 [build-dependencies]
 cc = "1.0"
+
+[dev-dependencies]
+criterion = {version= "0.3.5",features=["html_reports"] }
+
+[[bench]]
+name = "sha2"
+harness = false

sha2/benches/sha2.rs

@@ -0,0 +1,17 @@
+use criterion::{black_box, criterion_group, criterion_main, Criterion};
+
+pub fn criterion_benchmark(c: &mut Criterion) {
+    let mut state = Default::default();
+    let data = [[0u8; 64]];
+    c.bench_function("sha256", |b| {
+        b.iter(|| sha2_asm::compress256(black_box(&mut state), black_box(&data)))
+    });
+    let mut state = Default::default();
+    let data = [[0u8; 128]];
+    c.bench_function("sha512", |b| {
+        b.iter(|| sha2_asm::compress512(black_box(&mut state), black_box(&data)))
+    });
+}
+
+criterion_group!(benches, criterion_benchmark);
+criterion_main!(benches);

sha2/build.rs

@@ -1,14 +1,34 @@
+use std::collections::HashSet;
+
 fn main() {
     use std::env;
 
     let target_arch = env::var("CARGO_CFG_TARGET_ARCH").unwrap_or_default();
     let target_vendor = env::var("CARGO_CFG_TARGET_VENDOR").unwrap_or_default();
+    let features: HashSet<String> = env::var("CARGO_CFG_TARGET_FEATURE")
+        .unwrap_or_default()
+        .split(',')
+        .map(|x| x.to_string())
+        .collect();
 
     let mut build256 = cc::Build::new();
     let (sha256_path, sha512_path) = if target_arch == "x86" {
         ("src/sha256_x86.S", "src/sha512_x86.S")
     } else if target_arch == "x86_64" {
-        ("src/sha256_x64.S", "src/sha512_x64.S")
+        let sha512 = if features.contains("avx2") {
+            "src/sha512_x64_avx2.S"
+        } else {
+            "src/sha512_x64.S"
+        };
+        // Prioritizing aes-ni, cause it's fastest
+        let sha256 = if features.contains("aes") {
+            "src/sha256_x64_ni.S"
+        } else if features.contains("avx2") {
+            "src/sha256_x64_avx2.S"
+        } else {
+            "src/sha256_x64.S"
+        };
+        (sha256, sha512)
     } else if target_arch == "aarch64" && target_vendor == "apple" {
         build256.flag("-march=armv8-a+crypto");
         ("src/sha256_aarch64_apple.S", "")

sha2/src/lib.rs

@@ -14,29 +14,68 @@
 compile_error!("crate can only be used on x86, x86-64 and aarch64 architectures");
 
 #[link(name = "sha256", kind = "static")]
+#[allow(dead_code)]
 extern "C" {
     fn sha256_compress(state: &mut [u32; 8], block: &[u8; 64]);
+    #[cfg(target_feature = "avx2")]
+    fn sha256_transform_rorx(state: &mut [u32; 8], block: *const [u8; 64], num_blocks: usize);
+    #[cfg(target_feature = "aes")]
+    fn sha256_ni_transform(digest: &mut [u32; 8], data: *const [u8; 64], nblk: u64);
 }
 
-/// Safe wrapper around assembly implementation of SHA256 compression function
-#[inline]
-pub fn compress256(state: &mut [u32; 8], blocks: &[[u8; 64]]) {
-    for block in blocks {
-        unsafe { sha256_compress(state, block) }
+cfg_if::cfg_if! {
+    if #[cfg(target_feature = "aes")]
+    {
+        #[inline]
+        pub fn compress256(state: &mut [u32; 8], blocks: &[[u8; 64]]) {
+            if !blocks.is_empty() {
+                unsafe { sha256_ni_transform(state, blocks.as_ptr(), blocks.len() as u64) }
+            }
+        }
+    }
+    else if #[cfg(target_feature = "avx2")]
+    {
+        #[inline]
+        pub fn compress256(state: &mut [u32; 8], blocks: &[[u8; 64]]) {
+            if !blocks.is_empty() {
+                unsafe { sha256_transform_rorx(state, blocks.as_ptr(), blocks.len()) }
+            }
+        }
+    }
+    else{
+        #[inline]
+        pub fn compress256(state: &mut [u32; 8], blocks: &[[u8; 64]]) {
+            for block in blocks {
+                unsafe { sha256_compress(state, block) }
+            }
+        }
     }
 }
 
 #[cfg(not(target_arch = "aarch64"))]
 #[link(name = "sha512", kind = "static")]
 extern "C" {
+    #[cfg(not(target_feature = "avx2"))]
     fn sha512_compress(state: &mut [u64; 8], block: &[u8; 128]);
+    #[cfg(target_feature = "avx2")]
+    fn sha512_transform_rorx(state: &mut [u64; 8], block: *const [u8; 128], num_blocks: usize);
 }
 
 /// Safe wrapper around assembly implementation of SHA512 compression function
 ///
 /// This function is available only on x86 and x86-64 targets.
+#[inline]
 #[cfg(not(target_arch = "aarch64"))]
+#[cfg(target_feature = "avx2")]
+pub fn compress512(state: &mut [u64; 8], blocks: &[[u8; 128]]) {
+    if !blocks.is_empty() {
+        unsafe { sha512_transform_rorx(state, blocks.as_ptr(), blocks.len()) }
+    }
+}
+
 #[inline]
+#[cfg(not(target_arch = "aarch64"))]
+#[cfg(not(target_feature = "avx2"))]
 pub fn compress512(state: &mut [u64; 8], blocks: &[[u8; 128]]) {
     for block in blocks {
         unsafe { sha512_compress(state, block) }