Include hash functions (SHA-2, SHA-3, Keccak and BLAKE2)

Cargo.toml

@@ -42,7 +42,8 @@ blst = "0.3.10"
 digest = "0.10.5"
 once_cell = "1.14.0"
 readonly = "0.2.2"
-sha2 = "0.10.2"
+sha2 = "0.10.6"
+sha3 = "0.10.5"
 thiserror = "1.0.32"
 
 fastcrypto-derive = { path = "./fastcrypto-derive", version = "0.1.0" }
@@ -63,6 +64,5 @@ k256 = { version = "0.11.5", features = ["ecdsa", "sha256", "keccak256"] }
 proptest = "1.0.0"
 proptest-derive = "0.3.0"
 serde_json = "1.0.85"
-sha3 = "0.10.2"
 serde-reflection = "0.3.6"
 wycheproof = "0.4.0"

benches/crypto.rs

@@ -7,15 +7,16 @@ extern crate rand;
 
 mod ed25519_benches {
     use super::*;
-    use blake2::digest::Update;
     use criterion::*;
     use fastcrypto::{
         bls12381::{BLS12381KeyPair, BLS12381Signature},
         ed25519::*,
+        hash::Blake2b,
         secp256k1::{Secp256k1KeyPair, Secp256k1Signature},
         traits::{KeyPair, VerifyingKey},
         Verifier,
     };
+    use generic_array::typenum::U32;
     use rand::{prelude::ThreadRng, thread_rng};
     use signature::Signer;
 
@@ -74,12 +75,10 @@ mod ed25519_benches {
                 .map(|_| BLS12381KeyPair::generate(&mut csprng))
                 .collect();
 
-            let msg: Vec<u8> = {
-                fastcrypto::blake2b_256(|hasher| {
-                    hasher.update(b"aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa")
-                })
-                .to_vec()
-            };
+            let msg: Vec<u8> = fastcrypto::hash::Hashable::digest::<Blake2b<U32>>(
+                b"aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa".as_slice(),
+            )
+            .to_vec();
 
             let ed_signatures: Vec<_> = ed_keypairs.iter().map(|key| key.sign(&msg)).collect();
             let ed_public_keys: Vec<_> =

src/hash.rs

@@ -0,0 +1,144 @@
+// Copyright (c) 2022, Mysten Labs, Inc.
+// SPDX-License-Identifier: Apache-2.0
+
+use base64ct::{Base64, Encoding};
+use blake2::{
+    digest::{Update, VariableOutput},
+    VarBlake2b,
+};
+use digest::OutputSizeUser;
+use generic_array::{ArrayLength, GenericArray};
+use serde::{Deserialize, Serialize};
+use std::{fmt, marker::PhantomData};
+
+/// Represents a hash digest of `DigestLength` bytes.
+#[derive(Hash, PartialEq, Eq, Clone, Deserialize, Serialize, Ord, PartialOrd)]
+pub struct Digest<DigestLength: ArrayLength<u8> + 'static>(pub GenericArray<u8, DigestLength>);
+
+impl<DigestLength: ArrayLength<u8> + 'static> Digest<DigestLength> {
+    /// Clone the given slice into a new `Digest`.
+    pub fn from_bytes(val: &[u8]) -> Self {
+        Digest(GenericArray::<u8, DigestLength>::clone_from_slice(val))
+    }
+
+    /// Copy the digest into a new vector.
+    pub fn to_vec(&self) -> Vec<u8> {
+        self.0.to_vec()
+    }
+
+    /// The size of this digest in bytes.
+    pub fn size(&self) -> usize {
+        DigestLength::USIZE
+    }
+}
+
+impl<DigestLength: ArrayLength<u8> + 'static> fmt::Debug for Digest<DigestLength> {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> Result<(), fmt::Error> {
+        write!(f, "{}", Base64::encode_string(&self.0))
+    }
+}
+
+impl<DigestLength: ArrayLength<u8> + 'static> fmt::Display for Digest<DigestLength> {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> Result<(), fmt::Error> {
+        write!(
+            f,
+            "{}",
+            Base64::encode_string(&self.0)
+                .get(0..DigestLength::USIZE)
+                .unwrap()
+        )
+    }
+}
+
+impl<DigestLength: ArrayLength<u8> + 'static> AsRef<[u8]> for Digest<DigestLength> {
+    fn as_ref(&self) -> &[u8] {
+        self.0.as_ref()
+    }
+}
+
+/// This trait is implemented by all messages that can be hashed.
+pub trait Hashable<DigestLength: ArrayLength<u8> + 'static> {
+    fn digest<H: HashFunction<DigestLength>>(self) -> Digest<DigestLength>;
+}
+
+impl<DigestLength: ArrayLength<u8> + 'static> Hashable<DigestLength> for &[u8] {
+    /// Hash this data using the given hash function.
+    fn digest<H: HashFunction<DigestLength>>(self) -> Digest<DigestLength> {
+        H::digest(self)
+    }
+}
+
+/// Trait implemented by hash functions providing a output of fixed length
+pub trait HashFunction<DigestLength: ArrayLength<u8>>: OutputSizeUser + Sized + Default {
+    /// Process the given data, and update the internal of the hash function.
+    fn update(&mut self, data: &[u8]);
+
+    /// Retrieve result and consume hash function.
+    fn finalize(self) -> Digest<DigestLength>;
+
+    fn digest(data: &[u8]) -> Digest<DigestLength> {
+        let mut h = Self::default();
+        h.update(data);
+        h.finalize()
+    }
+}
+
+#[derive(Default)]
+pub struct HashFunctionWrapper<Variant: digest::Digest + 'static>(Variant);
+
+impl<Variant: digest::Digest + 'static> OutputSizeUser for HashFunctionWrapper<Variant> {
+    type OutputSize = Variant::OutputSize;
+}
+
+impl<Variant: digest::Digest + 'static + Default> HashFunction<Variant::OutputSize>
+    for HashFunctionWrapper<Variant>
+{
+    fn update(&mut self, data: &[u8]) {
+        self.0.update(data);
+    }
+
+    fn finalize(self) -> Digest<Variant::OutputSize> {
+        Digest(self.0.finalize())
+    }
+}
+
+// SHA-2
+pub type Sha256 = HashFunctionWrapper<sha2::Sha256>;
+
+// SHA-3
+pub type Sha3_256 = HashFunctionWrapper<sha3::Sha3_256>;
+
+// KECCAK
+pub type Keccak256 = HashFunctionWrapper<sha3::Keccak256>;
+
+// BLAKE2
+pub struct Blake2b<DigestLength: ArrayLength<u8>> {
+    variant: blake2::VarBlake2b,
+    digest_length: PhantomData<DigestLength>,
+}
+
+impl<DigestLength: ArrayLength<u8>> OutputSizeUser for Blake2b<DigestLength> {
+    type OutputSize = DigestLength;
+}
+
+impl<DigestLength: ArrayLength<u8> + Sized> HashFunction<DigestLength> for Blake2b<DigestLength> {
+    fn update(&mut self, data: &[u8]) {
+        self.variant.update(data)
+    }
+
+    fn finalize(self) -> Digest<DigestLength> {
+        let mut array = GenericArray::<u8, DigestLength>::default();
+        self.variant
+            .finalize_variable(|buffer| array.copy_from_slice(buffer));
+        Digest(array)
+    }
+}
+
+impl<DigestLength: ArrayLength<u8>> Default for Blake2b<DigestLength> {
+    fn default() -> Self {
+        Self {
+            variant: VarBlake2b::new(DigestLength::USIZE).unwrap(),
+            digest_length: Default::default(),
+        }
+    }
+}

src/hmac.rs

@@ -3,7 +3,7 @@
 use crate::error::FastCryptoError;
 use crate::{
     traits::{KeyPair, SigningKey, ToFromBytes},
-    Digest, DIGEST_LEN,
+    Digest,
 };
 use digest::{
     block_buffer::Eager,
@@ -95,11 +95,11 @@ const HKDF_KEY_RECOMMENDED_LENGTH: usize = 32;
 pub type HmacKey = PrivateSeed<HMAC_KEY_RECOMMENDED_LENGTH, false>;
 pub type HkdfIkm = PrivateSeed<HKDF_KEY_RECOMMENDED_LENGTH, false>;
 
-pub fn hmac(key: &HmacKey, message: &[u8]) -> Digest {
+pub fn hmac(key: &HmacKey, message: &[u8]) -> Digest<typenum::U32> {
     let mut hash = Hmac::<sha2::Sha256>::new_from_slice(key.as_bytes()).unwrap();
     hash.update(message);
-    let output: [u8; DIGEST_LEN] = hash.finalize().into_bytes().as_slice().try_into().unwrap();
-    Digest::new(output)
+    let output = hash.finalize();
+    Digest(output.into_bytes())
 }
 
 pub fn hkdf(