use crate::Hasher;
#[cfg(not(feature = "std"))]
use alloc::vec;
use bytes::{Buf, BufMut};
use commonware_codec::{
DecodeExt, Error as CodecError, FixedArray, FixedSize, Read, ReadExt, Write,
};
use commonware_formatting::Hex;
use commonware_math::algebra::Random;
use commonware_utils::{Array, Span};
use core::{
fmt::{Debug, Display},
ops::Deref,
};
use rand_core::CryptoRng;
use sha2::{Digest as _, Sha256 as ISha256, block_api::compress256};
use zeroize::Zeroize;
#[cfg(any(target_arch = "aarch64", target_arch = "x86_64"))]
mod simd;
pub type CoreSha256 = ISha256;
const DIGEST_LENGTH: usize = 32;
const BLOCK_LENGTH: usize = 64;
const MAX_FIXED: usize = 2 * BLOCK_LENGTH - 9;
const IV: [u32; 8] = [
0x6a09e667, 0xbb67ae85, 0x3c6ef372, 0xa54ff53a, 0x510e527f, 0x9b05688c, 0x1f83d9ab, 0x5be0cd19,
];
#[inline]
fn digest_from_state(state: [u32; 8]) -> [u8; DIGEST_LENGTH] {
let mut out = [0u8; DIGEST_LENGTH];
for (chunk, word) in out.as_chunks_mut::<4>().0.iter_mut().zip(state) {
*chunk = word.to_be_bytes();
}
out
}
#[inline]
fn finalize_fixed_fresh(scratch: &mut [u8; 2 * BLOCK_LENGTH], len: usize) -> [u8; DIGEST_LENGTH] {
assert!(len <= MAX_FIXED);
let bit_len = ((len as u64) * 8).to_be_bytes();
scratch[len] = 0x80;
let mut state = IV;
if len < BLOCK_LENGTH - 8 {
scratch[BLOCK_LENGTH - 8..BLOCK_LENGTH].copy_from_slice(&bit_len);
let (blocks, _) = scratch[..BLOCK_LENGTH].as_chunks::<BLOCK_LENGTH>();
compress256(&mut state, blocks);
} else {
scratch[2 * BLOCK_LENGTH - 8..].copy_from_slice(&bit_len);
let (blocks, _) = scratch.as_chunks::<BLOCK_LENGTH>();
compress256(&mut state, blocks);
}
digest_from_state(state)
}
#[inline(always)]
fn hash_specialized(parts: &[&[u8]]) -> Digest {
match parts {
[p, l, r] if p.len() == 8 && l.len() == 32 && r.len() == 32 => {
let mut scratch = [0u8; 2 * BLOCK_LENGTH];
scratch[..8].copy_from_slice(p);
scratch[8..40].copy_from_slice(l);
scratch[40..72].copy_from_slice(r);
Digest(finalize_fixed_fresh(&mut scratch, 72))
}
[a, b] if a.len() == 32 && b.len() == 32 => {
let mut scratch = [0u8; 2 * BLOCK_LENGTH];
scratch[..32].copy_from_slice(a);
scratch[32..64].copy_from_slice(b);
Digest(finalize_fixed_fresh(&mut scratch, 64))
}
[p, d] if p.len() == 8 && d.len() == 32 => {
let mut scratch = [0u8; 2 * BLOCK_LENGTH];
scratch[..8].copy_from_slice(p);
scratch[8..40].copy_from_slice(d);
Digest(finalize_fixed_fresh(&mut scratch, 40))
}
[p, d] if p.len() == 4 && d.len() == 32 => {
let mut scratch = [0u8; 2 * BLOCK_LENGTH];
scratch[..4].copy_from_slice(p);
scratch[4..36].copy_from_slice(d);
Digest(finalize_fixed_fresh(&mut scratch, 36))
}
_ => hash_general(parts),
}
}
#[inline(never)]
fn hash_general(parts: &[&[u8]]) -> Digest {
let mut scratch = [0u8; 2 * BLOCK_LENGTH];
let mut len = 0usize;
let mut parts = parts.iter();
loop {
match parts.next() {
Some(part) if len + part.len() <= MAX_FIXED => {
scratch[len..len + part.len()].copy_from_slice(part);
len += part.len();
}
Some(part) => {
let mut hasher = ISha256::new();
hasher.update(&scratch[..len]);
hasher.update(part);
for part in parts {
hasher.update(part);
}
let array: [u8; DIGEST_LENGTH] = hasher.finalize().into();
return Digest(array);
}
None => break,
}
}
Digest(finalize_fixed_fresh(&mut scratch, len))
}
#[derive(Debug, Default)]
pub struct Sha256 {
hasher: ISha256,
}
impl Sha256 {
pub fn fill(b: u8) -> <Self as Hasher>::Digest {
<Self as Hasher>::Digest::decode(vec![b; DIGEST_LENGTH].as_ref()).unwrap()
}
}
impl Hasher for Sha256 {
type Digest = Digest;
#[inline]
fn hash(parts: &[&[u8]]) -> Self::Digest {
hash_specialized(parts)
}
#[inline]
fn hash_pair(left: &[&[u8]], right: &[&[u8]]) -> (Self::Digest, Self::Digest) {
#[cfg(any(target_arch = "aarch64", target_arch = "x86_64"))]
if let Some(pair) = simd::hash_pair(left, right) {
return pair;
}
(Self::hash(left), Self::hash(right))
}
#[inline]
fn update(&mut self, message: &[u8]) -> &mut Self {
self.hasher.update(message);
self
}
#[inline]
fn finalize(mut self) -> (Self, Self::Digest) {
let finalized = self.hasher.finalize_reset();
let array: [u8; DIGEST_LENGTH] = finalized.into();
(self, Digest(array))
}
}
#[derive(Clone, Copy, Eq, PartialEq, Ord, PartialOrd, Hash, FixedArray)]
#[fixed_array(infallible)]
#[repr(transparent)]
pub struct Digest(pub [u8; DIGEST_LENGTH]);
#[cfg(feature = "arbitrary")]
impl<'a> arbitrary::Arbitrary<'a> for Digest {
fn arbitrary(u: &mut arbitrary::Unstructured<'a>) -> arbitrary::Result<Self> {
let len = u.int_in_range(0..=256)?;
let data = u.bytes(len)?;
Ok(Sha256::hash(&[data]))
}
}
impl Write for Digest {
fn write(&self, buf: &mut impl BufMut) {
self.0.write(buf);
}
}
impl Read for Digest {
type Cfg = ();
fn read_cfg(buf: &mut impl Buf, _: &()) -> Result<Self, CodecError> {
let array = <[u8; DIGEST_LENGTH]>::read(buf)?;
Ok(Self(array))
}
}
impl FixedSize for Digest {
const SIZE: usize = DIGEST_LENGTH;
}
impl Span for Digest {}
impl Array for Digest {}
impl AsRef<[u8]> for Digest {
fn as_ref(&self) -> &[u8] {
&self.0
}
}
impl Deref for Digest {
type Target = [u8];
fn deref(&self) -> &[u8] {
&self.0
}
}
impl Debug for Digest {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
write!(f, "{}", Hex(&self.0))
}
}
impl Display for Digest {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
write!(f, "{}", Hex(&self.0))
}
}
impl crate::Digest for Digest {
const EMPTY: Self = Self([0u8; DIGEST_LENGTH]);
}
impl Random for Digest {
fn random(mut rng: impl CryptoRng) -> Self {
let mut array = [0u8; DIGEST_LENGTH];
rng.fill_bytes(&mut array);
Self(array)
}
}
impl Zeroize for Digest {
fn zeroize(&mut self) {
self.0.zeroize();
}
}
#[cfg(test)]
mod tests {
use super::*;
use commonware_codec::{DecodeExt, Encode};
const HELLO_DIGEST: [u8; DIGEST_LENGTH] = commonware_formatting::hex!(
"b94d27b9934d3e08a52e52d7da7dabfac484efe37a5380ee9088f7ace2efcde9"
);
#[test]
fn test_sha256() {
let msg = b"hello world";
let mut hasher = Sha256::default();
hasher.update(msg);
let (_, digest) = hasher.finalize();
assert!(Digest::decode(digest.as_ref()).is_ok());
assert_eq!(digest.as_ref(), HELLO_DIGEST);
let hash = Sha256::hash(&[msg]);
assert_eq!(hash.as_ref(), HELLO_DIGEST);
let hash = Sha256::hash(&[b"hello", b" world"]);
assert_eq!(hash.as_ref(), HELLO_DIGEST);
}
#[test]
fn test_sha256_hash_parts_boundaries() {
for total in 0..=300usize {
let data: Vec<u8> = (0..total).map(|i| i as u8).collect();
let mid = total / 3;
let parts: [&[u8]; 3] = [&data[..mid], &data[mid..2 * mid], &data[2 * mid..]];
let oneshot = Sha256::hash(&parts);
let mut hasher = Sha256::default();
for part in &parts {
hasher.update(part);
}
let (_, streamed) = hasher.finalize();
assert_eq!(oneshot, streamed, "mismatch for total={total}");
}
}
#[test]
fn test_sha256_hash_specialized_arms() {
let data: Vec<u8> = (0u8..72).collect();
let shapes: [&[&[u8]]; 4] = [
&[&data[..8], &data[8..40], &data[40..72]],
&[&data[..32], &data[32..64]],
&[&data[..8], &data[8..40]],
&[&data[..4], &data[4..36]],
];
for parts in shapes {
let oneshot = Sha256::hash(parts);
let mut hasher = Sha256::default();
for part in parts {
hasher.update(part);
}
let (_, streamed) = hasher.finalize();
assert_eq!(oneshot, streamed, "mismatch for shape {parts:?}");
}
}
#[test]
fn test_sha256_len() {
assert_eq!(Digest::SIZE, DIGEST_LENGTH);
}
#[test]
fn test_hash_pair_mmr_node_shape_matches_streaming() {
fn node(position: u64, fill: u8) -> Vec<Vec<u8>> {
vec![
position.to_be_bytes().to_vec(),
vec![fill; 32],
vec![fill + 1; 32],
]
}
crate::fuzz::Plan::<Sha256>::new(node(42, 0x11), node(43, 0x33)).run();
}
#[test]
fn test_hash_pair_bmt_node_shape_matches_streaming() {
fn node(fill: u8) -> Vec<Vec<u8>> {
vec![vec![fill; 32], vec![fill + 1; 32]]
}
crate::fuzz::Plan::<Sha256>::new(node(0x11), node(0x33)).run();
}
#[test]
fn test_codec() {
let msg = b"hello world";
let mut hasher = Sha256::default();
hasher.update(msg);
let (_, digest) = hasher.finalize();
let encoded = digest.encode();
assert_eq!(encoded.len(), DIGEST_LENGTH);
assert_eq!(encoded, digest.as_ref());
let decoded = Digest::decode(encoded).unwrap();
assert_eq!(digest, decoded);
}
#[cfg(feature = "arbitrary")]
mod conformance {
use super::*;
use commonware_codec::conformance::CodecConformance;
commonware_conformance::conformance_tests! {
CodecConformance<Digest>,
}
}
}