From e72286899ddbae2f93815c2d9250885122189862 Mon Sep 17 00:00:00 2001 From: Nikola Pajkovsky Date: Wed, 26 Aug 2026 09:42:49 +0200 Subject: [PATCH] sha256: add AArch64 SHA-2 crypto-extension hardware backend MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds an opt-in (`asm` feature) hardware compression backend for SHA-256 on little-endian aarch64, using the ARMv8 Cryptographic Extension (SHA256H/SHA256H2/SHA256SU0/SHA256SU1) via core::arch intrinsics. Runtime support is detected with is_aarch64_feature_detected!, falling back to the portable scalar implementation everywhere else. - New crypto/sha2/src/asm/ module: one submodule per algorithm (asm::sha256), dispatching by (target_arch, target_endian) at its own `mod` declaration; try_compress is always callable and is a no-op on targets with no backend, so callers never need their own #[cfg]. - #![forbid(unsafe_code)] moves from the crate root into sha256.rs/ sha512.rs (unconditional, independent of the `asm` feature); the crate root drops to #![deny(unsafe_code)] and asm/mod.rs grants itself the sole #![allow(unsafe_code)] exception, since `forbid` can never be overridden by a nested #[allow]. - Forwards the feature through the workspace Cargo.toml as `sha2-asm`. - Extends the sha256 benchmark with a 1KiB-chunked and a 1MiB one-shot case alongside the existing 16KiB one-shot. Benchmark (16KiB one-shot, criterion): baseline (scalar): 33.147 µs ( 471.38 MiB/s) asm (aarch64 SHA-2): 4.757 µs (3.2077 GiB/s) ~7x faster Assisted-by: Claude:sonnet-5 Signed-off-by: Nikola Pajkovsky --- Cargo.toml | 6 +- crypto/sha2/Cargo.toml | 5 + crypto/sha2/benches/sha2_benches.rs | 33 +++- crypto/sha2/src/asm/mod.rs | 15 ++ crypto/sha2/src/asm/sha256/aarch64_le.rs | 202 +++++++++++++++++++++++ crypto/sha2/src/asm/sha256/mod.rs | 19 +++ crypto/sha2/src/lib.rs | 8 +- crypto/sha2/src/sha256.rs | 29 +++- crypto/sha2/src/sha512.rs | 2 + 9 files changed, 303 insertions(+), 16 deletions(-) create mode 100644 crypto/sha2/src/asm/mod.rs create mode 100644 crypto/sha2/src/asm/sha256/aarch64_le.rs create mode 100644 crypto/sha2/src/asm/sha256/mod.rs diff --git a/Cargo.toml b/Cargo.toml index 6e2ed3f9..54971e10 100644 --- a/Cargo.toml +++ b/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 \ No newline at end of file +bouncycastle-sha3.workspace = true + +[features] +# Forwarded opt-in for bouncycastle-sha2's hardware SHA-256 acceleration. +sha2-asm = ["bouncycastle-sha2/asm"] diff --git a/crypto/sha2/Cargo.toml b/crypto/sha2/Cargo.toml index 6de8662d..e5d01339 100644 --- a/crypto/sha2/Cargo.toml +++ b/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 diff --git a/crypto/sha2/benches/sha2_benches.rs b/crypto/sha2/benches/sha2_benches.rs index 0d12a00a..9983bf0c 100644 --- a/crypto/sha2/benches/sha2_benches.rs +++ b/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) { diff --git a/crypto/sha2/src/asm/mod.rs b/crypto/sha2/src/asm/mod.rs new file mode 100644 index 00000000..fb5fd860 --- /dev/null +++ b/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; diff --git a/crypto/sha2/src/asm/sha256/aarch64_le.rs b/crypto/sha2/src/asm/sha256/aarch64_le.rs new file mode 100644 index 00000000..19a48cce --- /dev/null +++ b/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::::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 + } +} diff --git a/crypto/sha2/src/asm/sha256/mod.rs b/crypto/sha2/src/asm/sha256/mod.rs new file mode 100644 index 00000000..4121790e --- /dev/null +++ b/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 +} diff --git a/crypto/sha2/src/lib.rs b/crypto/sha2/src/lib.rs index 6906e0c6..18383e9d 100644 --- a/crypto/sha2/src/lib.rs +++ b/crypto/sha2/src/lib.rs @@ -65,10 +65,16 @@ //! let h: Vec = 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; diff --git a/crypto/sha2/src/sha256.rs b/crypto/sha2/src/sha256.rs index 7cee95ca..e99e42d6 100644 --- a/crypto/sha2/src/sha256.rs +++ b/crypto/sha2/src/sha256.rs @@ -1,3 +1,5 @@ +#![forbid(unsafe_code)] + use crate::SHA2Params; use bouncycastle_core::errors::{HashError, SuspendableError}; use bouncycastle_core::suspendable_state::{add_lib_ver, check_lib_ver}; @@ -5,7 +7,7 @@ use bouncycastle_core::traits::{Algorithm, Hash, SecurityStrength, Suspendable}; use bouncycastle_utils::{min, secret::Secret}; use core::slice; -const SHA256_K: [u32; 64] = [ +pub(crate) const SHA256_K: [u32; 64] = [ 0x428A2F98, 0x71374491, 0xB5C0FBCF, 0xE9B5DBA5, 0x3956C25B, 0x59F111F1, 0x923F82A4, 0xAB1C5ED5, 0xD807AA98, 0x12835B01, 0x243185BE, 0x550C7DC3, 0x72BE5D74, 0x80DEB1FE, 0x9BDC06A7, 0xC19BF174, 0xE49B69C1, 0xEFBE4786, 0x0FC19DC6, 0x240CA1CC, 0x2DE92C6F, 0x4A7484AA, 0x5CB0A9DC, 0x76F988DA, @@ -16,32 +18,32 @@ const SHA256_K: [u32; 64] = [ 0x748F82EE, 0x78A5636F, 0x84C87814, 0x8CC70208, 0x90BEFFFA, 0xA4506CEB, 0xBEF9A3F7, 0xC67178F2, ]; -#[inline] +#[inline(always)] fn ch(x: u32, y: u32, z: u32) -> u32 { (x & y) ^ (!x & z) } -#[inline] +#[inline(always)] fn maj(x: u32, y: u32, z: u32) -> u32 { (x & y) | (z & (x ^ y)) } -#[inline] +#[inline(always)] fn sum0(x: u32) -> u32 { x.rotate_right(2) ^ x.rotate_right(13) ^ x.rotate_right(22) } -#[inline] +#[inline(always)] fn sum1(x: u32) -> u32 { x.rotate_right(6) ^ x.rotate_right(11) ^ x.rotate_right(25) } -#[inline] +#[inline(always)] fn theta0(x: u32) -> u32 { x.rotate_right(7) ^ x.rotate_right(18) ^ (x >> 3) } -#[inline] +#[inline(always)] fn theta1(x: u32) -> u32 { x.rotate_right(17) ^ x.rotate_right(19) ^ (x >> 10) } @@ -49,7 +51,7 @@ fn theta1(x: u32) -> u32 { #[derive(Clone)] pub(crate) struct Sha256State { _params: core::marker::PhantomData, - h: Secret<[u32; 8]>, + pub(crate) h: Secret<[u32; 8]>, } impl Sha256State { @@ -75,6 +77,15 @@ impl Sha256State { } fn compress(&mut self, blocks: &[[u8; 64]]) { + if crate::asm::sha256::try_compress(&mut self.h, blocks) { + return; + } + + self.compress_scalar(blocks); + } + + /// Portable compression function, FIPS 180-4 section 6.2.2. + pub(crate) fn compress_scalar(&mut self, blocks: &[[u8; 64]]) { let mut x = [0u32; 64]; // infallible; just unwrapping the [u32; 8] and re-casting to itself. @@ -150,7 +161,7 @@ pub struct SHA256Internal { byte_count: u64, x_buf: Secret<[u8; 64]>, x_buf_off: usize, - // TODO: Investigate whether maximum message size (according to FIPS 180-4) should be added + // TODO: Investigate whether maximum message size (according to FIPS 180-4) should be added // (2^64 for SHA256 and 2^128 for SHA512) } diff --git a/crypto/sha2/src/sha512.rs b/crypto/sha2/src/sha512.rs index 60207eb7..1c078507 100644 --- a/crypto/sha2/src/sha512.rs +++ b/crypto/sha2/src/sha512.rs @@ -1,3 +1,5 @@ +#![forbid(unsafe_code)] + use crate::SHA2Params; use bouncycastle_core::errors::{HashError, SuspendableError}; use bouncycastle_core::suspendable_state::{add_lib_ver, check_lib_ver};