sha256: add AArch64 SHA-2 crypto-extension hardware backend

Cargo.toml

@@ -52,4 +52,8 @@ bouncycastle-mlkem.workspace = true
 bouncycastle-mlkem-lowmemory.workspace = true
 bouncycastle-rng.workspace = true
 bouncycastle-sha2.workspace = true
-bouncycastle-sha3.workspace = true
+bouncycastle-sha3.workspace = true
+
+[features]
+# Forwarded opt-in for bouncycastle-sha2's hardware SHA-256 acceleration.
+sha2-asm = ["bouncycastle-sha2/asm"]

crypto/sha2/Cargo.toml

@@ -3,6 +3,11 @@ name = "bouncycastle-sha2"
 version = "0.1.2"
 edition.workspace = true
 
+[features]
+# Opt-in hardware SHA-256 compression via core::arch. Currently it uses runtime
+# SHA-2 detection
+asm = []
+
 [dependencies]
 bouncycastle-core.workspace = true
 bouncycastle-utils.workspace = true

crypto/sha2/benches/sha2_benches.rs

@@ -6,24 +6,47 @@ use bouncycastle_rng as rng;
 use bouncycastle_sha2::*;
 
 fn bench_sha256(c: &mut Criterion) {
-    let mut data = [0_u8; 1024];
-    rng::DefaultRNG::default().next_bytes_out(&mut data).unwrap();
+    let mut data = vec![0_u8; 1024 * 1024];
+    let mut rng = rng::DefaultRNG::default();
+    for chunk in data.chunks_mut(1024) {
+        rng.next_bytes_out(chunk).unwrap();
+    }
 
     let mut digest = vec![0; SHA256::new().output_len()];
 
     let mut group = c.benchmark_group("sha2::sha256");
     group.throughput(Throughput::Bytes(16 * 1024));
-    group.bench_function("16KiB", |b| {
+    group.bench_function("16KiB/one-shot", |b| {
         b.iter(|| {
             let mut md = SHA256::new();
-            for _ in 0..16 {
-                md.do_update(black_box(&data));
+            md.do_update(black_box(&data[..16 * 1024]));
+            _ = md.do_final_out(&mut digest);
+            black_box(&digest);
+        })
+    });
+    group.bench_function("16KiB/1KiB-chunks", |b| {
+        b.iter(|| {
+            let mut md = SHA256::new();
+            for chunk in data[..16 * 1024].chunks(1024) {
+                md.do_update(black_box(chunk));
             }
             _ = md.do_final_out(&mut digest);
             black_box(&digest);
         })
     });
     group.finish();
+
+    let mut group = c.benchmark_group("sha2::sha256");
+    group.throughput(Throughput::Bytes(1024 * 1024));
+    group.bench_function("1MiB/one-shot", |b| {
+        b.iter(|| {
+            let mut md = SHA256::new();
+            md.do_update(black_box(&data));
+            _ = md.do_final_out(&mut digest);
+            black_box(&digest);
+        })
+    });
+    group.finish();
 }
 
 fn bench_sha512(c: &mut Criterion) {

crypto/sha2/src/asm/mod.rs

@@ -0,0 +1,15 @@
+//! Hardware-accelerated compression backends.
+//!
+//! One submodule per algorithm; each algorithm submodule dispatches further
+//! by `(target_arch, target_endian)` — see [`sha256`] for the shape every
+//! backend follows.
+//!
+//! This is the one place in the crate that grants itself an exception from
+//! the `sha256`/`sha512` modules' `#![forbid(unsafe_code)]`: every backend
+//! under here is gated on its own `(feature = "asm", target_arch,
+//! target_endian)` cfg and carries a `// SAFETY:` comment at its unsafe
+//! block, but none of it can compile unless this crate's `asm` feature is
+//! explicitly opted into.
+#![allow(unsafe_code)]
+
+pub(crate) mod sha256;

crypto/sha2/src/asm/sha256/aarch64_le.rs

@@ -0,0 +1,202 @@
+//! Hardware SHA-256 compression via the ARMv8 Cryptographic Extension
+//! (`SHA256H`/`SHA256H2`/`SHA256SU0`/`SHA256SU1`) — the `intrinsics`-feature
+//! counterpart of the scalar compression function.
+//!
+//! Instruction availability is checked at runtime.
+
+use crate::sha256::SHA256_K;
+
+#[repr(align(16))]
+struct AlignedK([u32; 64]);
+
+static SHA256_K_HW: AlignedK = AlignedK(SHA256_K);
+
+/// Compresses `blocks` into the state `h` (FIPS 180-4 section 6.2.2,
+/// four rounds per `SHA256H`/`SHA256H2` pair), returning whether the
+/// hardware implementation was available.
+pub(crate) fn try_compress(h: &mut [u32; 8], blocks: &[[u8; 64]]) -> bool {
+    if !is_supported() {
+        return false;
+    }
+
+    // SAFETY: `is_supported` established that this CPU implements the
+    // SHA-2 instructions required by `compress_blocks`.
+    unsafe { compress_blocks(h, blocks) }
+    true
+}
+
+#[inline]
+fn is_supported() -> bool {
+    cfg!(target_feature = "sha2") || std::arch::is_aarch64_feature_detected!("sha2")
+}
+
+#[target_feature(enable = "sha2")]
+unsafe fn compress_blocks(h: &mut [u32; 8], blocks: &[[u8; 64]]) {
+    use core::arch::aarch64::{
+        vaddq_u32, vld1q_u32, vld1q_u8, vreinterpretq_u32_u8, vrev32q_u8, vsha256h2q_u32,
+        vsha256hq_u32, vsha256su0q_u32, vsha256su1q_u32, vst1q_u32,
+    };
+    use core::arch::asm;
+
+    // SAFETY: all loads/stores are within `h` ([u32; 8], read/written as
+    // two 4-lane halves), the current 64-byte `block` (read as four
+    // 16-byte quarters), and `SHA256_K_HW` ([u32; 64], read as sixteen
+    // 4-lane rows). The caller established SHA-2 instruction support.
+    unsafe {
+        let k = SHA256_K_HW.0.as_ptr();
+        let mut abcd = vld1q_u32(h.as_ptr());
+        let mut efgh = vld1q_u32(h.as_ptr().add(4));
+
+        for block in blocks {
+            let p = block.as_ptr();
+            let abcd_save = abcd;
+            let efgh_save = efgh;
+
+            // Load + byte-swap the 16 message words for this block, same
+            // as the scalar version's `x[0..16]` — just 4 at a time.
+            let mut m0 = vreinterpretq_u32_u8(vrev32q_u8(vld1q_u8(p)));
+            let mut m1 = vreinterpretq_u32_u8(vrev32q_u8(vld1q_u8(p.add(16))));
+            let mut m2 = vreinterpretq_u32_u8(vrev32q_u8(vld1q_u8(p.add(32))));
+            let mut m3 = vreinterpretq_u32_u8(vrev32q_u8(vld1q_u8(p.add(48))));
+
+            // K row for the group in flight; each group fetches the next
+            // row before its own hash pair issues, so the load is off the
+            // critical path (same software pipeline as the hand-written
+            // assembly this replaced).
+            let mut k_cur = vld1q_u32(k);
+
+            // SHA256H + SHA256H2 together do 4 rounds of mixing in one go
+            // — the scalar version does the same 4 rounds one at a time
+            // (via 4 calls to `sha256_round!`), reshuffling which
+            // variable plays which role each round so it doesn't have to
+            // physically move 8 values around. The hardware does that
+            // mixing and shuffling internally, so there's no reshuffling
+            // to write here — `abcd`/`efgh` just get overwritten in place.
+            //
+            // (The empty-asm block below is not part of the algorithm —
+            // it's a compiler hint. Without it, LLVM's register allocator
+            // makes a suboptimal choice that costs ~17% performance; the
+            // hint just pins a temporary copy in its own register.)
+            macro_rules! hash_pair {
+                ($wk:expr) => {
+                    let wk = $wk;
+                    let mut prev = efgh;
+                    // Empty register barrier: pins `prev` in a physical
+                    // register of its own, so the allocator cannot
+                    // coalesce the copy onto sha256h2's tied destination.
+                    asm!(
+                        "// {prev:q} register barrier",
+                        prev = inout(vreg) prev,
+                        options(pure, nomem, nostack, preserves_flags),
+                    );
+                    efgh = vsha256h2q_u32(efgh, abcd, wk);
+                    abcd = vsha256hq_u32(abcd, prev, wk);
+                };
+            }
+
+            // One "group" = 4 rounds, plus computing the next 4 message
+            // words while we're at it (SU0 before the mixing step, SU1
+            // after) — same overall work as 4 loop iterations of
+            // compress_scalar, just batched by 4.
+            //
+            // `$mc/$mn/$ma/$mb` are just the 4 message-word vectors,
+            // named by how far each one is from the one being extended
+            // right now: mc = current, mn = next, ma/mb = the other two.
+            // Each call below passes `m0..m3` shifted by one, so the same
+            // 4 vectors rotate through all 4 roles as we go.
+            macro_rules! group {
+                ($mc:ident, $mn:ident, $ma:ident, $mb:ident, $next:expr) => {
+                    let k_next = vld1q_u32(k.add($next * 4));
+                    // Round constant + message word for these 4 rounds,
+                    // added together in one step (scalar: K[t] + x[t],
+                    // one `t` at a time).
+                    let wk = vaddq_u32(k_cur, $mc);
+                    // Start computing the next 4 message words.
+                    $mc = vsha256su0q_u32($mc, $mn);
+                    hash_pair!(wk);
+                    // Finish computing them.
+                    $mc = vsha256su1q_u32($mc, $ma, $mb);
+                    k_cur = k_next;
+                };
+            }
+            // The last few rounds don't need new message words anymore
+            // (we've already computed all 64), so this skips the SU0/SU1
+            // step and just does the mixing.
+            macro_rules! tail_group {
+                ($mc:ident, $next:expr) => {
+                    let k_next = vld1q_u32(k.add($next * 4));
+                    hash_pair!(vaddq_u32(k_cur, $mc));
+                    k_cur = k_next;
+                };
+            }
+
+            group!(m0, m1, m2, m3, 1);
+            group!(m1, m2, m3, m0, 2);
+            group!(m2, m3, m0, m1, 3);
+            group!(m3, m0, m1, m2, 4);
+            group!(m0, m1, m2, m3, 5);
+            group!(m1, m2, m3, m0, 6);
+            group!(m2, m3, m0, m1, 7);
+            group!(m3, m0, m1, m2, 8);
+            group!(m0, m1, m2, m3, 9);
+            group!(m1, m2, m3, m0, 10);
+            group!(m2, m3, m0, m1, 11);
+            group!(m3, m0, m1, m2, 12);
+
+            tail_group!(m0, 13);
+            tail_group!(m1, 14);
+            tail_group!(m2, 15);
+            hash_pair!(vaddq_u32(k_cur, m3));
+
+            // Add the state we started this block with back in — same
+            // last step as compress_scalar's `s[i] += a` (etc.) loop.
+            abcd = vaddq_u32(abcd, abcd_save);
+            efgh = vaddq_u32(efgh, efgh_save);
+        }
+
+        vst1q_u32(h.as_mut_ptr(), abcd);
+        vst1q_u32(h.as_mut_ptr().add(4), efgh);
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::{is_supported, try_compress};
+    use crate::sha256::Sha256State;
+    use crate::SHA256Params;
+
+    #[test]
+    fn hardware_compression_matches_scalar() {
+        if !is_supported() {
+            return;
+        }
+
+        let mut seed = 0x6a09_e667u32;
+
+        for block_count in [0, 1, 2, 3, 8] {
+            let mut scalar = Sha256State::<SHA256Params>::new();
+            for word in scalar.h.iter_mut() {
+                seed = xorshift32(seed);
+                *word = seed;
+            }
+
+            let mut accelerated = scalar.clone();
+            let mut blocks = vec![[0u8; 64]; block_count];
+            for byte in blocks.iter_mut().flatten() {
+                seed = xorshift32(seed);
+                *byte = seed as u8;
+            }
+
+            scalar.compress_scalar(&blocks);
+            assert!(try_compress(&mut accelerated.h, &blocks));
+            assert_eq!(&*accelerated.h, &*scalar.h);
+        }
+    }
+
+    fn xorshift32(mut value: u32) -> u32 {
+        value ^= value << 13;
+        value ^= value >> 17;
+        value ^= value << 5;
+        value
+    }
+}

crypto/sha2/src/asm/sha256/mod.rs

@@ -0,0 +1,19 @@
+//! Hardware SHA-256 compression dispatch.
+//!
+//! Each backend module below is gated on its own `(target_arch,
+//! target_endian)` pair at the `mod` declaration and exposes
+//! `try_compress(h: &mut [u32; 8], blocks: &[[u8; 64]]) -> bool`, always
+//! available regardless of target or the `asm` feature — the fallback below
+//! returns `false` unconditionally when no backend module applies, so
+//! callers never need their own `#[cfg]`.
+
+#[cfg(all(feature = "asm", target_arch = "aarch64", target_endian = "little"))]
+mod aarch64_le;
+
+#[cfg(all(feature = "asm", target_arch = "aarch64", target_endian = "little"))]
+pub(crate) use aarch64_le::try_compress;
+
+#[cfg(not(all(feature = "asm", target_arch = "aarch64", target_endian = "little")))]
+pub(crate) fn try_compress(_h: &mut [u32; 8], _blocks: &[[u8; 64]]) -> bool {
+    false
+}

crypto/sha2/src/lib.rs

@@ -65,10 +65,16 @@
 //! let h: Vec<u8> = sha2_resumed.do_final();
 //! ```
 
-#![forbid(unsafe_code)]
+// Crate-wide default: unsafe code needs an explicit, reviewable `#[allow]`.
+// `sha256` and `sha512` upgrade this to `forbid` themselves (unconditionally,
+// regardless of the `asm` feature) since `forbid` can't be overridden even
+// by a local `#[allow]` — that's also why it can't be set here, since `asm`
+// (the audited hardware backends) needs to grant itself an exception.
+#![deny(unsafe_code)]
 #![forbid(missing_docs)]
 #![allow(private_bounds)]
 
+mod asm;
 mod sha256;
 mod sha512;